scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

feat(redis): move router cooldown and usage stores into Rust
Some checks are pending
LiteLLM Rust DB / db-queries (push) Waiting to run
Details
LiteLLM Rust / rust-lint (push) Waiting to run
Details
LiteLLM Rust / rust-test (push) Waiting to run
Details
LiteLLM Rust / rust-wheel (push) Waiting to run
Details

Browse source
This commit is contained in:
Devin AI 2026-11-02 00:17:51 +00:00
parent c713b92d89
commit 0a37916f92
77 changed files with 4311 additions and 2719 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
litellm-rust/Cargo.lock generated
View file

@ -3474,12 +3474,14 @@ name = "litellm-cache"
version = "0.1.0"
dependencies = [
"litellm-error",
"litellm-redis-client",
"rstest",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
@ -4357,6 +4359,8 @@ dependencies = [
"litellm-cache-redis",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-circuit-breaker",
"litellm-config",
"litellm-core-utils",
"litellm-error",
"litellm-host",
@ -4371,6 +4375,9 @@ dependencies = [
"litellm-llms",
"litellm-llms-types",
"litellm-model-catalog",
"litellm-redis-client",
"litellm-redis-connect",
"litellm-redis-store",
"litellm-router-types",
"litellm-secret-value",
"litellm-secrets",
@ -4432,6 +4439,7 @@ dependencies = [
"socket2 0.6.5",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
@ -4464,6 +4472,24 @@ dependencies = [
"veil",
]
[[package]]
name = "litellm-redis-store"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-memory",
"litellm-circuit-breaker",
"litellm-python-compat",
"litellm-redis-client",
"litellm-redis-testing",
"litellm-secret-value",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
name = "litellm-redis-testing"
version = "0.1.0"

1
litellm-rust/Cargo.toml
View file

@ -74,6 +74,7 @@ litellm-python-compat = { path = "crates/python-compat" }
litellm-circuit-breaker = { path = "crates/circuit-breaker" }
litellm-redis-client = { path = "crates/redis-client" }
litellm-redis-connect = { path = "crates/redis-connect" }
litellm-redis-store = { path = "crates/redis-store" }
litellm-redis-testing = { path = "crates/redis-testing" }
litellm-redis-types = { path = "crates/redis-types" }
litellm-cost = { path = "crates/cost" }

17
litellm-rust/crates/cache-memory/src/cache.rs
View file

@ -200,6 +200,23 @@ impl<V: Clone> InMemoryCache<V> {
.collect())
}
/// Every live entry, for callers that inspect the whole tier.
pub fn entries(&self) -> Result<Vec<(String, V)>, Error> {
let now = (self.now)();
let state = self.lock()?;
Ok(state
.values
.iter()
.filter(|(key, _)| {
state
.expirations
.get(*key)
.is_none_or(|expiration| *expiration >= now)
})
.map(|(key, value)| (key.clone(), value.clone()))
.collect())
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.lock()?;
Self::remove(&mut state, key);

2
litellm-rust/crates/cache/Cargo.toml vendored
View file

@ -7,10 +7,12 @@ repository.workspace = true
[dependencies]
litellm-error.workspace = true
litellm-redis-client.workspace = true
serde.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
strum.workspace = true
thiserror.workspace = true
tracing.workspace = true
[dev-dependencies]
rstest.workspace = true

679
litellm-rust/crates/cache/src/dual.rs vendored
View file

@ -1,519 +1,248 @@
use std::{sync::Arc, time::Duration};
use crate::{
BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheContext, ClaimCache, CounterCache,
DeleteCache, Error, FlushCache, IncrementOperation, SetCache, ShortBatch, TtlCache,
use std::{
collections::{HashMap, VecDeque},
sync::{Arc, Mutex, MutexGuard},
time::{Duration, SystemTime, UNIX_EPOCH},
};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum ReadPolicy {
#[default]
LocalThenRemote,
LocalOnly,
use litellm_redis_client::{Error as RedisError, RedisClient};
use crate::{BaseCache, DeleteCache, Error, ExactCacheContext, FlushCache};
pub type Clock = Arc<dyn Fn() -> Duration + Send + Sync>;
/// How a read-through family uses its memory tier, declared once by the family.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TierPolicy {
/// Python's `redis_batch_cache_expiry`: a key memory misses is read from Redis at most this often.
pub refresh_after: Duration,
/// How long a value read from Redis stays in memory.
pub backfill_ttl: Duration,
pub on_breaker_open: OnBreakerOpen,
/// Python's `default_max_redis_batch_cache_size`: how many keys the refresh throttle tracks.
pub tracked_reads: usize,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum WritePolicy {
#[default]
Both,
LocalOnly,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum RemoteFailurePolicy {
#[default]
Propagate,
UseLocal,
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OnBreakerOpen {
/// Serve what memory holds and let the next read try Redis again.
ServeMemory,
}
/// The memory tier in front of Redis that a read-through family builds for itself, as Python's
/// `DualCache` with a `RedisCache` does.
pub struct DualCache<L1, L2> {
l1: Arc<L1>,
l2: Arc<L2>,
read_policy: ReadPolicy,
write_policy: WritePolicy,
remote_failure_policy: RemoteFailurePolicy,
promotion_ttl: Option<Duration>,
delete_batch_size: usize,
memory: L1,
remote: Option<L2>,
policy: TierPolicy,
reads: Mutex<RecentReads>,
updates: Mutex<()>,
now: Clock,
}
/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const DEFAULT_DELETE_BATCH_SIZE: usize = 1000;
impl<L1, L2> DualCache<L1, L2> {
pub fn new(l1: Arc<L1>, l2: Arc<L2>) -> Self {
Self {
l1,
l2,
read_policy: ReadPolicy::default(),
write_policy: WritePolicy::default(),
remote_failure_policy: RemoteFailurePolicy::default(),
promotion_ttl: None,
delete_batch_size: DEFAULT_DELETE_BATCH_SIZE,
}
}
pub fn with_delete_batch_size(self, delete_batch_size: usize) -> Self {
Self {
delete_batch_size,
..self
}
}
pub fn with_read_policy(self, read_policy: ReadPolicy) -> Self {
Self {
read_policy,
..self
}
}
pub fn with_write_policy(self, write_policy: WritePolicy) -> Self {
Self {
write_policy,
..self
}
}
pub fn with_remote_failure_policy(self, remote_failure_policy: RemoteFailurePolicy) -> Self {
Self {
remote_failure_policy,
..self
}
}
pub fn with_promotion_ttl(self, promotion_ttl: Duration) -> Self {
Self {
promotion_ttl: Some(promotion_ttl),
..self
}
}
fn reads_remote(&self) -> bool {
self.read_policy == ReadPolicy::LocalThenRemote
}
fn writes_remote(&self) -> bool {
self.write_policy == WritePolicy::Both
}
fn remote<T>(&self, result: Result<T, Error>) -> Result<Option<T>, Error> {
match result {
Ok(value) => Ok(Some(value)),
Err(Error::Unavailable(_))
if self.remote_failure_policy == RemoteFailurePolicy::UseLocal =>
{
Ok(None)
}
Err(error) => Err(error),
}
}
fn promotion_context<C: CacheContext>(&self, context: &C) -> C {
context.with_ttl(self.promotion_ttl.or(context.ttl()))
}
pub fn system_clock() -> Clock {
Arc::new(|| {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
impl<V, C, L1, L2> DualCache<L1, L2>
impl<V, L1> DualCache<L1, RedisClient>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: BaseCache<Value = V, Context = C>,
L2: BaseCache<Value = V, Context = C>,
L1: BaseCache<Value = V, Context = ExactCacheContext> + DeleteCache + FlushCache,
{
fn missing(entries: &[BatchEntry<V>]) -> Vec<usize> {
entries
.iter()
.enumerate()
.filter_map(|(index, entry)| (!matches!(entry, BatchEntry::Hit(_))).then_some(index))
.collect()
pub fn read_through(
memory: L1,
remote: Option<RedisClient>,
policy: TierPolicy,
now: Clock,
) -> Self {
Self {
memory,
remote,
reads: Mutex::new(RecentReads::new(policy.tracked_reads)),
policy,
updates: Mutex::new(()),
now,
}
}
fn merge_batch(
pub fn memory(&self) -> &L1 {
&self.memory
}
pub fn remote(&self) -> Option<&RedisClient> {
self.remote.as_ref()
}
pub fn now(&self) -> Duration {
(self.now)()
}
pub fn set_memory(&self, key: &str, value: V, ttl: Duration) -> Result<(), Error> {
self.memory
.set_cache(key, value, &ExactCacheContext { ttl: Some(ttl) })
}
/// Read-modify-write of one memory entry, atomic against every other update of this tier.
pub fn update_memory(
&self,
key: &str,
ttl: Duration,
update: impl FnOnce(Option<V>) -> V,
) -> Result<V, Error> {
let _guard = lock(&self.updates)?;
let value = update(self.memory.get_cache(key, &ExactCacheContext::default())?);
self.set_memory(key, value.clone(), ttl)?;
Ok(value)
}
/// Python's `DualCache.async_batch_get_cache`: memory first, then one MGET for the misses
/// whose refresh is due, written back to memory. A failed Redis read serves memory, as
/// `RedisCache` swallows it, and only a breaker refusal releases the refresh it reserved.
pub async fn get_many<E: std::fmt::Display>(
&self,
keys: &[String],
context: &C,
mut entries: Vec<BatchEntry<V>>,
missing: Vec<usize>,
remote: Vec<BatchEntry<V>>,
) -> Result<Vec<BatchEntry<V>>, Error> {
if missing.len() != remote.len() {
return Err(Error::unavailable(ShortBatch {
requested: missing.len(),
returned: remote.len(),
}));
}
for (index, entry) in missing.into_iter().zip(remote) {
if let BatchEntry::Hit(value) = &entry {
let promotion_context = self.promotion_context(context);
self.l1
.set_cache(&keys[index], value.clone(), &promotion_context)?;
}
entries[index] = entry;
}
Ok(entries)
}
}
impl<V, C, L1, L2> BaseCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: BaseCache<Value = V, Context = C>,
L2: BaseCache<Value = V, Context = C>,
{
type Value = V;
type Context = C;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
self.l2.get_ttl(context)
}
fn set_cache(&self, key: &str, value: V, context: &C) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.set_cache(key, value.clone(), context))?;
}
self.l1.set_cache(key, value, context)
}
fn get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
if let Some(value) = self.l1.get_cache(key, context)? {
return Ok(Some(value));
}
if !self.reads_remote() {
return Ok(None);
}
let value = self.remote(self.l2.get_cache(key, context))?.flatten();
if let Some(value) = &value {
let promotion_context = self.promotion_context(context);
self.l1.set_cache(key, value.clone(), &promotion_context)?;
}
Ok(value)
}
async fn async_set_cache(&self, key: &str, value: V, context: C) -> Result<(), Error> {
if self.writes_remote() {
self.remote(
self.l2
.async_set_cache(key, value.clone(), context.clone())
.await,
)?;
}
self.l1.async_set_cache(key, value, context).await
}
async fn async_get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
if let Some(value) = self.l1.async_get_cache(key, context).await? {
return Ok(Some(value));
}
if !self.reads_remote() {
return Ok(None);
}
let value = self
.remote(self.l2.async_get_cache(key, context).await)?
.flatten();
if let Some(value) = &value {
self.l1
.async_set_cache(key, value.clone(), self.promotion_context(context))
.await?;
}
Ok(value)
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, V)>,
context: C,
) -> Result<(), Error> {
if self.writes_remote() {
self.remote(
self.l2
.async_set_cache_pipeline(entries.clone(), context.clone())
.await,
)?;
}
self.l1.async_set_cache_pipeline(entries, context).await
}
}
impl<V, C, L1, L2> BatchCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: BatchCache<Value = V, Context = C>,
L2: BatchCache<Value = V, Context = C>,
{
fn batch_get_cache(&self, keys: &[String], context: &C) -> Result<Vec<BatchEntry<V>>, Error> {
let entries = self.l1.batch_get_cache(keys, context)?;
let missing = Self::missing(&entries);
if missing.is_empty() || !self.reads_remote() {
return Ok(entries);
}
let remote_keys = missing
decode: impl Fn(&[u8]) -> Result<V, E>,
) -> Result<Vec<Option<V>>, Error> {
let context = ExactCacheContext::default();
let local = keys
.iter()
.map(|index| keys[*index].clone())
.collect::<Vec<_>>();
match self.remote(self.l2.batch_get_cache(&remote_keys, context))? {
Some(remote) => self.merge_batch(keys, context, entries, missing, remote),
None => Ok(entries),
.map(|key| self.memory.get_cache(key, &context))
.collect::<Result<Vec<_>, _>>()?;
let Some(remote) = &self.remote else {
return Ok(local);
};
let reserved = self.reserve(keys, &local)?;
if reserved.keys.is_empty() {
return Ok(local);
}
let due = reserved.keys.iter().map(String::as_str).collect::<Vec<_>>();
let fetched = match remote.mget::<Vec<u8>>(&due).await {
Ok(fetched) => fetched,
Err(RedisError::BreakerOpen) => {
self.release(reserved)?;
return Ok(local);
}
Err(error) => {
tracing::warn!(%error, "Redis read failed, serving the memory tier");
return Ok(local);
}
};
let decoded = match fetched
.iter()
.map(|raw| raw.as_deref().map(&decode).transpose())
.collect::<Result<Vec<_>, _>>()
{
Ok(decoded) => decoded,
Err(error) => {
tracing::warn!(%error, "a Redis value did not decode, serving the memory tier");
return Ok(local);
}
};
let read = reserved
.keys
.into_iter()
.zip(decoded)
.collect::<HashMap<_, _>>();
for (key, value) in &read {
if let Some(value) = value {
self.set_memory(key, value.clone(), self.policy.backfill_ttl)?;
}
}
Ok(keys
.iter()
.zip(local)
.map(|(key, local)| read.get(key).cloned().unwrap_or(local))
.collect())
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
context: C,
) -> Result<Vec<BatchEntry<V>>, Error> {
let entries = self
.l1
.async_batch_get_cache(keys.clone(), context.clone())
.await?;
let missing = Self::missing(&entries);
if missing.is_empty() || !self.reads_remote() {
return Ok(entries);
}
let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect();
match self.remote(
self.l2
.async_batch_get_cache(remote_keys, context.clone())
.await,
)? {
Some(remote) => self.merge_batch(&keys, &context, entries, missing, remote),
None => Ok(entries),
}
}
}
impl<V, C, L1, L2> DeleteCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: DeleteCache<Value = V, Context = C>,
L2: DeleteCache<Value = V, Context = C>,
{
fn delete_cache(&self, key: &str) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.delete_cache(key))?;
}
self.l1.delete_cache(key)
pub fn delete_memory(&self, key: &str) -> Result<(), Error> {
self.memory.delete_cache(key)
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.async_delete_cache(key).await)?;
}
self.l1.async_delete_cache(key).await
}
}
impl<V, C, L1, L2> FlushCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: FlushCache<Value = V, Context = C>,
L2: FlushCache<Value = V, Context = C>,
{
fn flush_cache(&self) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.flush_cache())?;
}
self.l1.flush_cache()
pub fn flush_memory(&self) -> Result<(), Error> {
self.memory.flush_cache()
}
async fn async_flush_cache(&self) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.async_flush_cache().await)?;
}
self.l1.async_flush_cache().await
}
}
impl<C, L1, L2> DualCache<L1, L2>
where
C: CacheContext,
L1: BaseCache<Value = f64, Context = C>,
{
/// Python's `local_only=True` increment: the local tier alone, read then written back.
fn increment_local(&self, key: &str, amount: f64, context: &C) -> Result<f64, Error> {
let value = self.l1.get_cache(key, context)?.unwrap_or(0.0) + amount;
self.l1.set_cache(key, value, context)?;
Ok(value)
}
}
impl<C, L1, L2> CounterCache for DualCache<L1, L2>
where
C: CacheContext,
L1: BaseCache<Value = f64, Context = C>,
L2: CounterCache<Value = f64, Context = C>,
{
fn increment_cache(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
if !self.writes_remote() {
return self.increment_local(key, amount, &context);
}
let value = self.l2.increment_cache(key, amount, context.clone())?;
self.l1.set_cache(key, value, &context)?;
Ok(value)
}
async fn async_increment(
&self,
key: &str,
amount: f64,
context: C,
refresh_ttl: bool,
) -> Result<f64, Error> {
if !self.writes_remote() {
return self.increment_local(key, amount, &context);
}
let value = self
.l2
.async_increment(key, amount, context.clone(), refresh_ttl)
.await?;
self.l1.async_set_cache(key, value, context).await?;
Ok(value)
}
/// `async_increment_cache_pipeline`, L2-first like single increments: the local tier takes
/// each remote result.
async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error>
where
C: Default,
{
if !self.writes_remote() {
return operations
.iter()
.map(|operation| {
let context = C::default().with_ttl(operation.ttl);
self.increment_local(&operation.key, operation.amount, &context)
fn reserve(&self, keys: &[String], local: &[Option<V>]) -> Result<Reservation, Error> {
let now = self.now();
let mut reads = lock(&self.reads)?;
let due = keys
.iter()
.zip(local)
.filter(|(_, value)| value.is_none())
.filter_map(|(key, _)| {
let previous = reads.get(key);
let due = previous
.is_none_or(|read| now.saturating_sub(read) >= self.policy.refresh_after);
due.then(|| {
reads.insert(key.clone(), now);
(key.clone(), previous)
})
.collect();
}
let values = self.l2.async_increment_pipeline(operations.clone()).await?;
if values.len() != operations.len() {
return Err(Error::unavailable(ShortBatch {
requested: operations.len(),
returned: values.len(),
}));
}
for (operation, value) in operations.iter().zip(&values) {
self.l1
.async_set_cache(&operation.key, *value, C::default().with_ttl(operation.ttl))
.await?;
}
Ok(values)
})
.collect::<Vec<_>>();
let (keys, previous) = due.into_iter().unzip();
Ok(Reservation { keys, previous })
}
}
/// `async_set_cache_sadd`: local set first, then the remote one unless writes stay local.
impl<V, C, S, L1, L2> SetCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
S: Clone + Send + Sync + 'static,
L1: SetCache<Value = V, Context = C, SetValue = S>,
L2: SetCache<Value = V, Context = C, SetValue = S>,
{
type SetValue = S;
type SetResult = ();
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<S>,
ttl: Option<Duration>,
) -> Result<(), Error> {
self.l1
.async_set_cache_sadd(key, values.clone(), ttl)
.await?;
if self.writes_remote() {
self.l2.async_set_cache_sadd(key, values, ttl).await?;
fn release(&self, reservation: Reservation) -> Result<(), Error> {
let mut reads = lock(&self.reads)?;
for (key, previous) in reservation.keys.into_iter().zip(reservation.previous) {
match previous {
Some(read) => reads.insert(key, read),
None => reads.remove(&key),
}
}
Ok(())
}
}
/// `async_delete_cache_keys`: every key leaves the local tier, then the remote tier in chunks
/// of `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, since Redis takes a chunk as one command.
impl<V, C, L1, L2> BulkDeleteCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: DeleteCache<Value = V, Context = C>,
L2: BulkDeleteCache<Value = V, Context = C>,
{
async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
for key in &keys {
self.l1.delete_cache(key)?;
}
let mut deleted = 0;
for chunk in keys.chunks(self.delete_batch_size.max(1)) {
deleted += self.l2.delete_cache_keys(chunk.to_vec()).await?;
}
Ok(deleted)
}
fn lock<T>(mutex: &Mutex<T>) -> Result<MutexGuard<'_, T>, Error> {
Ok(mutex.lock()?)
}
/// `async_get_ttl`: the local TTL, or the remote one when the local tier has none.
impl<V, C, L1, L2> TtlCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: TtlCache<Value = V, Context = C>,
L2: TtlCache<Value = V, Context = C>,
{
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
match self.l1.async_get_ttl(key).await? {
Some(ttl) => Ok(Some(ttl)),
None => self.l2.async_get_ttl(key).await,
}
}
struct Reservation {
keys: Vec<String>,
previous: Vec<Option<Duration>>,
}
impl<V, C, L1, L2> ClaimCache for DualCache<L1, L2>
where
V: Clone + PartialEq + Send + Sync + 'static,
C: CacheContext,
L1: ClaimCache<Value = V, Context = C>,
L2: ClaimCache<Value = V, Context = C>,
{
fn claim_cache(&self, key: &str, candidate: V, eligible: &[V], context: C) -> Result<V, Error> {
match self.remote(
self.l2
.claim_cache(key, candidate.clone(), eligible, context.clone()),
)? {
Some(winner) => {
self.l1.set_cache(key, winner.clone(), &context)?;
Ok(winner)
}
None => self.l1.claim_cache(key, candidate, eligible, context),
/// Python's `LimitedSizeOrderedDict`: insertion-ordered, and any write at capacity drops the
/// oldest entry first, even when the written key is already present.
struct RecentReads {
capacity: usize,
order: VecDeque<String>,
times: HashMap<String, Duration>,
}
impl RecentReads {
fn new(capacity: usize) -> Self {
Self {
capacity,
order: VecDeque::new(),
times: HashMap::new(),
}
}
async fn async_claim_cache(
&self,
key: &str,
candidate: V,
eligible: Vec<V>,
context: C,
) -> Result<V, Error> {
match self.remote(
self.l2
.async_claim_cache(key, candidate.clone(), eligible.clone(), context.clone())
.await,
)? {
Some(winner) => {
self.l1
.async_set_cache(key, winner.clone(), context)
.await?;
Ok(winner)
}
None => {
self.l1
.async_claim_cache(key, candidate, eligible, context)
.await
}
fn get(&self, key: &str) -> Option<Duration> {
self.times.get(key).copied()
}
fn insert(&mut self, key: String, time: Duration) {
if self.times.len() >= self.capacity
&& let Some(oldest) = self.order.pop_front()
{
self.times.remove(&oldest);
}
if self.times.insert(key.clone(), time).is_none() {
self.order.push_back(key);
}
}
fn remove(&mut self, key: &str) {
if self.times.remove(key).is_some() {
self.order.retain(|tracked| tracked != key);
}
}
}

4
litellm-rust/crates/cache/src/lib.rs vendored
View file

@ -20,7 +20,5 @@ pub use capabilities::{
RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache, TtlPipelineCache,
};
pub use codec::{CacheCodec, JsonCodec};
pub use dual::{
DEFAULT_DELETE_BATCH_SIZE, DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy,
};
pub use dual::{Clock, DualCache, OnBreakerOpen, TierPolicy, system_clock};
pub use error::{Error, ShortBatch};

839
litellm-rust/crates/cache/tests/dual.rs vendored
View file

@ -1,839 +0,0 @@
use std::{
collections::HashMap,
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, BulkDeleteCache, ClaimCache, CounterCache, DeleteCache,
DualCache, Error, ExactCacheContext, FlushCache, IncrementOperation, ReadPolicy,
RemoteFailurePolicy, SetCache, TtlCache, WritePolicy,
};
use rstest::{fixture, rstest};
struct TestCache<V> {
value: Mutex<Option<V>>,
fail: bool,
}
impl<V> TestCache<V> {
fn new(value: Option<V>, fail: bool) -> Self {
Self {
value: Mutex::new(value),
fail,
}
}
}
impl<V> BaseCache for TestCache<V>
where
V: Clone + Send + Sync + 'static,
{
type Value = V;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(Duration::from_secs(60)))
}
fn set_cache(&self, _: &str, value: V, _: &ExactCacheContext) -> Result<(), Error> {
*self.value.lock().unwrap() = Some(value);
Ok(())
}
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<V>, Error> {
Ok(self.value.lock().unwrap().clone())
}
}
impl<V> BatchCache for TestCache<V> where V: Clone + Send + Sync + 'static {}
impl<V> DeleteCache for TestCache<V>
where
V: Clone + Send + Sync + 'static,
{
fn delete_cache(&self, _: &str) -> Result<(), Error> {
*self.value.lock().unwrap() = None;
Ok(())
}
}
impl<V> FlushCache for TestCache<V>
where
V: Clone + Send + Sync + 'static,
{
fn flush_cache(&self) -> Result<(), Error> {
*self.value.lock().unwrap() = None;
Ok(())
}
}
impl CounterCache for TestCache<f64> {
fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result<f64, Error> {
if self.fail {
return Err(Error::Unavailable(None));
}
let mut value = self.value.lock().unwrap();
let incremented = value.unwrap_or_default() + amount;
*value = Some(incremented);
Ok(incremented)
}
}
impl<V> ClaimCache for TestCache<V>
where
V: Clone + PartialEq + Send + Sync + 'static,
{
fn claim_cache(
&self,
_: &str,
candidate: V,
eligible: &[V],
_: ExactCacheContext,
) -> Result<V, Error> {
if self.fail {
return Err(Error::Unavailable(None));
}
let mut value = self.value.lock().unwrap();
let winner = match value.as_ref() {
Some(existing) if eligible.is_empty() || eligible.contains(existing) => {
existing.clone()
}
_ => candidate,
};
*value = Some(winner.clone());
Ok(winner)
}
}
#[rstest]
fn failed_l2_increment_leaves_l1_unchanged() {
let l1 = Arc::new(TestCache::new(Some(10.0), false));
let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true)));
assert_eq!(
cache.increment_cache("counter", 2.0, ExactCacheContext::default()),
Err(Error::Unavailable(None))
);
assert_eq!(
l1.get_cache("counter", &ExactCacheContext::default())
.unwrap(),
Some(10.0)
);
}
#[rstest]
fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() {
let l1 = Arc::new(TestCache::new(Some("first".to_string()), false));
let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true)))
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(
cache
.claim_cache(
"affinity",
"second".into(),
&["first".into(), "second".into()],
ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
},
)
.unwrap(),
"first"
);
}
struct SyncPanics(TestCache<String>);
impl BaseCache for SyncPanics {
type Value = String;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
self.0.get_ttl(context)
}
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
panic!("sync L2 write on an async path")
}
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
panic!("sync L2 read on an async path")
}
async fn async_set_cache(
&self,
key: &str,
value: String,
context: ExactCacheContext,
) -> Result<(), Error> {
self.0.set_cache(key, value, &context)
}
async fn async_get_cache(
&self,
key: &str,
context: &ExactCacheContext,
) -> Result<Option<String>, Error> {
self.0.get_cache(key, context)
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, String)>,
context: ExactCacheContext,
) -> Result<(), Error> {
for (key, value) in cache_list {
self.0.set_cache(&key, value, &context)?;
}
Ok(())
}
}
impl BatchCache for SyncPanics {
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
context: ExactCacheContext,
) -> Result<Vec<BatchEntry<String>>, Error> {
assert_eq!(keys, ["missing"]);
Ok(vec![match self.0.get_cache("missing", &context)? {
Some(value) => BatchEntry::Hit(value),
None => BatchEntry::Miss,
}])
}
}
impl DeleteCache for SyncPanics {
fn delete_cache(&self, _: &str) -> Result<(), Error> {
panic!("sync L2 delete on an async path")
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
self.0.delete_cache(key)
}
}
impl FlushCache for SyncPanics {
fn flush_cache(&self) -> Result<(), Error> {
panic!("sync L2 flush on an async path")
}
}
#[rstest]
#[tokio::test]
async fn async_operations_use_the_async_l2_methods() {
let l1 = Arc::new(TestCache::new(None, false));
let cache = DualCache::new(
l1.clone(),
Arc::new(SyncPanics(TestCache::new(
Some("remote".to_string()),
false,
))),
);
let context = ExactCacheContext::default();
assert_eq!(
cache.async_get_cache("missing", &context).await.unwrap(),
Some("remote".into())
);
assert_eq!(
l1.get_cache("missing", &context).unwrap(),
Some("remote".into())
);
l1.delete_cache("missing").unwrap();
assert_eq!(
cache
.async_batch_get_cache(vec!["missing".into()], context.clone())
.await
.unwrap(),
[BatchEntry::Hit("remote".to_string())]
);
cache
.async_set_cache("missing", "written".into(), context.clone())
.await
.unwrap();
cache
.async_set_cache_pipeline(vec![("missing".into(), "piped".into())], context.clone())
.await
.unwrap();
cache.async_delete_cache("missing").await.unwrap();
assert_eq!(
cache.async_get_cache("missing", &context).await.unwrap(),
None
);
}
struct Unavailable;
impl BaseCache for Unavailable {
type Value = String;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
Err(Error::Unavailable(None))
}
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
Err(Error::Unavailable(None))
}
}
impl BatchCache for Unavailable {}
impl DeleteCache for Unavailable {
fn delete_cache(&self, _: &str) -> Result<(), Error> {
Err(Error::Unavailable(None))
}
}
impl FlushCache for Unavailable {
fn flush_cache(&self) -> Result<(), Error> {
Err(Error::Unavailable(None))
}
}
impl ClaimCache for Unavailable {
fn claim_cache(
&self,
_: &str,
_: String,
_: &[String],
_: ExactCacheContext,
) -> Result<String, Error> {
Err(Error::InvalidEntry(None))
}
}
#[rstest]
fn remote_failure_policy_selects_propagation_or_the_local_tier() {
let context = ExactCacheContext::default();
let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable));
assert_eq!(
strict.set_cache("key", "value".into(), &context),
Err(Error::Unavailable(None))
);
assert_eq!(
strict.get_cache("key", &context),
Err(Error::Unavailable(None))
);
let l1 = Arc::new(TestCache::new(None, false));
let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable))
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(degraded.get_cache("key", &context), Ok(None));
degraded.set_cache("key", "value".into(), &context).unwrap();
assert_eq!(
degraded.get_cache("key", &context),
Ok(Some("value".into()))
);
degraded.delete_cache("key").unwrap();
assert_eq!(l1.get_cache("key", &context), Ok(None));
}
#[rstest]
fn claim_fallback_does_not_hide_non_availability_errors() {
let cache = DualCache::new(
Arc::new(TestCache::new(Some("first".to_string()), false)),
Arc::new(Unavailable),
)
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(
cache.claim_cache(
"affinity",
"second".into(),
&[],
ExactCacheContext::default()
),
Err(Error::InvalidEntry(None))
);
}
#[rstest]
fn local_only_policies_never_touch_l2() {
let l2 = Arc::new(TestCache::new(Some("remote".to_string()), false));
let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone())
.with_read_policy(ReadPolicy::LocalOnly)
.with_write_policy(WritePolicy::LocalOnly);
let context = ExactCacheContext::default();
assert_eq!(cache.get_cache("key", &context), Ok(None));
cache.set_cache("key", "local".into(), &context).unwrap();
assert_eq!(l2.get_cache("key", &context), Ok(Some("remote".into())));
}
type Log = Arc<Mutex<Vec<String>>>;
type StoredSet = (Vec<String>, Option<Duration>);
/// A keyed tier that logs every call, so tests can assert which tier ran and in what order.
struct Tier {
name: &'static str,
log: Log,
fail: bool,
counters: Mutex<HashMap<String, (f64, ExactCacheContext)>>,
sets: Mutex<HashMap<String, StoredSet>>,
ttls: HashMap<String, Duration>,
}
impl Tier {
fn new(name: &'static str, log: &Log) -> Self {
Self {
name,
log: log.clone(),
fail: false,
counters: Mutex::default(),
sets: Mutex::default(),
ttls: HashMap::new(),
}
}
fn failing(self) -> Self {
Self { fail: true, ..self }
}
fn with_counter(self, key: &str, value: f64) -> Self {
self.counters
.lock()
.unwrap()
.insert(key.into(), (value, ExactCacheContext::default()));
self
}
fn with_ttl(mut self, key: &str, seconds: u64) -> Self {
self.ttls.insert(key.into(), Duration::from_secs(seconds));
self
}
fn record(&self, event: String) {
self.log
.lock()
.unwrap()
.push(format!("{} {event}", self.name));
}
fn counter(&self, key: &str) -> Option<(f64, ExactCacheContext)> {
self.counters.lock().unwrap().get(key).cloned()
}
fn check(&self) -> Result<(), Error> {
if self.fail {
return Err(Error::Unavailable(None));
}
Ok(())
}
}
impl BaseCache for Tier {
type Value = f64;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(&self, key: &str, value: f64, context: &ExactCacheContext) -> Result<(), Error> {
self.check()?;
self.record(format!("set {key}={value}"));
self.counters
.lock()
.unwrap()
.insert(key.into(), (value, context.clone()));
Ok(())
}
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<f64>, Error> {
Ok(self.counter(key).map(|(value, _)| value))
}
}
impl CounterCache for Tier {
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
self.check()?;
let mut counters = self.counters.lock().unwrap();
let value = counters.get(key).map_or(0.0, |(value, _)| *value) + amount;
counters.insert(key.into(), (value, context));
Ok(value)
}
async fn async_increment(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
refresh_ttl: bool,
) -> Result<f64, Error> {
self.record(format!(
"increment {key}+{amount} refresh_ttl={refresh_ttl}"
));
self.increment_cache(key, amount, context)
}
async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error> {
let keys = operations
.iter()
.map(|operation| operation.key.as_str())
.collect::<Vec<_>>();
self.record(format!("pipeline {}", keys.join(",")));
operations
.into_iter()
.map(|operation| {
self.increment_cache(
&operation.key,
operation.amount,
ExactCacheContext { ttl: operation.ttl },
)
})
.collect()
}
}
impl DeleteCache for Tier {
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.check()?;
self.record(format!("delete {key}"));
self.counters.lock().unwrap().remove(key);
Ok(())
}
}
impl BulkDeleteCache for Tier {
async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
self.check()?;
self.record(format!("delete_keys {}", keys.join(",")));
let mut counters = self.counters.lock().unwrap();
Ok(keys
.iter()
.filter(|key| counters.remove(key.as_str()).is_some())
.count())
}
}
impl SetCache for Tier {
type SetValue = String;
type SetResult = ();
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<String>,
ttl: Option<Duration>,
) -> Result<(), Error> {
self.check()?;
self.record(format!("sadd {key} {}", values.join(",")));
self.sets.lock().unwrap().insert(key.into(), (values, ttl));
Ok(())
}
}
impl TtlCache for Tier {
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
self.check()?;
self.record(format!("ttl {key}"));
Ok(self.ttls.get(key).copied())
}
}
#[fixture]
fn log() -> Log {
Log::default()
}
fn events(log: &Log) -> Vec<String> {
log.lock().unwrap().clone()
}
fn seconds(ttl: u64) -> ExactCacheContext {
ExactCacheContext {
ttl: Some(Duration::from_secs(ttl)),
}
}
#[rstest]
#[case::window_semantics(false)]
#[case::refresh_on_every_write(true)]
#[tokio::test]
async fn async_increment_passes_refresh_ttl_to_l2_and_stores_its_result_locally(
log: Log,
#[case] refresh_ttl: bool,
) {
let l1 = Arc::new(Tier::new("l1", &log).with_counter("counter", 1.0));
let l2 = Arc::new(Tier::new("l2", &log).with_counter("counter", 10.0));
let cache = DualCache::new(l1.clone(), l2.clone());
assert_eq!(
cache
.async_increment("counter", 2.0, seconds(30), refresh_ttl)
.await,
Ok(12.0)
);
assert_eq!(
events(&log),
[
format!("l2 increment counter+2 refresh_ttl={refresh_ttl}"),
"l1 set counter=12".into(),
]
);
assert_eq!(l1.counter("counter"), Some((12.0, seconds(30))));
assert_eq!(l2.counter("counter"), Some((12.0, seconds(30))));
}
#[rstest]
#[tokio::test]
async fn failed_async_l2_increment_leaves_l1_unchanged(log: Log) {
let l1 = Arc::new(Tier::new("l1", &log).with_counter("counter", 1.0));
let cache = DualCache::new(l1.clone(), Arc::new(Tier::new("l2", &log).failing()))
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(
cache
.async_increment("counter", 2.0, seconds(30), true)
.await,
Err(Error::Unavailable(None))
);
assert_eq!(
l1.counter("counter"),
Some((1.0, ExactCacheContext::default()))
);
}
#[rstest]
#[case::empty(Vec::new(), Vec::new())]
#[case::one_key(vec![("a", 1.0, Some(10))], vec![6.0])]
#[case::repeated_and_mixed_ttls(
vec![("a", 1.0, Some(10)), ("b", 2.0, None), ("a", 3.0, Some(20))],
vec![6.0, 2.0, 9.0],
)]
#[tokio::test]
async fn async_increment_pipeline_runs_l2_first_and_l1_takes_each_remote_result(
log: Log,
#[case] operations: Vec<(&str, f64, Option<u64>)>,
#[case] expected: Vec<f64>,
) {
let operations = operations
.into_iter()
.map(|(key, amount, ttl)| IncrementOperation {
key: key.into(),
amount,
ttl: ttl.map(Duration::from_secs),
})
.collect::<Vec<_>>();
let l1 = Arc::new(Tier::new("l1", &log).with_counter("a", 100.0));
let l2 = Arc::new(Tier::new("l2", &log).with_counter("a", 5.0));
let cache = DualCache::new(l1.clone(), l2);
assert_eq!(
cache.async_increment_pipeline(operations.clone()).await,
Ok(expected.clone())
);
let keys = operations
.iter()
.map(|operation| operation.key.as_str())
.collect::<Vec<_>>();
let mut expected_events = vec![format!("l2 pipeline {}", keys.join(","))];
expected_events.extend(
operations
.iter()
.zip(&expected)
.map(|(operation, value)| format!("l1 set {}={value}", operation.key)),
);
assert_eq!(events(&log), expected_events);
if let Some((operation, value)) = operations.iter().zip(&expected).next_back() {
assert_eq!(
l1.counter(&operation.key),
Some((*value, ExactCacheContext { ttl: operation.ttl }))
);
}
}
#[rstest]
#[tokio::test]
async fn failed_l2_increment_pipeline_leaves_l1_unchanged(log: Log) {
let l1 = Arc::new(Tier::new("l1", &log).with_counter("a", 1.0));
let cache = DualCache::new(l1.clone(), Arc::new(Tier::new("l2", &log).failing()));
assert_eq!(
cache
.async_increment_pipeline(vec![IncrementOperation {
key: "a".into(),
amount: 1.0,
ttl: None,
}])
.await,
Err(Error::Unavailable(None))
);
assert_eq!(events(&log), ["l2 pipeline a"]);
assert_eq!(l1.counter("a"), Some((1.0, ExactCacheContext::default())));
}
#[rstest]
#[case::both_tiers(WritePolicy::Both, &["l1 sadd members a,b", "l2 sadd members a,b"])]
#[case::local_only(WritePolicy::LocalOnly, &["l1 sadd members a,b"])]
#[tokio::test]
async fn set_add_writes_locally_then_remotely_unless_local_only(
log: Log,
#[case] write_policy: WritePolicy,
#[case] expected: &[&str],
) {
let l1 = Arc::new(Tier::new("l1", &log));
let l2 = Arc::new(Tier::new("l2", &log));
let cache = DualCache::new(l1.clone(), l2.clone()).with_write_policy(write_policy);
let ttl = Some(Duration::from_secs(45));
cache
.async_set_cache_sadd("members", vec!["a".into(), "b".into()], ttl)
.await
.unwrap();
assert_eq!(events(&log), expected);
let stored = Some((vec!["a".to_string(), "b".to_string()], ttl));
assert_eq!(l1.sets.lock().unwrap().get("members").cloned(), stored);
assert_eq!(
l2.sets.lock().unwrap().get("members").cloned(),
stored.filter(|_| write_policy == WritePolicy::Both)
);
}
#[rstest]
#[tokio::test]
async fn failed_local_set_add_never_reaches_l2(log: Log) {
let cache = DualCache::new(
Arc::new(Tier::new("l1", &log).failing()),
Arc::new(Tier::new("l2", &log)),
);
assert_eq!(
cache
.async_set_cache_sadd("members", vec!["a".into()], None)
.await,
Err(Error::Unavailable(None))
);
assert!(events(&log).is_empty());
}
#[rstest]
#[case::empty(None, &[], &[])]
#[case::default_batch_size(None, &["a", "b", "c"], &["a,b,c"])]
#[case::chunked(Some(2), &["a", "b", "c", "d", "e"], &["a,b", "c,d", "e"])]
#[case::exact_chunks(Some(2), &["a", "b", "c", "d"], &["a,b", "c,d"])]
#[case::zero_means_one_per_chunk(Some(0), &["a", "b"], &["a", "b"])]
#[tokio::test]
async fn bulk_delete_removes_every_key_locally_then_remotely_in_chunks(
log: Log,
#[case] batch_size: Option<usize>,
#[case] keys: &[&str],
#[case] chunks: &[&str],
) {
let l1 = Tier::new("l1", &log);
let l2 = Tier::new("l2", &log);
for key in keys.iter().step_by(2) {
l2.counters
.lock()
.unwrap()
.insert((*key).into(), (1.0, ExactCacheContext::default()));
}
let l1 = Arc::new(l1);
let mut cache = DualCache::new(l1, Arc::new(l2));
if let Some(batch_size) = batch_size {
cache = cache.with_delete_batch_size(batch_size);
}
assert_eq!(
cache
.delete_cache_keys(keys.iter().map(|key| (*key).into()).collect())
.await,
Ok(keys.len().div_ceil(2))
);
let expected = keys
.iter()
.map(|key| format!("l1 delete {key}"))
.chain(chunks.iter().map(|chunk| format!("l2 delete_keys {chunk}")))
.collect::<Vec<_>>();
assert_eq!(events(&log), expected);
}
#[rstest]
#[tokio::test]
async fn bulk_delete_stops_before_l2_when_the_local_delete_fails(log: Log) {
let cache = DualCache::new(
Arc::new(Tier::new("l1", &log).failing()),
Arc::new(Tier::new("l2", &log)),
);
assert_eq!(
cache.delete_cache_keys(vec!["a".into()]).await,
Err(Error::Unavailable(None))
);
assert!(events(&log).is_empty());
}
#[rstest]
#[case::local_hit("both", Some(10), &["l1 ttl both"])]
#[case::remote_fallback("remote", Some(20), &["l1 ttl remote", "l2 ttl remote"])]
#[case::missing_everywhere("missing", None, &["l1 ttl missing", "l2 ttl missing"])]
#[tokio::test]
async fn ttl_reads_local_then_remote(
log: Log,
#[case] key: &str,
#[case] expected: Option<u64>,
#[case] expected_events: &[&str],
) {
let cache = DualCache::new(
Arc::new(Tier::new("l1", &log).with_ttl("both", 10)),
Arc::new(
Tier::new("l2", &log)
.with_ttl("both", 99)
.with_ttl("remote", 20),
),
);
assert_eq!(
cache.async_get_ttl(key).await,
Ok(expected.map(Duration::from_secs))
);
assert_eq!(events(&log), expected_events);
}
/// Python `local_only=True`: the increment and the pipeline stay on the local tier.
#[rstest]
#[tokio::test]
async fn local_only_writes_increment_the_local_tier_alone(log: Log) {
let l1 = Arc::new(Tier::new("l1", &log).with_counter("a", 1.0));
let l2 = Arc::new(Tier::new("l2", &log).with_counter("a", 10.0));
let cache = DualCache::new(l1.clone(), l2.clone()).with_write_policy(WritePolicy::LocalOnly);
assert_eq!(cache.increment_cache("a", 2.0, seconds(30)), Ok(3.0));
assert_eq!(
cache.async_increment("a", 1.0, seconds(30), true).await,
Ok(4.0)
);
let operations = vec![
IncrementOperation {
key: "a".into(),
amount: 1.0,
ttl: Some(Duration::from_secs(10)),
},
IncrementOperation {
key: "b".into(),
amount: 2.0,
ttl: None,
},
];
assert_eq!(
cache.async_increment_pipeline(operations).await,
Ok(vec![5.0, 2.0])
);
assert_eq!(
events(&log),
["l1 set a=3", "l1 set a=4", "l1 set a=5", "l1 set b=2"]
);
assert_eq!(l2.counter("a"), Some((10.0, ExactCacheContext::default())));
}

16
litellm-rust/crates/config/AGENTS.md
View file

@ -3,7 +3,7 @@
- `RedisSettings::resolve` follows `litellm/_redis.py` `_get_redis_client_logic` plus `_credential_provider_auth_kwargs` and the per-client allowlists of `get_redis_async_client` (redis-py 5.3.1). Python and Rust pods share one Redis during the switch, so a change there is a change to a wire contract
- `tests/redis_resolution.json` is that contract. Generate and check it with redis-py 5.3.1. The cross-version compatibility suite skips only this version-specific snapshot check, since later redis-py releases change defaults such as TLS hostname verification. `tests/generate_redis_resolution.py` builds the client and its actual connection objects without connecting, reading their TLS contexts, target fields and effective timeouts. Rejections retain their exception type and message, and Rust must return the matching error variant and setting. `tests/unit/test_redis.py` detects Python drift and `resolution_matches_python_redis` replays the contract here. New resolution rules get generator cases, not hand-written Rust expectations
- `resolve` takes a `SecretSnapshot` and an environment `Lookup`. The snapshot holds every name `RedisSettings::secret_names` lists, read the way Python's `get_secret` reads it (secret manager, then environment). The host fills it once before resolving, python-bridge through `get_secret` and rust-proxy from its environment, so the resolver never calls back into the host. The environment lookup is only for what Python reads straight from `os.environ` (`REDIS_USERNAME` for Azure). Reading a name the snapshot does not list fails a debug assertion
- Every Redis setting Python reads is one row in `SETTINGS` in `src/redis/keys.rs`: how Rust handles it (read, ignore, warn, reject), whether `_get_redis_env_kwarg_mapping` reads it as `REDIS_<KWARG>`, and the names `_get_redis_client_logic` falls back to. The env layer, the rejections, the warnings and `secret_names` all come from that table, so a new setting is one row. It is the shape the `#[derive(Layered)]` attributes in `litellm-layered` would generate, and moves onto them once that crate lands
- Every Redis setting Python reads is one row in `SETTINGS` in `src/redis/keys.rs`: how Rust handles it (read, or warn and drop), whether `_get_redis_env_kwarg_mapping` reads it as `REDIS_<KWARG>`, and the names `_get_redis_client_logic` falls back to. The env layer, the warnings and `secret_names` all come from that table, so a new setting is one row. It is the shape the `#[derive(Layered)]` attributes in `litellm-layered` would generate, and moves onto them once that crate lands
- Lookups go through `litellm_core_utils::settings::Lookup`, never a closure type of our own
- A url carries its own TLS. redis-py's url allowlist is the plain `Connection` signature, so every `ssl_*` setting outside the url is dropped, and a credential provider strips the url's userinfo. A `rediss://` url's `ssl_cert_reqs`, `ssl_ca_certs`, `ssl_certfile`, `ssl_keyfile` and `ssl_check_hostname` query parameters are read like `parse_url` reads them
- TLS options and `db` are only read for the endpoint that uses them, so a bad `ssl_cert_reqs` never fails a url target. A cluster with a non-zero `db`, and Sentinel with a `host` or `port`, are errors because redis-py raises building those clients
@ -11,9 +11,11 @@
- Precedence matches Python. The target is cluster, then URL, then Sentinel (`sentinel_nodes` plus `service_name`), then host. Auth is an explicit `credential_provider`, then GCP, then Azure, then AWS, and the first one wins with a warning
- A value `int()` or `float()` cannot parse (`port: abc`, `socket_timeout: slow`) is dropped with a warning and the default is used, like `_coerce_redis_kwargs_types`. A number that parses but is out of range (negative seconds, a port above 65535) fails with `RedisError::Invalid`
- `validate_coordination` mirrors the proxy's `CoordinationRedisParams` pydantic model, which rejects a malformed `port` or `ssl` in `general_settings.coordination_redis` before redis-py coerces anything. Hosts call it for that block only
- Known divergences, kept on purpose. Each is a case in `redis_resolution.json` and an entry in `DIVERGENCES` in `tests/redis.rs`
- A url query parameter other than the TLS ones, or any query parameter on a plain `redis://` url, is `RedisError::Unsupported` naming `url ?<name>`. redis-py turns each one into a client kwarg
- Settings that would change the server, identity or TLS trust if ignored are `RedisError::Unsupported` naming the setting, including `ssl_ciphers`, `ssl_min_version`, `ssl_password`, `ssl_ca_data`, OCSP options, `unix_socket_path`, `connection_pool`, `event_dispatcher`, `redis_connect_func`, `retry`
- `ssl_ca_path` is rejected on both sides. redis-py 5.3.1's Redis constructor raises TypeError and Rust returns Unsupported naming the setting
- `retry_on_timeout`, `retry_on_error`, `protocol`, `cache`, `cache_config`, `health_check_interval`, `socket_keepalive` and `max_connections` are accepted with a warning and ignored. Cluster clients still turn TCP keepalive on, like Python's async cluster default, but redis-rs has no periodic health check, so Python's 25s cluster default has no equivalent
- A boolean `REDIS_RETRY_ON_ERROR` reaches redis-py's connection constructor as a non-iterable and raises TypeError. Rust ignores it with a warning and otherwise resolves the default host connection, verified against the host-default fixture
- The resolver never fails on a setting it cannot honor. A `Warn` row is dropped with a warning that names its source (`ssl_ciphers` or `REDIS_SSL_CIPHERS`) and the consequence from `keys::consequence`, because a Python deployment that worked must keep working when its stores move to Rust. Only a value that is malformed (`RedisError::Invalid`) or a target that is missing fails. Promoting a `Warn` row to `Read` is then a pure improvement with no behavior cliff
- `ssl_ciphers`, `ssl_min_version`, `ssl_password`, the OCSP options, `connection_pool`, `event_dispatcher`, `redis_connect_func` and `retry` are warned and dropped. Ignoring OCSP connects without revocation checks, and the warning says so
- `ssl_ca_data` and `ssl_ca_path` are read into `RedisTls` and trusted next to `ssl_ca_certs`. `unix_socket_path` is read into `RedisHost` and, like redis-py's `UnixDomainSocketConnection`, replaces host, port and TLS. Python still requires `host` next to it, so Rust does too
- A url query parameter other than the TLS ones, or any query parameter on a plain `redis://` url, is warned and dropped naming `url ?<name>`. redis-py turns each one into a client kwarg
- `retry_on_timeout`, `retry_on_error`, `protocol`, `cache`, `cache_config`, `socket_keepalive` and `max_connections` are warned and dropped. Cluster clients still turn TCP keepalive on, like Python's async cluster default
- `health_check_interval` is read as an `int` of seconds like `_coerce_redis_kwargs_types`, zero meaning off. `redis-connect` gives cluster clients Python's 25s default and `redis-client` pings every node before the next command once the client sat idle that long, outside the breaker and the call events
- A `credential_provider` object resolves to `RedisAuth::Host`, the way Python's `{"kind": "host"}` does. `redis-connect` has no host credential source yet, so python-bridge leaves that family on memory with a warning, and rust-proxy never sees one because yaml cannot carry an object
- Cases where Rust and Python disagree on purpose are matched explicitly in `tests/redis.rs` against the generated Python outcome. `ssl_ca_path` is one, redis-py 5.3.1's async `Redis` has no such parameter and raises `TypeError` building the client, while Rust trusts the directory like `load_verify_locations(capath=...)`

28
litellm-rust/crates/config/src/redis.rs
View file

@ -123,6 +123,8 @@ pub struct ResolvedRedis {
pub namespace: Option<String>,
pub timeout: Option<Duration>,
pub connect_timeout: Option<Duration>,
/// redis-py's `health_check_interval`, which `_coerce_redis_kwargs_types` reads as an `int`
pub health_check_interval: Option<Duration>,
}
#[derive(Clone)]
@ -182,7 +184,6 @@ impl RedisSettings {
.filter(|(key, _)| !(key == "url" && drop_environment_url))
.chain(block)
.collect();
reject_unsupported(&merged)?;
warn_ignored(&merged);
let startup_nodes = match configured(&merged, "startup_nodes") {
@ -247,6 +248,7 @@ impl RedisSettings {
port: integer(&merged, "port")?.unwrap_or(6379),
tls: tls()?,
db: db()?,
unix_socket_path: text_value(&merged, "unix_socket_path").map(Into::into),
}),
},
};
@ -263,6 +265,9 @@ impl RedisSettings {
namespace: text_value(&merged, "namespace"),
timeout: seconds(&merged, "socket_timeout")?,
connect_timeout: seconds(&merged, "socket_connect_timeout")?,
health_check_interval: integer::<i64>(&merged, "health_check_interval")?
.filter(|seconds| *seconds > 0)
.map(|seconds| Duration::from_secs(seconds.unsigned_abs())),
})
}
@ -383,15 +388,6 @@ pub fn environment_has_target(lookup: &dyn Lookup) -> bool {
.any(|name| lookup.get(name).is_some())
}
fn reject_unsupported(merged: &Kwargs) -> Result<(), RedisError> {
match merged.iter().find(|(key, _)| is_unsupported(key)) {
Some((_, entry)) => Err(RedisError::Unsupported {
setting: entry.source.clone(),
}),
None => Ok(()),
}
}
fn invalid(setting: &str, expected: Expected) -> RedisError {
RedisError::Invalid {
setting: setting.to_owned(),
@ -449,17 +445,17 @@ fn coerced<N: std::str::FromStr + serde::de::DeserializeOwned, T>(
}
fn warn_ignored(merged: &Kwargs) {
for key in merged.keys() {
for (key, entry) in merged {
if keys::setting(key).is_some_and(|setting| setting.handling == Handling::Warn) {
tracing::warn!(setting = %key, "Redis setting is accepted but ignored");
tracing::warn!(
"Redis setting {} is not supported by the Rust client and is ignored: {}",
entry.source,
keys::consequence(key)
);
}
}
}
fn is_unsupported(key: &str) -> bool {
keys::setting(key).is_some_and(|setting| setting.handling == Handling::Reject)
}
fn environment_kwargs(lookup: &dyn Lookup) -> Kwargs {
keys::SETTINGS
.iter()

2
litellm-rust/crates/config/src/redis/error.rs
View file

@ -1,7 +1,5 @@
#[derive(Debug, thiserror::Error)]
pub enum RedisError {
#[error("Redis host does not yet support configured setting {setting}")]
Unsupported { setting: String },
#[error("invalid Redis setting {setting}: expected {expected}")]
Invalid { setting: String, expected: Expected },
#[error("either host or url must be specified for redis")]

75
litellm-rust/crates/config/src/redis/keys.rs
View file

@ -1,14 +1,13 @@
use litellm_core_utils::settings::Lookup;
// Every Redis setting litellm/_redis.py reads, with how Rust handles it and where else Python looks for it.
// Warn rows are read by Python and dropped here with a warning naming the consequence, never an error.
// environment marks the kwargs _get_redis_env_kwarg_mapping reads as REDIS_<KWARG>, and fallback lists the names
// _get_redis_client_logic consults with get_secret when the merged kwargs lack the setting
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Handling {
Read,
Ignore,
Warn,
Reject,
}
pub(super) struct Setting {
@ -76,16 +75,16 @@ pub(super) const SETTINGS: &[Setting] = &[
),
Setting::new("cache", Handling::Warn, true, &[]),
Setting::new("cache_config", Handling::Warn, true, &[]),
Setting::new("charset", Handling::Ignore, true, &[]),
Setting::new("client_name", Handling::Ignore, true, &[]),
Setting::new("connection_pool", Handling::Reject, false, &[]),
Setting::new("charset", Handling::Warn, true, &[]),
Setting::new("client_name", Handling::Warn, true, &[]),
Setting::new("connection_pool", Handling::Warn, false, &[]),
Setting::new("credential_provider", Handling::Read, false, &[]),
Setting::new("db", Handling::Read, true, &[]),
Setting::new("decode_responses", Handling::Ignore, true, &[]),
Setting::new("encoding", Handling::Ignore, true, &[]),
Setting::new("encoding_errors", Handling::Ignore, true, &[]),
Setting::new("errors", Handling::Ignore, true, &[]),
Setting::new("event_dispatcher", Handling::Reject, true, &[]),
Setting::new("decode_responses", Handling::Warn, true, &[]),
Setting::new("encoding", Handling::Warn, true, &[]),
Setting::new("encoding_errors", Handling::Warn, true, &[]),
Setting::new("errors", Handling::Warn, true, &[]),
Setting::new("event_dispatcher", Handling::Warn, true, &[]),
Setting::new(
"gcp_service_account",
Handling::Read,
@ -98,17 +97,17 @@ pub(super) const SETTINGS: &[Setting] = &[
true,
&["REDIS_GCP_SSL_CA_CERTS"],
),
Setting::new("health_check_interval", Handling::Warn, true, &[]),
Setting::new("health_check_interval", Handling::Read, true, &[]),
Setting::new("host", Handling::Read, true, &[]),
Setting::new("lib_name", Handling::Ignore, true, &[]),
Setting::new("lib_version", Handling::Ignore, true, &[]),
Setting::new("lib_name", Handling::Warn, true, &[]),
Setting::new("lib_version", Handling::Warn, true, &[]),
Setting::new("max_connections", Handling::Warn, true, &[]),
Setting::new("namespace", Handling::Read, false, &[]),
Setting::new("password", Handling::Read, true, &[]),
Setting::new("port", Handling::Read, true, &[]),
Setting::new("protocol", Handling::Warn, true, &[]),
Setting::new("redis_connect_func", Handling::Reject, true, &[]),
Setting::new("retry", Handling::Reject, false, &[]),
Setting::new("redis_connect_func", Handling::Warn, true, &[]),
Setting::new("retry", Handling::Warn, false, &[]),
Setting::new("retry_on_error", Handling::Warn, true, &[]),
Setting::new("retry_on_timeout", Handling::Warn, true, &[]),
Setting::new(
@ -129,37 +128,61 @@ pub(super) const SETTINGS: &[Setting] = &[
false,
&["REDIS_SERVICE_NAME"],
),
Setting::new("single_connection_client", Handling::Ignore, true, &[]),
Setting::new("single_connection_client", Handling::Warn, true, &[]),
Setting::new("socket_connect_timeout", Handling::Read, true, &[]),
Setting::new("socket_keepalive", Handling::Warn, true, &[]),
Setting::new("socket_keepalive_options", Handling::Warn, true, &[]),
Setting::new("socket_timeout", Handling::Read, true, &[]),
Setting::new("ssl", Handling::Read, true, &[]),
Setting::new("ssl_ca_certs", Handling::Read, true, &[]),
Setting::new("ssl_ca_data", Handling::Reject, true, &[]),
Setting::new("ssl_ca_path", Handling::Reject, true, &[]),
Setting::new("ssl_ca_data", Handling::Read, true, &[]),
Setting::new("ssl_ca_path", Handling::Read, true, &[]),
Setting::new("ssl_cert_reqs", Handling::Read, true, &[]),
Setting::new("ssl_certfile", Handling::Read, true, &[]),
Setting::new("ssl_check_hostname", Handling::Read, true, &[]),
Setting::new("ssl_ciphers", Handling::Reject, true, &[]),
Setting::new("ssl_ciphers", Handling::Warn, true, &[]),
Setting::new("ssl_keyfile", Handling::Read, true, &[]),
Setting::new("ssl_min_version", Handling::Reject, true, &[]),
Setting::new("ssl_ocsp_context", Handling::Reject, true, &[]),
Setting::new("ssl_ocsp_expected_cert", Handling::Reject, true, &[]),
Setting::new("ssl_password", Handling::Reject, true, &[]),
Setting::new("ssl_validate_ocsp", Handling::Reject, true, &[]),
Setting::new("ssl_validate_ocsp_stapled", Handling::Reject, true, &[]),
Setting::new("ssl_min_version", Handling::Warn, true, &[]),
Setting::new("ssl_ocsp_context", Handling::Warn, true, &[]),
Setting::new("ssl_ocsp_expected_cert", Handling::Warn, true, &[]),
Setting::new("ssl_password", Handling::Warn, true, &[]),
Setting::new("ssl_validate_ocsp", Handling::Warn, true, &[]),
Setting::new("ssl_validate_ocsp_stapled", Handling::Warn, true, &[]),
Setting::new(
"startup_nodes",
Handling::Read,
false,
&["REDIS_CLUSTER_NODES"],
),
Setting::new("unix_socket_path", Handling::Reject, true, &[]),
Setting::new("unix_socket_path", Handling::Read, true, &[]),
Setting::new("url", Handling::Read, true, &[]),
Setting::new("username", Handling::Read, true, &[]),
];
/// What ignoring a warned setting means for the connection, so the operator reads the consequence
/// rather than only the key name.
pub(super) fn consequence(kwarg: &str) -> &'static str {
match kwarg {
"ssl_validate_ocsp"
| "ssl_validate_ocsp_stapled"
| "ssl_ocsp_context"
| "ssl_ocsp_expected_cert" => {
"rustls has no client-side OCSP, so the connection is made without revocation checks"
}
"ssl_ciphers" | "ssl_min_version" => {
"rustls negotiates TLS 1.2+ with its own ECDHE AEAD suites"
}
"ssl_password" => "an encrypted client key cannot be read, so the TLS handshake fails",
"credential_provider" | "redis_connect_func" => {
"the Python object is not callable from Rust, so static username and password are used"
}
"connection_pool" | "retry" | "event_dispatcher" => {
"the Rust client uses its own pool, retries and events"
}
_ => "the Rust client has no equivalent",
}
}
pub(super) fn setting(kwarg: &str) -> Option<&'static Setting> {
SETTINGS.iter().find(|setting| setting.kwarg == kwarg)
}

22
litellm-rust/crates/config/src/redis/tls.rs
View file

@ -4,10 +4,12 @@ use super::{Entry, Expected, Kwargs, RedisError, bool_value, parsed, text_value}
use crate::Value;
// redis-py's parse_url turns every query parameter into a client kwarg. Only the TLS ones have a redis-rs
// equivalent, so any other parameter, and any parameter on a plain redis:// url, is rejected
// equivalent, so any other parameter, and any parameter on a plain redis:// url, is warned and dropped
const URL_TLS_SETTINGS: &[&str] = &[
"ssl_cert_reqs",
"ssl_ca_certs",
"ssl_ca_data",
"ssl_ca_path",
"ssl_certfile",
"ssl_keyfile",
"ssl_check_hostname",
@ -35,6 +37,8 @@ pub(super) fn tls_options(
.map(Into::into),
certfile: text_value(merged, "ssl_certfile").map(Into::into),
keyfile: text_value(merged, "ssl_keyfile").map(Into::into),
ca_data: text_value(merged, "ssl_ca_data"),
ca_path: text_value(merged, "ssl_ca_path").map(Into::into),
verify_certificate,
check_hostname: bool_value(merged, "ssl_check_hostname"),
})
@ -57,14 +61,16 @@ pub(super) fn url_tls(url: &str) -> Result<TlsOptions, RedisError> {
)
})
.collect();
if let Some((_, entry)) = query
.iter()
.find(|(name, _)| !tls_scheme || !URL_TLS_SETTINGS.contains(&name.as_str()))
{
return Err(RedisError::Unsupported {
setting: entry.source.clone(),
});
let (honored, ignored): (Kwargs, Kwargs) = query
.into_iter()
.partition(|(name, _)| tls_scheme && URL_TLS_SETTINGS.contains(&name.as_str()));
for entry in ignored.values() {
tracing::warn!(
"Redis setting {} is not supported by the Rust client and is ignored",
entry.source
);
}
let query = honored;
Ok(TlsOptions {
check_hostname: text_value(&query, "ssl_check_hostname")
.is_some_and(|text| url_bool(&text)),

24
litellm-rust/crates/config/tests/generate_redis_resolution.py
View file

@ -258,6 +258,25 @@ CASES: Final[tuple[tuple[str, Mapping[str, object], Mapping[str, str]], ...]] =
("protocol_3", {"host": "redis", "protocol": 3}, {}),
("retry_on_error_from_the_environment", {"host": "redis"}, {"REDIS_RETRY_ON_ERROR": "true"}),
("unix_socket_path", {"host": "redis", "unix_socket_path": "/tmp/redis.sock"}, {}),
(
"unix_socket_path_keeps_tls_settings_out_of_the_target",
{"host": "redis", "ssl": True, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"},
{},
),
(
"ssl_ca_data",
{"host": "redis", "ssl": True, "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----"},
{},
),
(
"ssl_ca_path_from_the_environment",
{"host": "redis"},
{"REDIS_SSL": "true", "REDIS_SSL_CA_PATH": "/etc/ssl/private-cas"},
),
("health_check_interval_is_whole_seconds", {"health_check_interval": "7", "host": "redis"}, {}),
("health_check_interval_zero_is_off", {"health_check_interval": 0, "host": "redis"}, {}),
("a_cluster_rejects_a_nonzero_db", {"startup_nodes": [{"host": "node-a", "port": 7000}], "db": 3}, {}),
("unix_socket_path_without_a_host_fails_like_python", {"unix_socket_path": "/tmp/redis.sock"}, {}),
)
@ -280,6 +299,8 @@ def _tls(connection: AbstractConnection) -> dict[str, object] | None:
context: Final = connection.ssl_context
return {
"ca_certs": context.ca_certs,
"ca_data": context.ca_data,
"ca_path": None,
"certfile": context.certfile,
"keyfile": context.keyfile,
"verify_certificate": context.cert_reqs != ssl.CERT_NONE,
@ -318,6 +339,7 @@ def _cluster(client: async_redis.RedisCluster[bytes]) -> dict[str, object]:
"auth": _auth(kwargs),
"socket_timeout": connection.socket_timeout,
"socket_connect_timeout": connection.socket_connect_timeout,
"health_check_interval": float(connection.health_check_interval) or None,
}
@ -338,6 +360,7 @@ def _sentinel(pool: async_redis.SentinelConnectionPool[async_redis.Connection])
"auth": _auth(kwargs),
"socket_timeout": connection.socket_timeout,
"socket_connect_timeout": connection.socket_connect_timeout,
"health_check_interval": float(connection.health_check_interval) or None,
}
@ -359,6 +382,7 @@ def _single(pool: async_redis.ConnectionPool[AbstractConnection], url: object) -
"auth": _auth(kwargs),
"socket_timeout": connection.socket_timeout,
"socket_connect_timeout": connection.socket_connect_timeout,
"health_check_interval": float(connection.health_check_interval) or None,
}

140
litellm-rust/crates/config/tests/redis.rs
View file

@ -2,6 +2,7 @@ use std::{
collections::HashMap,
fmt::Write,
sync::{Arc, Mutex},
time::Duration,
};
use litellm_config::{
@ -135,14 +136,31 @@ fn resolves_redis_through_the_parsed_config(
&[],
"connection_pool"
)]
fn rejected_settings_name_their_source(
fn unsupported_settings_resolve_with_a_warning_naming_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
let events = Arc::new(Mutex::new(Vec::new()));
let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone()));
let resolved = tracing::subscriber::with_default(subscriber, || resolve(yaml, environment));
let warnings = events.lock().unwrap().join(" ");
assert!(matches!(error, RedisError::Unsupported { setting } if setting == expected));
assert_eq!(
resolved.unwrap().endpoint,
Endpoint::Host(litellm_redis_types::HostEndpoint {
host: "redis".into(),
port: 6379,
db: 0,
tls: None,
unix_socket_path: None,
})
);
assert!(
warnings.contains(&format!("Redis setting {expected} is not supported")),
"{warnings}"
);
assert!(warnings.contains("is ignored: "), "{warnings}");
}
fn static_auth(username: Option<&str>, password: Option<&str>) -> RedisAuth {
@ -257,22 +275,22 @@ fn incomplete_aws_iam_auth_fails(#[case] yaml: &str, #[case] message: &str) {
#[case::defaults_verify_the_certificate_but_not_the_hostname(
" ssl: true",
&[],
Some(TlsOptions { ca_certs: None, certfile: None, keyfile: None, verify_certificate: true, check_hostname: false })
Some(TlsOptions { ca_certs: None, ca_data: None, ca_path: None, certfile: None, keyfile: None, verify_certificate: true, check_hostname: false })
)]
#[case::every_option(
" ssl: true\n ssl_ca_certs: /ca.pem\n ssl_certfile: /client.pem\n ssl_keyfile: /client.key\n ssl_cert_reqs: none\n ssl_check_hostname: true",
&[],
Some(TlsOptions { ca_certs: Some("/ca.pem".into()), certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), verify_certificate: false, check_hostname: true })
Some(TlsOptions { ca_certs: Some("/ca.pem".into()), ca_data: None, ca_path: None, certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), verify_certificate: false, check_hostname: true })
)]
#[case::options_from_the_environment(
" ssl: true",
&[("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")],
Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), certfile: None, keyfile: None, verify_certificate: true, check_hostname: false })
Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, verify_certificate: true, check_hostname: false })
)]
#[case::gcp_ca_certs_replace_the_configured_ones(
" ssl: true\n ssl_ca_certs: /ca.pem\n gcp_service_account: account\n gcp_ssl_ca_certs: /gcp-ca.pem",
&[],
Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), certfile: None, keyfile: None, verify_certificate: true, check_hostname: false })
Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, verify_certificate: true, check_hostname: false })
)]
#[case::options_without_ssl_stay_plaintext(" ssl_ca_certs: /ca.pem", &[], None)]
fn tls_options_resolve_like_redis_py(
@ -521,19 +539,23 @@ fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
}
#[rstest]
fn rejected_environment_settings_never_print_the_value() {
let error = resolve(
"model_list: []",
&[
("REDIS_HOST", "redis"),
("REDIS_SSL_CIPHERS", "private-option-value"),
],
)
.unwrap_err();
let message = error.to_string();
fn unsupported_environment_settings_never_print_the_value() {
let events = Arc::new(Mutex::new(Vec::new()));
let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone()));
tracing::subscriber::with_default(subscriber, || {
resolve(
"model_list: []",
&[
("REDIS_HOST", "redis"),
("REDIS_SSL_CIPHERS", "private-option-value"),
],
)
.unwrap();
});
let warnings = events.lock().unwrap().join(" ");
assert!(message.contains("REDIS_SSL_CIPHERS"));
assert!(!message.contains("private-option-value"));
assert!(warnings.contains("REDIS_SSL_CIPHERS"), "{warnings}");
assert!(!warnings.contains("private-option-value"), "{warnings}");
}
#[rstest]
@ -659,20 +681,12 @@ struct ContractCase {
const CONTRACT: &str = include_str!("redis_resolution.json");
// Settings redis-py accepts and Rust rejects at boot on purpose, each listed in crates/config/AGENTS.md
const DIVERGENCES: &[(&str, &str)] = &[
("ssl_ciphers", "ssl_ciphers"),
("unix_socket_path", "unix_socket_path"),
(
"a_url_query_setting_without_a_redis_rs_equivalent",
"url ?socket_timeout",
),
];
fn tls_view(tls: Option<&TlsOptions>) -> serde_json::Value {
tls.map_or(serde_json::Value::Null, |tls| {
serde_json::json!({
"ca_certs": tls.ca_certs,
"ca_data": tls.ca_data,
"ca_path": tls.ca_path,
"certfile": tls.certfile,
"keyfile": tls.keyfile,
"verify_certificate": tls.verify_certificate,
@ -694,6 +708,9 @@ fn expose(value: Option<&litellm_secret_value::SecretValue>) -> Option<&str> {
fn view(resolved: &ResolvedRedis) -> serde_json::Value {
let target = match &resolved.endpoint {
Endpoint::Host(host) if host.unix_socket_path.is_some() => serde_json::json!({
"mode": "unix", "path": host.unix_socket_path, "tls": null,
}),
Endpoint::Host(host) => serde_json::json!({
"mode": "host",
"host": host.host,
@ -748,6 +765,10 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
"auth": auth,
"socket_timeout": resolved.timeout.map_or(5.0, |timeout| timeout.as_secs_f64()),
"socket_connect_timeout": resolved.connect_timeout.or(resolved.timeout).map_or(5.0, |timeout| timeout.as_secs_f64()),
"health_check_interval": resolved.health_check_interval
.or_else(|| matches!(resolved.endpoint, Endpoint::Cluster(_)).then_some(Duration::from_secs(25)))
.filter(|interval| !interval.is_zero())
.map(|interval| interval.as_secs_f64()),
});
let serde_json::Value::Object(rest) = rest else {
unreachable!("the view is an object")
@ -763,11 +784,6 @@ fn rust_outcome(case: &ContractCase) -> Result<serde_json::Value, RedisError> {
fn matches_python_rejection(error: &RedisError, python: &serde_json::Value) -> bool {
let rejection = &python["rejected"];
match error {
RedisError::Unsupported { setting } if setting == "ssl_ca_path" => {
rejection["type"] == "TypeError"
&& rejection["message"]
== "Redis.__init__() got an unexpected keyword argument 'ssl_ca_path'"
}
RedisError::NoTarget => {
rejection["type"] == "ValueError"
&& rejection["message"] == "Either 'host' or 'url' must be specified for redis."
@ -793,23 +809,34 @@ fn matches_python_rejection(error: &RedisError, python: &serde_json::Value) -> b
}
fn mismatch(case: &ContractCase, contract: &Contract) -> Option<String> {
let divergence = DIVERGENCES.iter().find(|(name, _)| *name == case.name);
let python_rejects = case.python.get("rejected").is_some();
match (rust_outcome(case), divergence) {
(
Err(RedisError::Unsupported { setting } | RedisError::Invalid { setting, .. }),
Some((_, expected)),
) if setting == *expected && !python_rejects => None,
(outcome, Some(_)) => Some(format!(
"{}: listed as a divergence, got {outcome:?}",
case.name
)),
(Err(ref error), None) if matches_python_rejection(error, &case.python) => None,
(Err(RedisError::Unsupported { setting }), None) => Some(format!(
"{}: Rust rejects {setting}, which Python accepts",
case.name
)),
(Ok(view), None)
match rust_outcome(case) {
Err(ref error) if matches_python_rejection(error, &case.python) => None,
Ok(view)
if matches!(
case.name.as_str(),
"ssl_ca_path_is_unsupported" | "ssl_ca_path_from_the_environment"
) && case.python["rejected"]["type"] == "TypeError"
&& case.python["rejected"]["message"]
== "Redis.__init__() got an unexpected keyword argument 'ssl_ca_path'"
&& view["tls"]["ca_path"]
== case.kwargs.get("ssl_ca_path").cloned().unwrap_or_else(|| {
serde_json::json!(case.environment["REDIS_SSL_CA_PATH"])
}) =>
{
None
}
Ok(view)
if case.name == "a_url_query_setting_without_a_redis_rs_equivalent"
&& view
== serde_json::json!({
"mode": "url", "url": case.kwargs["url"], "tls": null,
"auth": case.python["auth"], "socket_timeout": 5.0,
"socket_connect_timeout": 5.0, "health_check_interval": null,
}) =>
{
None
}
Ok(view)
if case.name == "retry_on_error_from_the_environment"
&& case.python["rejected"]["type"] == "TypeError"
&& case.python["rejected"]["message"] == "'bool' object is not iterable"
@ -820,8 +847,8 @@ fn mismatch(case: &ContractCase, contract: &Contract) -> Option<String> {
{
None
}
(Ok(view), None) if view == case.python => None,
(outcome, None) => Some(format!(
Ok(view) if view == case.python => None,
outcome => Some(format!(
"{}: rust {outcome:?}, python {}",
case.name, case.python
)),
@ -834,19 +861,12 @@ fn mismatch(case: &ContractCase, contract: &Contract) -> Option<String> {
fn resolution_matches_python_redis() {
let contract: Contract = serde_json::from_str(CONTRACT).unwrap();
let unknown_divergences: Vec<&str> = DIVERGENCES
.iter()
.map(|(name, _)| *name)
.filter(|name| !contract.cases.iter().any(|case| case.name == *name))
.collect();
let mismatches: Vec<String> = contract
.cases
.iter()
.filter_map(|case| mismatch(case, &contract))
.collect();
assert_eq!(unknown_divergences, Vec::<&str>::new());
assert!(mismatches.is_empty(), "{}", mismatches.join("\n"));
}

227
litellm-rust/crates/config/tests/redis_resolution.json
View file

@ -13,6 +13,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -35,6 +36,7 @@
"password": null,
"username": null
},
"health_check_interval": 25.0,
"mode": "cluster",
"nodes": [
"node-a:7000"
@ -60,6 +62,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"mode": "sentinel",
"nodes": [
"sentinel-a:26379"
@ -92,6 +95,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"mode": "sentinel",
"nodes": [
"sentinel-a:26379"
@ -117,6 +121,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -159,6 +164,7 @@
"username": "user"
},
"db": 2,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6380,
@ -166,6 +172,8 @@
"socket_timeout": 1.5,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -189,6 +197,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "environment",
"mode": "host",
"port": 6390,
@ -196,6 +205,8 @@
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -219,6 +230,7 @@
"username": null
},
"db": 3,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6380,
@ -242,6 +254,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -262,6 +275,7 @@
"password": null,
"username": null
},
"health_check_interval": null,
"mode": "url",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
@ -282,6 +296,7 @@
"password": null,
"username": null
},
"health_check_interval": null,
"mode": "url",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
@ -304,6 +319,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "configured",
"mode": "host",
"port": 6379,
@ -328,11 +344,14 @@
"password": null,
"username": "user"
},
"health_check_interval": null,
"mode": "url",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -356,11 +375,14 @@
"password": null,
"username": null
},
"health_check_interval": null,
"mode": "url",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -388,6 +410,7 @@
"password": "cluster-secret",
"username": null
},
"health_check_interval": 25.0,
"mode": "cluster",
"nodes": [
"node-a:7000"
@ -409,6 +432,7 @@
"password": null,
"username": null
},
"health_check_interval": 25.0,
"mode": "cluster",
"nodes": [
"node-a:7000",
@ -439,6 +463,7 @@
"password": null,
"username": null
},
"health_check_interval": 25.0,
"mode": "cluster",
"nodes": [
"node-a:7000"
@ -447,6 +472,8 @@
"socket_timeout": 5.0,
"tls": {
"ca_certs": "/ca.pem",
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": true,
"keyfile": null,
@ -479,6 +506,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"mode": "sentinel",
"nodes": [
"sentinel-a:26379",
@ -506,6 +534,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"mode": "sentinel",
"nodes": [
"sentinel-a:26379"
@ -556,6 +585,7 @@
"password": null,
"username": null
},
"health_check_interval": null,
"mode": "url",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
@ -580,6 +610,7 @@
"kind": "host"
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -605,6 +636,7 @@
"service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -629,6 +661,7 @@
"service_account": "account"
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -636,6 +669,8 @@
"socket_timeout": 5.0,
"tls": {
"ca_certs": "/gcp-ca.pem",
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -663,6 +698,7 @@
"username": "environment-user"
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -687,6 +723,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -712,6 +749,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -742,6 +780,7 @@
"user_name": "iam-user"
},
"db": 0,
"health_check_interval": null,
"host": "cache-1",
"mode": "host",
"port": 6379,
@ -749,6 +788,8 @@
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -778,6 +819,7 @@
"user_name": "iam-user"
},
"db": 0,
"health_check_interval": null,
"host": "cache-1",
"mode": "host",
"port": 6379,
@ -785,6 +827,8 @@
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -811,11 +855,14 @@
"serverless": true,
"user_name": "iam-user"
},
"health_check_interval": null,
"mode": "url",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -875,6 +922,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -882,6 +930,8 @@
"socket_timeout": 5.0,
"tls": {
"ca_certs": "/ca.pem",
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": true,
"keyfile": null,
@ -906,6 +956,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -913,6 +964,8 @@
"socket_timeout": 5.0,
"tls": {
"ca_certs": "/env-ca.pem",
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -935,6 +988,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -960,6 +1014,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "referenced",
"mode": "host",
"port": 6379,
@ -997,6 +1052,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -1004,6 +1060,8 @@
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": "/client.pem",
"check_hostname": false,
"keyfile": "/client.key",
@ -1023,11 +1081,14 @@
"password": null,
"username": null
},
"health_check_interval": null,
"mode": "url",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": {
"ca_certs": "/ca.pem",
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": true,
"keyfile": null,
@ -1048,6 +1109,7 @@
"password": null,
"username": null
},
"health_check_interval": null,
"mode": "url",
"socket_connect_timeout": 2.0,
"socket_timeout": 2.0,
@ -1069,11 +1131,14 @@
"password": null,
"username": null
},
"health_check_interval": null,
"mode": "url",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -1138,6 +1203,7 @@
"password": null,
"username": null
},
"health_check_interval": 25.0,
"mode": "cluster",
"nodes": [
"node-a:7000"
@ -1184,6 +1250,7 @@
"password": null,
"username": null
},
"health_check_interval": 25.0,
"mode": "cluster",
"nodes": [
"node-a:7000"
@ -1232,6 +1299,7 @@
"password": null,
"username": null
},
"health_check_interval": 25.0,
"mode": "cluster",
"nodes": [
"node-a:7000"
@ -1256,6 +1324,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -1263,6 +1332,8 @@
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": null,
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
@ -1284,6 +1355,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -1306,6 +1378,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -1328,6 +1401,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -1350,6 +1424,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -1372,6 +1447,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -1394,6 +1470,7 @@
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
@ -1430,12 +1507,162 @@
"password": null,
"username": null
},
"health_check_interval": null,
"mode": "unix",
"path": "/tmp/redis.sock",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": null
}
},
{
"environment": {},
"kwargs": {
"host": "redis",
"ssl": true,
"ssl_ca_certs": "/ca.pem",
"unix_socket_path": "/tmp/redis.sock"
},
"name": "unix_socket_path_keeps_tls_settings_out_of_the_target",
"python": {
"auth": {
"kind": "static",
"password": null,
"username": null
},
"health_check_interval": null,
"mode": "unix",
"path": "/tmp/redis.sock",
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": null
}
},
{
"environment": {},
"kwargs": {
"host": "redis",
"ssl": true,
"ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----"
},
"name": "ssl_ca_data",
"python": {
"auth": {
"kind": "static",
"password": null,
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": {
"ca_certs": null,
"ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----",
"ca_path": null,
"certfile": null,
"check_hostname": false,
"keyfile": null,
"verify_certificate": true
}
}
},
{
"environment": {
"REDIS_SSL": "true",
"REDIS_SSL_CA_PATH": "/etc/ssl/private-cas"
},
"kwargs": {
"host": "redis"
},
"name": "ssl_ca_path_from_the_environment",
"python": {
"rejected": {
"message": "Redis.__init__() got an unexpected keyword argument 'ssl_ca_path'",
"type": "TypeError"
}
}
},
{
"environment": {},
"kwargs": {
"health_check_interval": "7",
"host": "redis"
},
"name": "health_check_interval_is_whole_seconds",
"python": {
"auth": {
"kind": "static",
"password": null,
"username": null
},
"db": 0,
"health_check_interval": 7.0,
"host": "redis",
"mode": "host",
"port": 6379,
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": null
}
},
{
"environment": {},
"kwargs": {
"health_check_interval": 0,
"host": "redis"
},
"name": "health_check_interval_zero_is_off",
"python": {
"auth": {
"kind": "static",
"password": null,
"username": null
},
"db": 0,
"health_check_interval": null,
"host": "redis",
"mode": "host",
"port": 6379,
"socket_connect_timeout": 5.0,
"socket_timeout": 5.0,
"tls": null
}
},
{
"environment": {},
"kwargs": {
"db": 3,
"startup_nodes": [
{
"host": "node-a",
"port": 7000
}
]
},
"name": "a_cluster_rejects_a_nonzero_db",
"python": {
"rejected": {
"message": "Argument 'db' must be 0 or None in cluster mode",
"type": "RedisClusterException"
}
}
},
{
"environment": {},
"kwargs": {
"unix_socket_path": "/tmp/redis.sock"
},
"name": "unix_socket_path_without_a_host_fails_like_python",
"python": {
"rejected": {
"message": "Either 'host' or 'url' must be specified for redis.",
"type": "ValueError"
}
}
}
]
}

69
litellm-rust/crates/gateway/src/redis.rs
View file

@ -1,9 +1,8 @@
use std::{sync::Arc, time::Duration};
use litellm_circuit_breaker::Config as BreakerConfig;
use litellm_config::{
CacheSettings, Config, Expected, RedisError, RedisSettings, ResolvedRedis, SecretSnapshot,
Value, environment_has_target,
CacheSettings, Config, RedisSettings, ResolvedRedis, SecretSnapshot, Value,
environment_has_target,
};
use litellm_core_utils::settings::Lookup;
use litellm_redis_client::{RedisClient, RedisEvents};
@ -72,9 +71,11 @@ pub(crate) async fn connect(
resources: &Resources,
) -> Result<RedisClients, Error> {
let plan = plan(config, lookup)?;
let breaker = breaker_config(lookup)?;
let breaker = litellm_redis_connect::breaker_config(lookup)?;
let events: Arc<dyn RedisEvents> = Arc::new(TracingEvents::new(lookup));
let client = |params: Params, breaker: Option<BreakerConfig>| -> Result<RedisClient, Error> {
let client = |params: Params,
breaker: Option<litellm_circuit_breaker::Config>|
-> Result<RedisClient, Error> {
Ok(litellm_redis_connect::connect(
params,
HostServices {
@ -178,62 +179,6 @@ async fn reachable(client: RedisClient) -> Option<RedisClient> {
}
}
fn parse_seconds(text: &str) -> Option<Duration> {
text.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
}
fn breaker_config(lookup: &dyn Lookup) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup
.get("REDIS_CIRCUIT_BREAKER_ENABLED")
.is_none_or(|value| value.eq_ignore_ascii_case("true"));
if !enabled {
return Ok(None);
}
let defaults = BreakerConfig::default();
Ok(Some(BreakerConfig {
failure_threshold: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
Expected::Integer,
|text| text.parse().ok(),
)?
.unwrap_or(defaults.failure_threshold),
recovery_timeout: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.recovery_timeout),
timeout_min_duration: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.timeout_min_duration),
}))
}
fn environment_value<T>(
lookup: &dyn Lookup,
name: &str,
expected: Expected,
parse: impl Fn(&str) -> Option<T>,
) -> Result<Option<T>, RedisError> {
lookup
.get(name)
.map(|value| {
parse(&value).ok_or_else(|| RedisError::Invalid {
setting: name.to_owned(),
expected,
})
})
.transpose()
}
#[cfg(test)]
mod tests {
use litellm_core_utils::settings::Lookup;
@ -264,6 +209,7 @@ mod tests {
port,
tls: None,
db: 0,
unix_socket_path: None,
}),
auth: RedisAuth::Static {
username: None,
@ -272,6 +218,7 @@ mod tests {
namespace: None,
timeout: None,
connect_timeout: None,
health_check_interval: None,
}
}

6
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -31,6 +31,12 @@ litellm-cache.workspace = true
litellm-error.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-circuit-breaker.workspace = true
litellm-redis-client.workspace = true
litellm-redis-connect.workspace = true
litellm-redis-store.workspace = true
litellm-config.workspace = true
tracing.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true

10
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -12,6 +12,7 @@ mod marshal;
mod model_registry;
mod preflight;
mod python_settings;
mod redis;
mod routes;
mod secrets;
mod tokenizer;
@ -36,6 +37,12 @@ mod _native {
#[pymodule_export]
use crate::model_registry::bundled_model_cost_map;
#[pymodule_export]
use crate::redis::native_redis_breaker_states;
#[pymodule_export]
use crate::redis::router_cooldowns::NativeRouterCooldowns;
#[pymodule_export]
use crate::redis::router_usage::NativeRouterUsage;
#[pymodule_export]
use crate::routes::audio_transcription::{atranscription, transcription};
#[pymodule_export]
use crate::routes::chat_completions::{acompletion, completion};
@ -113,6 +120,9 @@ mod tests {
"aresponses",
"ResponsesWebSocketConnection",
"NativeDiagnosticProcessor",
"NativeRouterCooldowns",
"NativeRouterUsage",
"native_redis_breaker_states",
"NativeClickHouseSpendConfig",
"NativeClickHouseSpendStorage",
"NativeDatabase",

2
litellm-rust/crates/python-bridge/src/python_settings.rs
View file

@ -13,6 +13,8 @@ pub(crate) enum PythonSettings {
UrlPolicy,
#[strum(serialize = "provider_defaults")]
ProviderDefaults,
#[strum(serialize = "router_tiers")]
RouterTiers,
#[strum(serialize = "secret_manager")]
SecretManager,
#[strum(serialize = "secret_manager_binding")]

145
litellm-rust/crates/python-bridge/src/redis/client.rs Normal file
View file

@ -0,0 +1,145 @@
use std::{
collections::BTreeSet,
sync::{Arc, LazyLock, Mutex, PoisonError},
};
use litellm_config::{RedisSettings, ResolvedRedis, SecretSnapshot};
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_redis_client::{RedisClient, RedisEvents};
use litellm_redis_connect::{HostServices, SharedEnvironment};
use pyo3::{
prelude::*,
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple},
};
use super::events::{ServiceLoggingEvents, state_label};
/// The `redis_kwargs` dict `RedisCache` hands redis-py, read the way `litellm/_redis.py` reads it.
/// A value that is not plain data, such as a `credential_provider` object, reaches the resolver by
/// its type name, so the resolver or the connector names the setting it cannot honor.
pub(crate) struct PythonRedisSettings(RedisSettings);
impl<'a, 'py> FromPyObject<'a, 'py> for PythonRedisSettings {
type Error = PyErr;
fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let py = kwargs.py();
let plain = PyDict::new(py);
for (name, value) in kwargs.cast::<PyDict>()?.iter() {
if value.is_none() {
continue;
}
if is_plain(&value) {
plain.set_item(&name, value)?;
} else {
plain.set_item(&name, value.get_type().name()?)?;
}
}
Ok(Self(litellm_host_python::from_py_argument(plain.as_any())?))
}
}
fn is_plain(value: &Bound<'_, PyAny>) -> bool {
value.is_instance_of::<PyString>()
|| value.is_instance_of::<PyBool>()
|| value.is_instance_of::<PyInt>()
|| value.is_instance_of::<PyFloat>()
|| value.is_instance_of::<PyList>()
|| value.is_instance_of::<PyTuple>()
|| value.is_instance_of::<PyDict>()
}
struct Shared {
resolved: ResolvedRedis,
client: RedisClient,
}
static CLIENTS: LazyLock<Mutex<Vec<Shared>>> = LazyLock::new(|| Mutex::new(Vec::new()));
/// The one client for a resolved Redis config, so every family in this process shares its pool
/// and its breaker. Like redis-py, which never connects at construction, a config the resolver or
/// the connector cannot honor leaves the family on memory with a warning rather than failing the
/// `Router`.
pub(crate) fn shared_client(
py: Python<'_>,
kwargs: PythonRedisSettings,
namespace: Option<String>,
) -> PyResult<Option<RedisClient>> {
let environment: SharedEnvironment = Arc::new(ProcessEnvironment);
let secrets = match secret_snapshot(py, kwargs.0.secret_names())? {
Ok(secrets) => secrets,
Err(error) => return Ok(memory_only(py, &error)),
};
let logger = crate::logger::capture(py);
let resolved = match logger.scope(|| kwargs.0.resolve(&secrets, &*environment)) {
Ok(resolved) => ResolvedRedis {
namespace,
..resolved
},
Err(error) => return Ok(memory_only(py, &error)),
};
let mut clients = CLIENTS.lock().unwrap_or_else(PoisonError::into_inner);
if let Some(shared) = clients.iter().find(|shared| shared.resolved == resolved) {
return Ok(Some(shared.client.clone()));
}
let events: Arc<dyn RedisEvents> = Arc::new(ServiceLoggingEvents::new(py)?);
let breaker = match logger.scope(|| litellm_redis_connect::breaker_config(&*environment)) {
Ok(breaker) => breaker,
Err(error) => return Ok(memory_only(py, &error)),
};
let host = HostServices {
events: Some(events),
breaker,
http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?,
environment,
};
let target = resolved.clone();
let client = match crate::execution::run_sync_value(py, async move {
Ok(litellm_redis_connect::connect(target, host))
})? {
Ok(client) => client,
Err(error) => return Ok(memory_only(py, &error)),
};
clients.push(Shared {
resolved,
client: client.clone(),
});
Ok(Some(client))
}
/// The breaker state of every native Redis client this process built, so the Prometheus state
/// gauge counts them next to the Python `RedisCache` breakers.
#[pyfunction]
pub(crate) fn native_redis_breaker_states() -> Vec<&'static str> {
let clients = CLIENTS.lock().unwrap_or_else(PoisonError::into_inner);
clients
.iter()
.filter_map(|shared| shared.client.breaker_state())
.map(|state| state_label(&state))
.collect()
}
/// Every name the resolver may read, fetched once the way Python's `get_secret` fetches it: the
/// configured secret manager first, then the process environment.
fn secret_snapshot(
py: Python<'_>,
names: BTreeSet<String>,
) -> PyResult<Result<SecretSnapshot, litellm_secrets::Error>> {
let source = crate::secrets::source(py)?;
let owned: Vec<String> = names.iter().cloned().collect();
let lookup = crate::execution::run_sync_value(py, async move {
let borrowed: Vec<&str> = owned.iter().map(String::as_str).collect();
Ok(source.resolve(&borrowed).await)
})?;
Ok(lookup.map(|lookup| SecretSnapshot::read(names, &*lookup)))
}
fn memory_only(py: Python<'_>, error: &(dyn std::error::Error + 'static)) -> Option<RedisClient> {
let report = litellm_error::Report::new(error).to_string();
crate::logger::capture(py).scope(|| {
litellm_tracing::warn!(
"Redis is not usable from the Rust router stores, cooldowns and usage stay in memory: {report}"
)
});
None
}

251
litellm-rust/crates/python-bridge/src/redis/events.rs Normal file
View file

@ -0,0 +1,251 @@
use std::{sync::Arc, time::Duration};
use litellm_circuit_breaker::{State, Transition};
use litellm_host_python::{PythonContext, attach_blocking};
use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents};
use pyo3::{prelude::*, types::PyModule};
tokio::task_local! {
static CALL: RedisCall;
}
/// The Python caller of one store operation. Its context carries `with_service_target`, and its
/// `call_type` and `caller` are the labels `RedisCache` reports, so dashboards keep their series.
#[derive(Clone)]
pub(crate) struct RedisCall {
pub(crate) context: PythonContext,
pub(crate) call_type: &'static str,
pub(crate) caller: Option<String>,
}
impl RedisCall {
pub(crate) async fn scope<F: Future>(self, operation: F) -> F::Output {
CALL.scope(self, operation).await
}
}
/// Reports every native Redis call and breaker transition to `ServiceLogging`.
pub(crate) struct ServiceLoggingEvents {
record_call: Arc<Py<PyAny>>,
record_breaker: Arc<Py<PyAny>>,
context: PythonContext,
}
impl ServiceLoggingEvents {
pub(crate) fn new(py: Python<'_>) -> PyResult<Self> {
const MODULE: &str = "litellm.rust_bridge.host.redis_events";
Self::from_module(py, &py.import(MODULE)?)
}
fn from_module(py: Python<'_>, module: &Bound<'_, PyModule>) -> PyResult<Self> {
Ok(Self {
record_call: Arc::new(module.getattr("record_redis_call")?.unbind()),
record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()),
context: PythonContext::capture(py)?,
})
}
fn report(&self, context: PythonContext, record: Arc<Py<PyAny>>, args: ReportArgs) {
let Ok(runtime) = tokio::runtime::Handle::try_current() else {
return;
};
runtime.spawn(async move {
let reported = attach_blocking(context, move |py| match args {
ReportArgs::Call(call_type, duration, failure, caller) => record
.bind(py)
.call1((call_type, duration.as_secs_f64(), failure, caller))
.map(drop),
ReportArgs::Breaker(state) => record.bind(py).call1((state,)).map(drop),
})
.await
.and_then(|reported| reported);
if let Err(error) = reported {
tracing::warn!(error = %error, "reporting a Redis event to ServiceLogging failed");
}
});
}
}
enum ReportArgs {
Call(&'static str, Duration, Option<&'static str>, Option<String>),
Breaker(&'static str),
}
fn operation_call_type(operation: Operation) -> &'static str {
match operation {
Operation::Command(name) => name,
Operation::Script => "eval",
Operation::Pipeline { .. } => "pipeline",
}
}
fn failure_label(outcome: Outcome) -> Option<&'static str> {
match outcome {
Outcome::Succeeded => None,
Outcome::Failed(FailureKind::Timeout) => Some("timeout"),
Outcome::Failed(FailureKind::Transport) => Some("transport"),
Outcome::Failed(FailureKind::Server) => Some("server"),
Outcome::Failed(FailureKind::Decode) => Some("decode"),
Outcome::Failed(FailureKind::BreakerOpen) => Some("breaker_open"),
}
}
pub(crate) fn state_label(state: &State) -> &'static str {
match state {
State::Closed { .. } => "closed",
State::Open { .. } => "open",
State::HalfOpen => "half_open",
}
}
impl RedisEvents for ServiceLoggingEvents {
fn call(&self, event: &CallEvent) {
let caller = CALL.try_with(Clone::clone).ok();
let (context, call_type, caller) = match caller {
Some(call) => (call.context, call.call_type, call.caller),
None => (
self.context.clone(),
operation_call_type(event.operation),
None,
),
};
self.report(
context,
self.record_call.clone(),
ReportArgs::Call(
call_type,
event.duration,
failure_label(event.outcome),
caller,
),
);
}
fn breaker(&self, transition: &Transition<State>) {
self.report(
self.context.clone(),
self.record_breaker.clone(),
ReportArgs::Breaker(state_label(&transition.to)),
);
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use litellm_host_python::PythonContext;
use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents};
use pyo3::{prelude::*, types::PyModule};
use rstest::rstest;
use super::{RedisCall, ServiceLoggingEvents};
const RECORDER: &str = "
import contextvars
target = contextvars.ContextVar('target', default=None)
seen = []
def record_redis_call(call_type, duration, failure, caller):
seen.append((call_type, round(duration, 3), failure, caller, target.get()))
def record_redis_breaker(state):
seen.append(('breaker', state))
";
fn recorder(name: &str) -> (ServiceLoggingEvents, Py<PyModule>) {
Python::initialize();
Python::attach(|py| {
let name = std::ffi::CString::new(format!("redis_events_recorder_{name}")).unwrap();
let module =
PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name)
.unwrap();
(
ServiceLoggingEvents::from_module(py, &module).unwrap(),
module.unbind(),
)
})
}
async fn seen(module: &Py<PyModule>, count: usize) -> String {
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
loop {
let seen: (usize, String) = Python::attach(|py| {
let seen = module.bind(py).getattr("seen").unwrap();
(seen.len().unwrap(), seen.repr().unwrap().to_string())
});
if seen.0 >= count || tokio::time::Instant::now() > deadline {
return seen.1;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn a_scoped_call_reports_the_store_labels_in_the_callers_context() {
let (events, namespace) = recorder("scoped");
let call = Python::attach(|py| {
namespace
.bind(py)
.getattr("target")
.unwrap()
.call_method1("set", ("spend_counter",))
.unwrap();
RedisCall {
context: PythonContext::capture(py).unwrap(),
call_type: "async_increment",
caller: Some("router.py:42".to_owned()),
}
});
let event = CallEvent {
operation: Operation::Command("INCRBYFLOAT"),
duration: Duration::from_millis(250),
outcome: Outcome::Failed(FailureKind::Timeout),
};
call.scope(async { events.call(&event) }).await;
assert_eq!(
seen(&namespace, 1).await,
"[('async_increment', 0.25, 'timeout', 'router.py:42', 'spend_counter')]"
);
}
#[rstest]
#[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")]
#[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")]
#[tokio::test(flavor = "multi_thread")]
async fn breaker_transitions_report_python_state_names(
#[case] to: litellm_circuit_breaker::State,
#[case] expected: &str,
) {
let (events, namespace) = recorder(expected);
events.breaker(&litellm_circuit_breaker::Transition {
from: litellm_circuit_breaker::State::HalfOpen,
to,
generation: 1,
});
assert_eq!(
seen(&namespace, 1).await,
format!("[('breaker', '{expected}')]")
);
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn an_unscoped_call_falls_back_to_the_command_name() {
let (events, namespace) = recorder("unscoped");
events.call(&CallEvent {
operation: Operation::Pipeline { ops: 3 },
duration: Duration::from_millis(5),
outcome: Outcome::Succeeded,
});
assert_eq!(
seen(&namespace, 1).await,
"[('pipeline', 0.005, None, None, None)]"
);
}
}

60
litellm-rust/crates/python-bridge/src/redis/mod.rs Normal file
View file

@ -0,0 +1,60 @@
mod client;
mod events;
pub(crate) use client::native_redis_breaker_states;
mod policy;
pub(crate) mod router_cooldowns;
pub(crate) mod router_usage;
use litellm_host_python::PythonContext;
use litellm_redis_store::Number;
use pyo3::{
exceptions::{PyRuntimeError, PyTypeError},
prelude::*,
types::{PyBool, PyFloat, PyInt},
};
use self::events::RedisCall;
/// The labels `RedisCache` reports for this operation, in the caller's context so its
/// `service_target` reaches the service event.
fn call(py: Python<'_>, call_type: &'static str) -> PyResult<RedisCall> {
Ok(RedisCall {
context: PythonContext::capture(py)?,
call_type,
caller: None,
})
}
fn memory_error(error: litellm_cache::Error) -> PyErr {
PyRuntimeError::new_err(litellm_error::Report::new(&error).to_string())
}
/// A Python `int` or `float` across the bridge, keeping which one it is.
pub(crate) struct PyNumber(Number);
impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
if value.is_instance_of::<PyInt>() && !value.is_instance_of::<PyBool>() {
return Ok(Self(Number::Int(value.extract()?)));
}
if value.is_instance_of::<PyFloat>() {
return Ok(Self(Number::Float(value.extract()?)));
}
Err(PyTypeError::new_err("expected an int or a float"))
}
}
impl<'py> IntoPyObject<'py> for PyNumber {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
Ok(match self.0 {
Number::Int(value) => value.into_pyobject(py)?.into_any(),
Number::Float(value) => value.into_pyobject(py)?.into_any(),
})
}
}

130
litellm-rust/crates/python-bridge/src/redis/policy.rs Normal file
View file

@ -0,0 +1,130 @@
use std::time::Duration;
use litellm_cache::TierPolicy;
use litellm_redis_store::{router_cooldowns, router_usage};
use pyo3::prelude::*;
use crate::{
coercion::{Field, FieldSpec, ProjectionError},
python_settings::{PythonSettings, Snapshot},
};
const COOLDOWN_REFRESH: FieldSpec<Duration> =
FieldSpec::new("cooldown_redis_read_interval", seconds);
const USAGE_REFRESH: FieldSpec<Duration> = FieldSpec::new("usage_redis_read_interval", seconds);
const TRACKED_READS: FieldSpec<usize> = FieldSpec::new("max_redis_batch_cache_size", count);
/// The tier values this Python process runs with. `POLICY` in each family holds `main`'s
/// defaults, the host reports what its environment and `litellm` globals changed.
#[derive(Debug)]
pub(crate) struct RouterTierPolicies {
pub(crate) cooldowns: TierPolicy,
pub(crate) usage: TierPolicy,
}
pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult<RouterTierPolicies> {
Ok(project(&PythonSettings::RouterTiers.read(py)?)?)
}
fn project(snapshot: &Snapshot<'_>) -> Result<RouterTierPolicies, ProjectionError> {
let tracked_reads = snapshot.read(&TRACKED_READS)?;
Ok(RouterTierPolicies {
cooldowns: TierPolicy {
refresh_after: snapshot.read(&COOLDOWN_REFRESH)?,
tracked_reads,
..router_cooldowns::POLICY
},
usage: TierPolicy {
refresh_after: snapshot.read(&USAGE_REFRESH)?,
tracked_reads,
..router_usage::POLICY
},
})
}
fn seconds(field: &Field<'_>) -> Result<Duration, ProjectionError> {
match field.value().extract::<f64>() {
Ok(seconds) if seconds.is_finite() && seconds >= 0.0 => {
Ok(Duration::from_secs_f64(seconds))
}
_ => Err(field.invalid("a non-negative number of seconds")),
}
}
fn count(field: &Field<'_>) -> Result<usize, ProjectionError> {
field
.value()
.extract::<usize>()
.map_err(|_| field.invalid("a non-negative integer"))
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use pyo3::{prelude::*, types::PyDict};
use super::project;
use crate::python_settings::PythonSettings;
fn projected(py: Python<'_>, fields: &str) -> Result<super::RouterTierPolicies, String> {
let locals = PyDict::new(py);
let source = std::ffi::CString::new(format!(
"import types\nsettings = types.SimpleNamespace({fields})"
))
.unwrap();
py.run(&source, Some(&locals), Some(&locals)).unwrap();
let snapshot =
PythonSettings::RouterTiers.snapshot(locals.get_item("settings").unwrap().unwrap());
project(&snapshot).map_err(|error| PyErr::from(error).to_string())
}
#[test]
fn the_host_values_replace_only_the_refresh_and_tracked_reads_of_mains_defaults() {
Python::initialize();
Python::attach(|py| {
let policies = projected(
py,
"cooldown_redis_read_interval=2.5, usage_redis_read_interval=30, max_redis_batch_cache_size=42",
)
.unwrap();
assert_eq!(
policies.cooldowns.refresh_after,
Duration::from_millis(2500)
);
assert_eq!(policies.usage.refresh_after, Duration::from_secs(30));
assert_eq!(policies.cooldowns.tracked_reads, 42);
assert_eq!(policies.usage.tracked_reads, 42);
assert_eq!(
policies.cooldowns.backfill_ttl,
litellm_redis_store::router_cooldowns::POLICY.backfill_ttl
);
assert_eq!(
policies.usage.backfill_ttl,
litellm_redis_store::router_usage::POLICY.backfill_ttl
);
});
}
#[test]
fn a_value_python_would_not_run_with_is_a_configuration_error() {
Python::initialize();
Python::attach(|py| {
let negative = projected(
py,
"cooldown_redis_read_interval=-1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000",
)
.unwrap_err();
let text = projected(
py,
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many'",
)
.unwrap_err();
assert!(
negative.contains("cooldown_redis_read_interval"),
"{negative}"
);
assert!(text.contains("max_redis_batch_cache_size"), "{text}");
});
}
}

191
litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs Normal file
View file

@ -0,0 +1,191 @@
use std::{sync::Arc, time::Duration};
use litellm_redis_store::router_cooldowns::{Cooldown, RouterCooldowns};
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use super::{
PyNumber, call,
client::{PythonRedisSettings, shared_client},
memory_error,
policy::router_tier_policies,
};
use crate::execution::{run_async_value, run_sync_value};
/// One router's cooldown store. The router's sync completion path calls `put`, `active`,
/// `minimum` and `increment_allowed_fails`, so those block with the GIL released like redis-py.
#[pyclass(frozen, name = "NativeRouterCooldowns")]
pub(crate) struct NativeRouterCooldowns(Arc<RouterCooldowns>);
pub(crate) struct PyCooldown(Cooldown);
impl<'a, 'py> FromPyObject<'a, 'py> for PyCooldown {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let fields = value.cast::<PyDict>()?;
let item = |name: &str| fields.as_any().get_item(name);
Ok(Self(Cooldown {
exception_received: item("exception_received")?.extract()?,
status_code: item("status_code")?.extract()?,
timestamp: item("timestamp")?.extract::<PyNumber>()?.0,
cooldown_time: item("cooldown_time")?.extract::<PyNumber>()?.0,
}))
}
}
impl<'py> IntoPyObject<'py> for PyCooldown {
type Target = PyDict;
type Output = Bound<'py, PyDict>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
let fields = PyDict::new(py);
fields.set_item("exception_received", self.0.exception_received)?;
fields.set_item("status_code", self.0.status_code)?;
fields.set_item("timestamp", PyNumber(self.0.timestamp))?;
fields.set_item("cooldown_time", PyNumber(self.0.cooldown_time))?;
Ok(fields)
}
}
fn active_to_py(
py: Python<'_>,
active: Vec<(String, Cooldown)>,
) -> PyResult<Vec<Bound<'_, PyTuple>>> {
active
.into_iter()
.map(|(model_id, cooldown)| {
PyTuple::new(
py,
[
model_id.into_pyobject(py)?.into_any(),
PyCooldown(cooldown).into_pyobject(py)?.into_any(),
],
)
})
.collect()
}
#[pymethods]
impl NativeRouterCooldowns {
#[new]
#[pyo3(signature = (redis_kwargs, *, namespace = None))]
fn new(
py: Python<'_>,
redis_kwargs: Option<PythonRedisSettings>,
namespace: Option<String>,
) -> PyResult<Self> {
let client = redis_kwargs
.map(|kwargs| shared_client(py, kwargs, namespace))
.transpose()?
.flatten();
let policy = router_tier_policies(py)?.cooldowns;
Ok(Self(Arc::new(RouterCooldowns::new(client, policy))))
}
fn put(
&self,
py: Python<'_>,
model_id: String,
cooldown: PyCooldown,
ttl_seconds: f64,
) -> PyResult<()> {
let store = self.0.clone();
let scope = call(py, "set_cache")?;
run_sync_value(py, async move {
scope
.scope(store.put(
&model_id,
&cooldown.0,
Duration::from_secs_f64(ttl_seconds.max(0.0)),
))
.await
.map_err(memory_error)
})
}
fn increment_allowed_fails(
&self,
py: Python<'_>,
model_id: String,
suffix: Option<String>,
ttl_seconds: f64,
) -> PyResult<i64> {
let store = self.0.clone();
let scope = call(py, "increment_cache")?;
run_sync_value(py, async move {
scope
.scope(store.increment_allowed_fails(
&model_id,
suffix.as_deref(),
Duration::from_secs_f64(ttl_seconds.max(0.0)),
))
.await
.map_err(memory_error)
})
}
fn active<'py>(
&self,
py: Python<'py>,
model_ids: Vec<String>,
) -> PyResult<Vec<Bound<'py, PyTuple>>> {
let store = self.0.clone();
let scope = call(py, "batch_get_cache")?;
let active = run_sync_value(py, async move {
scope
.scope(store.active(&model_ids))
.await
.map_err(memory_error)
})?;
active_to_py(py, active)
}
fn aactive<'py>(&self, py: Python<'py>, model_ids: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let scope = call(py, "async_batch_get_cache")?;
run_async_value(py, async move {
let active = scope
.scope(store.active(&model_ids))
.await
.map_err(memory_error)?;
Python::attach(|py| {
Ok(active_to_py(py, active)?
.into_iter()
.map(Bound::unbind)
.collect::<Vec<_>>())
})
})
}
fn minimum(&self, py: Python<'_>, model_ids: Vec<String>) -> PyResult<Option<PyNumber>> {
let store = self.0.clone();
let scope = call(py, "batch_get_cache")?;
let minimum = run_sync_value(py, async move {
scope
.scope(store.minimum(&model_ids))
.await
.map_err(memory_error)
})?;
Ok(minimum.map(PyNumber))
}
fn evict_local(&self, model_id: &str) -> PyResult<()> {
self.0.evict_local(model_id).map_err(memory_error)
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(memory_error)
}
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, cooldown) in self.0.snapshot().map_err(memory_error)? {
snapshot.set_item(key, PyCooldown(cooldown))?;
}
Ok(snapshot)
}
}

139
litellm-rust/crates/python-bridge/src/redis/router_usage.rs Normal file
View file

@ -0,0 +1,139 @@
use std::{sync::Arc, time::Duration};
use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use super::{
PyNumber, call,
client::{PythonRedisSettings, shared_client},
memory_error,
policy::router_tier_policies,
};
use crate::execution::run_async_value;
/// One router's global TPM/RPM counters. Every caller on `main` is async.
#[pyclass(frozen, name = "NativeRouterUsage")]
pub(crate) struct NativeRouterUsage(Arc<RouterUsage>);
/// `(deployment_id, model, minute, "tpm" | "rpm")`.
pub(crate) struct PyCounter(Counter);
impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let (deployment_id, model, minute, kind): (String, Option<String>, String, String) =
value.extract()?;
let kind = match kind.as_str() {
"tpm" => Kind::Tpm,
"rpm" => Kind::Rpm,
_ => {
return Err(PyValueError::new_err(format!(
"unknown usage counter kind {kind:?}"
)));
}
};
Ok(Self(Counter {
deployment_id,
model,
minute,
kind,
}))
}
}
#[pymethods]
impl NativeRouterUsage {
#[new]
#[pyo3(signature = (redis_kwargs, *, namespace = None))]
fn new(
py: Python<'_>,
redis_kwargs: Option<PythonRedisSettings>,
namespace: Option<String>,
) -> PyResult<Self> {
let client = redis_kwargs
.map(|kwargs| shared_client(py, kwargs, namespace))
.transpose()?
.flatten();
let policy = router_tier_policies(py)?.usage;
Ok(Self(Arc::new(RouterUsage::new(client, policy))))
}
fn increment<'py>(
&self,
py: Python<'py>,
counter: PyCounter,
value: PyNumber,
ttl_seconds: u64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let scope = call(py, "async_increment")?;
run_async_value(py, async move {
scope
.scope(store.increment(&counter.0, value.0, Duration::from_secs(ttl_seconds)))
.await
.map(PyNumber)
.map_err(memory_error)
})
}
fn increment_deployment<'py>(
&self,
py: Python<'py>,
deployment_id: String,
model: Option<String>,
minute: String,
tokens: PyNumber,
requests: i64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let scope = call(py, "async_increment_pipeline")?;
run_async_value(py, async move {
scope
.scope(store.increment_deployment(
&deployment_id,
model.as_deref(),
&minute,
tokens.0,
requests,
))
.await
.map_err(memory_error)
})
}
fn get_many<'py>(
&self,
py: Python<'py>,
counters: Vec<PyCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = counters
.into_iter()
.map(|counter| counter.0)
.collect::<Vec<_>>();
let scope = call(py, "async_batch_get_cache")?;
run_async_value(py, async move {
let read = scope
.scope(store.get_many(&counters))
.await
.map_err(memory_error)?;
Ok(read
.into_iter()
.map(|count| count.map(PyNumber))
.collect::<Vec<_>>())
})
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(memory_error)
}
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, count) in self.0.snapshot().map_err(memory_error)? {
snapshot.set_item(key, PyNumber(count))?;
}
Ok(snapshot)
}
}

1
litellm-rust/crates/redis-client/Cargo.toml
View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
tracing.workspace = true
litellm-circuit-breaker.workspace = true
litellm-secret-value.workspace = true
litellm-redis-types.workspace = true

39
litellm-rust/crates/redis-client/src/client.rs
View file

@ -1,10 +1,10 @@
use std::{
future::Future,
sync::Arc,
sync::{Arc, Mutex, PoisonError},
time::{Duration, Instant},
};
use litellm_circuit_breaker::{Admission, Breaker, Failure};
use litellm_circuit_breaker::{Admission, Breaker, Failure, State};
use redis::{Cmd, FromRedisValue, RedisError, ToRedisArgs, ToSingleRedisArg, Value};
use crate::{
@ -20,6 +20,8 @@ pub struct RedisClient {
namespace: Namespace,
breaker: Option<Breaker>,
events: Option<Arc<dyn RedisEvents>>,
health_check: Option<Duration>,
last_used: Arc<Mutex<Instant>>,
}
#[derive(Debug, PartialEq, Eq)]
@ -36,9 +38,16 @@ impl RedisClient {
namespace: config.namespace,
breaker: config.breaker,
events: config.events,
health_check: config.health_check,
last_used: Arc::new(Mutex::new(Instant::now())),
})
}
/// The breaker's current state, or none when the breaker is disabled.
pub fn breaker_state(&self) -> Option<State> {
self.breaker.as_ref().map(Breaker::state)
}
pub async fn ping(&self) -> Result<(), Error> {
self.query("PING", &commands::ping()).await
}
@ -159,6 +168,7 @@ impl RedisClient {
.collect();
}
};
self.check_health().await;
let results = pipeline.run(&self.connection).await;
let transport = results.iter().find_map(|result| {
result
@ -216,6 +226,7 @@ impl RedisClient {
let started = Instant::now();
let result = match self.admit() {
Ok(admission) => {
self.check_health().await;
let result = call
.await
.map_err(|error| Error::classify(error, self.connection.deployment()));
@ -229,6 +240,10 @@ impl RedisClient {
}
fn observe(&self, operation: Operation, started: Instant, outcome: Outcome) {
*self
.last_used
.lock()
.unwrap_or_else(PoisonError::into_inner) = Instant::now();
if let Some(events) = &self.events {
events.call(&CallEvent {
operation,
@ -238,6 +253,26 @@ impl RedisClient {
}
}
/// redis-py pings a connection that sat idle longer than `health_check_interval` before
/// sending the next command on it, and reconnects when the ping fails. The ping is not a
/// call of the caller's and never reaches the breaker or the events.
async fn check_health(&self) {
let Some(interval) = self.health_check else {
return;
};
let idle = self
.last_used
.lock()
.unwrap_or_else(PoisonError::into_inner)
.elapsed();
if idle < interval {
return;
}
if let Err(error) = self.connection.ping_every_node().await {
tracing::debug!(%error, "idle Redis connection failed its health check, the next command reconnects");
}
}
fn admit(&self) -> Result<Option<Admission>, Error> {
self.breaker
.as_ref()

77
litellm-rust/crates/redis-client/src/config.rs
View file

@ -24,6 +24,10 @@ pub struct RedisConfig {
pub timeouts: Timeouts,
pub breaker: Option<Breaker>,
pub events: Option<Arc<dyn RedisEvents>>,
/// Ping every node before the next command once the client sat idle this long, like
/// redis-py's `health_check_interval`, so a connection the server dropped silently is
/// replaced before a real command pays for it.
pub health_check: Option<Duration>,
}
#[derive(Clone, Copy, Debug, PartialEq)]
@ -95,17 +99,44 @@ pub(crate) fn tls_certificates(tls: &TlsOptions) -> Result<TlsCertificates, Conn
}),
(None, _) => None,
};
let from_file: Option<Vec<u8>> = tls.ca_certs.as_ref().map(read).transpose()?;
let from_directory: Vec<Vec<u8>> = tls
.ca_path
.as_ref()
.map(read_directory)
.transpose()?
.unwrap_or_default();
let configured: Vec<Vec<u8>> = from_file
.into_iter()
.chain(tls.ca_data.iter().map(|pem| pem.clone().into_bytes()))
.chain(from_directory)
.collect();
Ok(TlsCertificates {
client_tls,
root_cert: tls
.ca_certs
.as_ref()
.map(read)
.transpose()?
.map(with_system_roots),
root_cert: (!configured.is_empty()).then(|| with_system_roots(configured.join(&b'\n'))),
})
}
/// Every regular file in `path`, in name order, as redis-py hands `ssl_ca_path` to
/// `load_verify_locations(capath=...)`
fn read_directory(path: &PathBuf) -> Result<Vec<Vec<u8>>, ConnectError> {
let tls_file = |source| ConnectError::TlsFile {
path: path.display().to_string(),
source,
};
let mut files: Vec<PathBuf> = std::fs::read_dir(path)
.map_err(tls_file)?
.map(|entry| entry.map(|entry| entry.path()))
.collect::<Result<_, _>>()
.map_err(tls_file)?;
files.sort();
files
.iter()
.filter(|file| file.is_file())
.map(read)
.collect()
}
pub(crate) fn apply_tls(
tls: &TlsOptions,
info: ConnectionInfo,
@ -179,9 +210,16 @@ pub(crate) fn connection_info(
let (info, tls) = match endpoint {
Endpoint::Url { url, tls } => (url.expose().into_connection_info()?, Some(tls)),
Endpoint::Host(host) => {
let info = addr(&host.host, host.port, host.tls.as_ref()).into_connection_info()?;
let (addr, tls) = match &host.unix_socket_path {
Some(path) => (ConnectionAddr::Unix(path.clone()), None),
None => (
addr(&host.host, host.port, host.tls.as_ref()),
host.tls.as_ref(),
),
};
let info = addr.into_connection_info()?;
let settings = info.redis_settings().clone().set_db(host.db);
(info.set_redis_settings(settings), host.tls.as_ref())
(info.set_redis_settings(settings), tls)
}
Endpoint::Cluster(_) | Endpoint::Sentinel(_) => {
unreachable!("cluster and sentinel endpoints build their own connection info")
@ -210,3 +248,26 @@ fn with_static_credentials(info: ConnectionInfo, credentials: &Credentials) -> C
};
info.set_redis_settings(settings)
}
#[cfg(test)]
mod tests {
use litellm_redis_types::HostEndpoint;
use super::*;
#[test]
fn a_unix_socket_path_replaces_host_port_and_tls_like_redis_py() {
let endpoint = Endpoint::Host(HostEndpoint {
host: "ignored".to_owned(),
port: 1,
tls: Some(TlsOptions::verified()),
db: 3,
unix_socket_path: Some("/tmp/redis.sock".into()),
});
let info = connection_info(&endpoint, &Credentials::NONE).unwrap();
assert_eq!(*info.addr(), ConnectionAddr::Unix("/tmp/redis.sock".into()));
assert_eq!(info.redis_settings().db(), 3);
}
}

26
litellm-rust/crates/redis-client/src/connection.rs
View file

@ -12,6 +12,7 @@ use redis::{
aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig, MultiplexedConnection},
cluster::{ClusterClientBuilder, ClusterConfig},
cluster_async::ClusterConnection,
cluster_routing::{MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo},
io::tcp::TcpSettings,
sentinel::{SentinelClient, SentinelClientBuilder, SentinelServerType},
};
@ -131,6 +132,30 @@ impl Connection {
}
}
impl Connection {
/// `PING` on the one multiplexed connection, or on every cluster node, so a dead socket is
/// found and replaced by redis-rs before a real command uses it.
pub(crate) async fn ping_every_node(&self) -> RedisResult<()> {
let ping = redis::cmd("PING");
match self {
Self::Cluster(connection) => connection
.clone()
.route_command(
ping,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllNodes,
Some(ResponsePolicy::AllSucceeded),
)),
)
.await
.map(drop),
Self::Standalone(_) | Self::Sentinel(_) => {
ping.query_async::<Value>(&mut self.clone()).await.map(drop)
}
}
}
}
impl ConnectionLike for Connection {
fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> {
match self {
@ -277,6 +302,7 @@ impl SentinelConnection {
port,
tls: self.0.endpoint.tls.clone(),
db: self.0.endpoint.db,
unix_socket_path: None,
});
let info =
config::connection_info(&endpoint, &self.0.credentials).map_err(

97
litellm-rust/crates/redis-client/tests/client.rs
View file

@ -5,7 +5,7 @@ use std::{
time::{Duration, Instant},
};
use litellm_redis_testing::{RedisDeployment, RedisServer, Topology, eventually};
use litellm_redis_testing::{RedisCluster, RedisDeployment, RedisServer, Topology, eventually};
use rstest::rstest;
use litellm_circuit_breaker::{Breaker, Config, State};
@ -254,6 +254,7 @@ async fn outage_events_report_the_failure_kind(
breaker: None,
events: Some(recorded.clone()),
credentials: Credentials::NONE,
health_check: None,
})
.unwrap();
client.ping().await.unwrap();
@ -310,6 +311,7 @@ async fn every_call_reports_its_operation_and_outcome(
timeouts: Timeouts::default(),
breaker: None,
events: Some(recorded.clone()),
health_check: None,
})
.unwrap();
client.rpush("list", "item", KEEP).await.unwrap();
@ -739,6 +741,7 @@ async fn a_host_endpoint_connects_to_its_database() {
timeouts: Timeouts::default(),
breaker: None,
events: None,
health_check: None,
})
};
let url = |url: String| Endpoint::Url {
@ -750,6 +753,7 @@ async fn a_host_endpoint_connects_to_its_database() {
port,
tls: None,
db: 3,
unix_socket_path: None,
}))
.unwrap();
let in_db_0 = connect(url(server.url().to_owned())).unwrap();
@ -766,3 +770,94 @@ async fn a_host_endpoint_connects_to_its_database() {
Some("3")
);
}
async fn ping_calls(cluster: &RedisCluster) -> Vec<u64> {
cluster
.cli_outputs(&["INFO", "commandstats"])
.await
.unwrap()
.iter()
.map(|info| {
info.lines()
.find_map(|line| line.strip_prefix("cmdstat_ping:calls="))
.and_then(|rest| rest.split(',').next())
.and_then(|calls| calls.trim().parse().ok())
.unwrap_or(0)
})
.collect()
}
fn health_checked(
server: &RedisDeployment,
interval: Duration,
timeout: Duration,
breaker: Option<Breaker>,
events: Option<Arc<Recorded>>,
) -> RedisClient {
RedisClient::new(RedisConfig {
endpoint: endpoint(server),
credentials: Credentials::NONE,
namespace: Namespace::default(),
timeouts: Timeouts::new(timeout, None),
breaker,
events: events.map(|events| events as Arc<dyn RedisEvents>),
health_check: Some(interval),
})
.unwrap()
}
// redis-py 5.3.1 redis/asyncio/connection.py check_health: a connection idle longer than
// health_check_interval gets a PING before the next command, a busy one does not
#[rstest]
#[tokio::test]
async fn an_idle_client_pings_every_node_before_the_next_command_like_redis_py() {
let server = RedisDeployment::start(Topology::Cluster).await.unwrap();
let RedisDeployment::Cluster(cluster) = &server else {
unreachable!("started a cluster")
};
let client = health_checked(&server, Duration::from_secs(3), DEFAULT_TIMEOUT, None, None);
client.set("k", "1", Duration::from_secs(60)).await.unwrap();
let before = ping_calls(cluster).await;
client.set("k", "2", Duration::from_secs(60)).await.unwrap();
let busy = ping_calls(cluster).await;
tokio::time::sleep(Duration::from_millis(3500)).await;
client.set("k", "3", Duration::from_secs(60)).await.unwrap();
let idle = ping_calls(cluster).await;
assert_eq!(busy, before);
assert_eq!(
idle,
before.iter().map(|calls| calls + 1).collect::<Vec<_>>(),
"before {before:?}"
);
}
#[rstest]
#[tokio::test]
async fn a_failed_health_check_is_neither_a_call_nor_a_breaker_failure() {
let server = RedisDeployment::start(Topology::Cluster).await.unwrap();
let recorded = Arc::new(Recorded::default());
let breaker = breaker(2);
let client = health_checked(
&server,
Duration::from_millis(200),
Duration::from_millis(300),
Some(breaker.clone()),
Some(recorded.clone()),
);
client.set("k", "1", Duration::from_secs(60)).await.unwrap();
server.freeze().await.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
let failed = client.get::<String>("k").await;
server.thaw().await.unwrap();
assert!(matches!(failed, Err(Error::Transport(_))), "{failed:?}");
assert!(
matches!(breaker.state(), State::Closed { timeouts: 1, .. }),
"{:?}",
breaker.state()
);
assert_eq!(recorded.0.lock().unwrap().len(), 2);
}

23
litellm-rust/crates/redis-client/tests/config.rs
View file

@ -21,9 +21,18 @@ fn debug_output_never_shows_the_secret(#[case] debug: String) {
#[derive(Clone, Copy, Debug)]
enum Ca {
Trusted,
TrustedData,
TrustedDirectory,
SystemRoots,
}
fn directory_holding(ca: &std::path::Path) -> std::path::PathBuf {
let directory = std::env::temp_dir().join(format!("litellm-ca-dir-{}", std::process::id()));
std::fs::create_dir_all(&directory).unwrap();
std::fs::copy(ca, directory.join("ca.pem")).unwrap();
directory
}
async fn ping_over_tls(
require_client_certificate: bool,
host: &str,
@ -41,6 +50,8 @@ async fn ping_over_tls(
let files = server.files();
let tls = TlsOptions {
ca_certs: matches!(ca, Ca::Trusted).then(|| files.ca.clone()),
ca_data: matches!(ca, Ca::TrustedData).then(|| std::fs::read_to_string(&files.ca).unwrap()),
ca_path: matches!(ca, Ca::TrustedDirectory).then(|| directory_holding(&files.ca)),
certfile: client_certificate.then(|| files.client_cert.clone()),
keyfile: client_certificate.then(|| files.client_key.clone()),
verify_certificate,
@ -52,12 +63,14 @@ async fn ping_over_tls(
port: server.port(),
tls: Some(tls),
db: 0,
unix_socket_path: None,
}),
credentials: Credentials::NONE,
namespace: litellm_redis_client::Namespace::default(),
timeouts: litellm_redis_client::Timeouts::default(),
breaker: None,
events: None,
health_check: None,
})
.unwrap();
client.ping().await
@ -71,6 +84,14 @@ async fn ping_over_tls(
false,
true
)]
#[case::trusted_ca_as_pem_text_like_ssl_ca_data("127.0.0.1", Ca::TrustedData, true, false, true)]
#[case::trusted_ca_in_a_directory_like_ssl_ca_path(
"127.0.0.1",
Ca::TrustedDirectory,
true,
false,
true
)]
#[case::hostname_check_with_a_matching_name("localhost", Ca::Trusted, true, true, true)]
#[case::hostname_check_against_an_unlisted_address("127.0.0.1", Ca::Trusted, true, true, false)]
#[case::an_untrusted_ca_is_refused("127.0.0.1", Ca::SystemRoots, true, false, false)]
@ -128,12 +149,14 @@ fn a_missing_tls_file_names_the_path() {
port: 6379,
tls: Some(tls),
db: 0,
unix_socket_path: None,
}),
credentials: Credentials::NONE,
namespace: litellm_redis_client::Namespace::default(),
timeouts: litellm_redis_client::Timeouts::default(),
breaker: None,
events: None,
health_check: None,
})
.err()
.unwrap();

3
litellm-rust/crates/redis-client/tests/connection.rs
View file

@ -91,12 +91,14 @@ async fn an_out_of_range_database_is_a_server_rejection() {
port,
tls: None,
db: 999,
unix_socket_path: None,
}),
credentials: Credentials::NONE,
namespace: Namespace::default(),
timeouts: Timeouts::new(Duration::from_secs(2), None),
breaker: None,
events: None,
health_check: None,
})
.unwrap();
let result = client.ping().await;
@ -145,6 +147,7 @@ fn sentinel_client(sentinel: &RedisSentinel) -> RedisClient {
timeouts: Timeouts::new(Duration::from_secs(1), None),
breaker: None,
events: None,
health_check: None,
})
.unwrap()
}

2
litellm-rust/crates/redis-client/tests/support/mod.rs
View file

@ -50,6 +50,7 @@ pub fn connect_guarded(
timeouts: Timeouts::new(timeout, None),
breaker,
events: None,
health_check: None,
})
.unwrap()
}
@ -62,6 +63,7 @@ pub fn connect_as(deployment: &RedisDeployment, credentials: Credentials) -> Red
timeouts: Timeouts::new(Duration::from_secs(2), None),
breaker: None,
events: None,
health_check: None,
})
.unwrap()
}

1
litellm-rust/crates/redis-connect/Cargo.toml
View file

@ -17,6 +17,7 @@ litellm-circuit-breaker.workspace = true
litellm-config.workspace = true
litellm-core-utils.workspace = true
litellm-http.workspace = true
rustls.workspace = true
litellm-redis-client.workspace = true
futures-util.workspace = true
serde.workspace = true

79
litellm-rust/crates/redis-connect/src/connect.rs
View file

@ -1,11 +1,13 @@
use std::time::Duration;
use std::sync::Arc;
use litellm_circuit_breaker::{Breaker, Config as BreakerConfig};
use litellm_config::{RedisAuth, ResolvedRedis};
use litellm_config::{Expected, RedisAuth, RedisError, ResolvedRedis};
use litellm_core_utils::settings::Lookup;
use litellm_redis_client::{
Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Namespace, RedisClient, RedisConfig,
RedisEvents, Timeouts,
Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Endpoint, Namespace, RedisClient,
RedisConfig, RedisEvents, Timeouts,
};
use crate::{
@ -22,7 +24,13 @@ pub struct HostServices {
pub environment: SharedEnvironment,
}
// litellm/constants.py REDIS_CLUSTER_HEALTH_CHECK_INTERVAL, the default _redis.py sets for cluster clients only
const CLUSTER_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(25);
pub fn connect(resolved: ResolvedRedis, host: HostServices) -> Result<RedisClient, Error> {
if rustls::crypto::CryptoProvider::get_default().is_none() {
let _ = rustls::crypto::ring::default_provider().install_default();
}
tracing::info!(endpoint = ?resolved.endpoint, namespace = ?resolved.namespace, "redis client configured");
let credentials = credentials(resolved.auth, &host)?;
let events = host.events;
@ -37,6 +45,13 @@ pub fn connect(resolved: ResolvedRedis, host: HostServices) -> Result<RedisClien
}),
)
});
let health_check = match &resolved.endpoint {
Endpoint::Cluster(_) => resolved
.health_check_interval
.or(Some(CLUSTER_HEALTH_CHECK_INTERVAL)),
_ => resolved.health_check_interval,
}
.filter(|interval| !interval.is_zero());
Ok(RedisClient::new(RedisConfig {
endpoint: resolved.endpoint,
credentials,
@ -47,6 +62,7 @@ pub fn connect(resolved: ResolvedRedis, host: HostServices) -> Result<RedisClien
),
breaker,
events,
health_check,
})?)
}
@ -70,3 +86,60 @@ fn credentials(auth: RedisAuth, host: &HostServices) -> Result<Credentials, Erro
)),
})
}
/// The breaker `RedisCache` builds from `REDIS_CIRCUIT_BREAKER_*`, or none when it is disabled.
pub fn breaker_config(lookup: &dyn Lookup) -> Result<Option<BreakerConfig>, RedisError> {
let enabled = lookup
.get("REDIS_CIRCUIT_BREAKER_ENABLED")
.is_none_or(|value| value.eq_ignore_ascii_case("true"));
if !enabled {
return Ok(None);
}
let defaults = BreakerConfig::default();
Ok(Some(BreakerConfig {
failure_threshold: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
Expected::Integer,
|text| text.parse().ok(),
)?
.unwrap_or(defaults.failure_threshold),
recovery_timeout: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.recovery_timeout),
timeout_min_duration: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.timeout_min_duration),
}))
}
fn parse_seconds(text: &str) -> Option<Duration> {
text.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
}
fn environment_value<T>(
lookup: &dyn Lookup,
name: &str,
expected: Expected,
parse: impl Fn(&str) -> Option<T>,
) -> Result<Option<T>, RedisError> {
lookup
.get(name)
.map(|value| {
parse(&value).ok_or_else(|| RedisError::Invalid {
setting: name.to_owned(),
expected,
})
})
.transpose()
}

2
litellm-rust/crates/redis-connect/src/lib.rs
View file

@ -4,7 +4,7 @@ mod events;
mod refresh;
mod sources;
pub use connect::{HostServices, SharedEnvironment, connect};
pub use connect::{HostServices, SharedEnvironment, breaker_config, connect};
pub use error::{CredentialStage, Error};
pub use events::TracingEvents;
pub(crate) use refresh::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture};

90
litellm-rust/crates/redis-connect/src/refresh.rs
View file

@ -250,9 +250,11 @@ mod tests {
use super::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture, refresh_after};
use litellm_redis_client::{
ClusterEndpoint, ClusterNode, Credentials, CredentialsProvider, Endpoint, Namespace,
RedisClient, RedisConfig, Timeouts, TlsOptions,
RedisClient, RedisConfig, SentinelEndpoint, Timeouts, TlsOptions,
};
use litellm_redis_testing::{
ClientConnection, RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually,
};
use litellm_redis_testing::{ClientConnection, RedisDeployment, Topology, eventually};
#[rstest]
#[case::three_quarters_of_the_lifetime(
@ -598,6 +600,7 @@ mod tests {
timeouts: Timeouts::default(),
breaker: None,
events: None,
health_check: None,
})
.unwrap();
client.set("k", "1", Duration::from_secs(60)).await.unwrap();
@ -633,4 +636,87 @@ mod tests {
assert_eq!(connection_users(&reconnected), BTreeSet::from(["bob"]));
assert!(connection_ids(&reconnected).is_disjoint(&connection_ids(&after)));
}
/// Cloud credentials rotate while Sentinel fails over: the refreshed user must be what the
/// client presents to the master Sentinel promotes, not the one it connected with first.
#[rstest]
#[tokio::test]
async fn a_refreshed_user_reaches_the_master_sentinel_promotes() {
let sentinel = RedisSentinel::start().await.unwrap();
sentinel
.cli_on_data_nodes(&[
"ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all",
])
.await
.unwrap();
sentinel
.cli_on_data_nodes(&[
"ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all",
])
.await
.unwrap();
let provider = RefreshingCredentials::new(Scripted {
replies: Mutex::new(vec![
user("alice", "first", Duration::from_secs(2)),
user("bob", "second", Duration::from_secs(3600)),
]),
calls: Arc::new(AtomicUsize::new(0)),
});
let (host, port) = sentinel.sentinel();
let client = RedisClient::new(RedisConfig {
endpoint: Endpoint::Sentinel(SentinelEndpoint {
sentinels: vec![ClusterNode {
host: host.to_owned(),
port,
}],
service_name: SERVICE_NAME.to_owned(),
sentinel_password: None,
tls: None,
db: 0,
}),
credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))),
namespace: Namespace::default(),
timeouts: Timeouts::new(Duration::from_secs(1), None),
breaker: None,
events: None,
health_check: None,
})
.unwrap();
client
.set("before", "1", Duration::from_secs(60))
.await
.unwrap();
let master = sentinel.master_port().await.unwrap();
let app_users = |connections: &[ClientConnection]| -> BTreeSet<String> {
connections
.iter()
.filter(|connection| connection.user != "default")
.map(|connection| connection.user.clone())
.collect()
};
eventually(async || {
let connections = sentinel.clients_on(master).await.ok()?;
(app_users(&connections) == BTreeSet::from(["bob".to_owned()])).then_some(())
})
.await;
sentinel
.cli_on_data_nodes(&["ACL", "DELUSER", "alice"])
.await
.unwrap();
let promoted = sentinel.failover().await.unwrap();
eventually(async || client.set("after", "2", Duration::from_secs(60)).await.ok()).await;
let on_promoted = sentinel.clients_on(promoted).await.unwrap();
assert_eq!(promoted, sentinel.replica_port());
assert_eq!(
app_users(&on_promoted),
BTreeSet::from(["bob".to_owned()]),
"{on_promoted:?}"
);
assert_eq!(
client.get::<String>("after").await.unwrap().as_deref(),
Some("2")
);
}
}

6
litellm-rust/crates/redis-store/AGENTS.md Normal file
View file

@ -0,0 +1,6 @@
- One module per Redis family (`router_cooldowns.rs`, `router_usage.rs`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys
- 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 read-through family builds its own `DualCache` from a `TierPolicy` the host passes in. Its `pub const POLICY` is `main`'s defaults, and the Python host overrides what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `RedisClient` (or none), never a cache. Redis-only families are plain functions over `&RedisClient`
- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not
- `tests/fixtures/router_wire.json` is the wire contract, checked here against real Redis and in `tests/unit/redis/stores/` against Python's own codec. Change both sides together
- Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages

22
litellm-rust/crates/redis-store/Cargo.toml Normal file
View file

@ -0,0 +1,22 @@
[package]
name = "litellm-redis-store"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-memory.workspace = true
litellm-python-compat.workspace = true
litellm-redis-client.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tracing.workspace = true
[dev-dependencies]
litellm-circuit-breaker.workspace = true
litellm-redis-testing.workspace = true
litellm-secret-value.workspace = true
rstest.workspace = true
tokio.workspace = true

20
litellm-rust/crates/redis-store/src/codec.rs Normal file
View file

@ -0,0 +1,20 @@
use litellm_python_compat::{Value, json, literal};
#[derive(Debug, thiserror::Error)]
pub enum DecodeError {
#[error("Redis value is not UTF-8")]
NotUtf8,
#[error("Redis value is neither JSON nor a Python literal: {0}")]
Unparsed(#[source] litellm_python_compat::Error),
#[error("Redis value has the wrong shape: {0}")]
Shape(&'static str),
}
/// `RedisCache._get_cache_logic`: `json.loads`, falling back to `ast.literal_eval`.
pub(crate) fn python_value(raw: &[u8]) -> Result<Value, DecodeError> {
let text = std::str::from_utf8(raw).map_err(|_| DecodeError::NotUtf8)?;
match serde_json::from_str::<serde_json::Value>(text) {
Ok(document) => Ok(json::from_json(document)),
Err(_) => literal::literal_eval(text).map_err(DecodeError::Unparsed),
}
}

7
litellm-rust/crates/redis-store/src/lib.rs Normal file
View file

@ -0,0 +1,7 @@
mod codec;
mod number;
pub mod router_cooldowns;
pub mod router_usage;
pub use codec::DecodeError;
pub use number::Number;

68
litellm-rust/crates/redis-store/src/number.rs Normal file
View file

@ -0,0 +1,68 @@
use std::{cmp::Ordering, ops::Add};
use litellm_python_compat::{BigInt, Value};
/// A Python `int` or `float`, kept apart because `repr` and `json.dumps` write them differently.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Number {
Int(i64),
Float(f64),
}
impl Number {
pub fn as_f64(self) -> f64 {
match self {
Self::Int(value) => value as f64,
Self::Float(value) => value,
}
}
/// Python's `int(value)`, which truncates toward zero.
pub fn truncate(self) -> i64 {
match self {
Self::Int(value) => value,
Self::Float(value) => value.trunc() as i64,
}
}
pub fn to_value(self) -> Value {
match self {
Self::Int(value) => Value::Int(BigInt::from(value)),
Self::Float(value) => Value::Float(value),
}
}
pub fn from_value(value: &Value) -> Option<Self> {
match value {
Value::Int(value) => i64::try_from(value).ok().map(Self::Int),
Value::Float(value) => Some(Self::Float(*value)),
_ => None,
}
}
/// Python's `min` over numbers: the first of the smallest, keeping its type.
pub fn min(values: impl IntoIterator<Item = Self>) -> Option<Self> {
values
.into_iter()
.fold(None, |smallest, value| match smallest {
Some(smallest)
if value.as_f64().partial_cmp(&smallest.as_f64()) != Some(Ordering::Less) =>
{
Some(smallest)
}
_ => Some(value),
})
}
}
/// Python's `+`: two ints stay an int, anything with a float is a float.
impl Add for Number {
type Output = Self;
fn add(self, other: Self) -> Self {
match (self, other) {
(Self::Int(left), Self::Int(right)) => Self::Int(left + right),
(left, right) => Self::Float(left.as_f64() + right.as_f64()),
}
}
}

245
litellm-rust/crates/redis-store/src/router_cooldowns.rs Normal file
View file

@ -0,0 +1,245 @@
use std::time::Duration;
use litellm_cache::{
Clock, DualCache, Error as MemoryError, OnBreakerOpen, TierPolicy, system_clock,
};
use litellm_cache_memory::InMemoryCache;
use litellm_python_compat::{Value, repr::repr};
use litellm_redis_client::{Expiry, RedisClient};
use crate::{DecodeError, Number, codec::python_value};
/// `main`'s cooldown tier: `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS` (1 s), `InMemoryCache`'s
/// 600 s default for backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(1),
backfill_ttl: Duration::from_secs(600),
on_breaker_open: OnBreakerOpen::ServeMemory,
tracked_reads: 1000,
};
const MEMORY_ENTRIES: usize = 200;
const MAX_CORRECTED_TTL: Duration = Duration::from_secs(60);
const CORRECTION_SLACK: Duration = Duration::from_secs(5);
#[derive(Clone, Debug, PartialEq)]
pub struct Cooldown {
pub exception_received: String,
pub status_code: String,
pub timestamp: Number,
pub cooldown_time: Number,
}
impl Cooldown {
fn remaining(&self, now: Duration) -> f64 {
self.timestamp.as_f64() + self.cooldown_time.as_f64() - now.as_secs_f64()
}
/// `str(dict)`, the bytes `RedisCache.set_cache` writes on `main`.
pub fn encode(&self) -> String {
let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value);
repr(&Value::Dict(vec![
field(
"exception_received",
Value::Str(self.exception_received.clone()),
),
field("status_code", Value::Str(self.status_code.clone())),
field("timestamp", self.timestamp.to_value()),
field("cooldown_time", self.cooldown_time.to_value()),
]))
}
/// Reads both `main`'s `str(dict)` and a JSON object, as `RedisCache._get_cache_logic` does.
pub fn decode(raw: &[u8]) -> Result<Self, DecodeError> {
let Value::Dict(fields) = python_value(raw)? else {
return Err(DecodeError::Shape("a cooldown is a dict"));
};
let field = |name: &str| {
fields
.iter()
.find(|(key, _)| matches!(key, Value::Str(key) if key == name))
.map(|(_, value)| value)
};
let text = |name: &'static str| match field(name) {
Some(Value::Str(text)) => Ok(text.clone()),
_ => Err(DecodeError::Shape(name)),
};
let number = |name: &'static str| {
field(name)
.and_then(Number::from_value)
.ok_or(DecodeError::Shape(name))
};
Ok(Self {
exception_received: text("exception_received")?,
status_code: text("status_code")?,
timestamp: number("timestamp")?,
cooldown_time: number("cooldown_time")?,
})
}
}
fn key(model_id: &str) -> String {
format!("deployment:{model_id}:cooldown")
}
fn allowed_fails_key(model_id: &str, suffix: Option<&str>) -> String {
match suffix.filter(|suffix| !suffix.is_empty()) {
Some(suffix) => format!("deployment:{model_id}:allowed_fails:{suffix}"),
None => format!("deployment:{model_id}:allowed_fails"),
}
}
fn memory<V: Clone>(policy: TierPolicy, now: &Clock) -> InMemoryCache<V> {
let now = now.clone();
InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || {
now()
})
}
/// `BaseCache.get_ttl` sends `int(ttl)`, so Redis sees whole seconds.
fn whole_seconds(ttl: Duration) -> Duration {
Duration::from_secs(ttl.as_secs())
}
/// One router's cooldowns and allowed-fails counters: memory per router, Redis shared by every pod.
pub struct RouterCooldowns {
cooldowns: DualCache<InMemoryCache<Cooldown>, RedisClient>,
fails: DualCache<InMemoryCache<i64>, RedisClient>,
}
impl RouterCooldowns {
pub fn new(client: Option<RedisClient>, policy: TierPolicy) -> Self {
Self::with_clock(client, policy, system_clock())
}
pub fn with_clock(client: Option<RedisClient>, policy: TierPolicy, now: Clock) -> Self {
Self {
cooldowns: DualCache::read_through(
memory(policy, &now),
client.clone(),
policy,
now.clone(),
),
fails: DualCache::read_through(memory(policy, &now), client, policy, now),
}
}
/// Memory first, then a `SET ... EX int(ttl)` whose failure only logs, like `DualCache.set_cache`.
pub async fn put(
&self,
model_id: &str,
cooldown: &Cooldown,
ttl: Duration,
) -> Result<(), MemoryError> {
let key = key(model_id);
self.cooldowns.set_memory(&key, cooldown.clone(), ttl)?;
if let Some(remote) = self.cooldowns.remote()
&& let Err(error) = remote
.set(&key, cooldown.encode(), whole_seconds(ttl))
.await
{
tracing::warn!(%error, "writing a cooldown to Redis failed");
}
Ok(())
}
/// The fleet-wide fail count. A Redis failure falls back to this router's own count.
pub async fn increment_allowed_fails(
&self,
model_id: &str,
suffix: Option<&str>,
ttl: Duration,
) -> Result<i64, MemoryError> {
let key = allowed_fails_key(model_id, suffix);
let local = self
.fails
.update_memory(&key, ttl, |count| count.unwrap_or(0) + 1)?;
let Some(remote) = self.fails.remote() else {
return Ok(local);
};
match remote
.incr_by(&key, 1, Expiry::KeepExisting(whole_seconds(ttl)))
.await
{
Ok(count) => Ok(count),
Err(error) => {
tracing::warn!(%error, "allowed_fails counter fell back to this worker's in-memory count");
Ok(local)
}
}
}
/// The deployments still cooling down. An entry whose cooldown already ended is dropped
/// from memory, and one memory would keep past its end is shortened, at most to 60 s.
pub async fn active(
&self,
model_ids: &[String],
) -> Result<Vec<(String, Cooldown)>, MemoryError> {
let read = self.read(model_ids).await?;
let now = self.cooldowns.now();
let corrected = model_ids
.iter()
.zip(read)
.filter_map(|(model_id, cooldown)| cooldown.map(|cooldown| (model_id, cooldown)))
.map(|(model_id, cooldown)| {
Ok(self
.correct(model_id, &cooldown, now)?
.then(|| (model_id.clone(), cooldown)))
})
.collect::<Result<Vec<_>, MemoryError>>()?;
Ok(corrected.into_iter().flatten().collect())
}
/// The shortest stored cooldown, ended ones included, as `main`'s `get_min_cooldown` reads it.
pub async fn minimum(&self, model_ids: &[String]) -> Result<Option<Number>, MemoryError> {
let read = self.read(model_ids).await?;
Ok(Number::min(
read.into_iter()
.flatten()
.map(|cooldown| cooldown.cooldown_time),
))
}
pub fn evict_local(&self, model_id: &str) -> Result<(), MemoryError> {
self.cooldowns.delete_memory(&key(model_id))
}
pub fn flush_local_counters(&self) -> Result<(), MemoryError> {
self.fails.flush_memory()
}
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
/// in-memory router state.
pub fn snapshot(&self) -> Result<Vec<(String, Cooldown)>, MemoryError> {
self.cooldowns.memory().entries()
}
async fn read(&self, model_ids: &[String]) -> Result<Vec<Option<Cooldown>>, MemoryError> {
let keys = model_ids
.iter()
.map(|model_id| key(model_id))
.collect::<Vec<_>>();
self.cooldowns.get_many(&keys, Cooldown::decode).await
}
fn correct(
&self,
model_id: &str,
cooldown: &Cooldown,
now: Duration,
) -> Result<bool, MemoryError> {
let key = key(model_id);
let remaining = cooldown.remaining(now);
if remaining <= 0.0 {
self.cooldowns.delete_memory(&key)?;
return Ok(false);
}
let remaining = Duration::from_secs_f64(remaining);
let memory = self.cooldowns.memory();
if memory.expires_at(&key)? > Some(now + remaining + CORRECTION_SLACK) {
self.cooldowns.delete_memory(&key)?;
self.cooldowns
.set_memory(&key, cooldown.clone(), remaining.min(MAX_CORRECTED_TTL))?;
}
Ok(true)
}
}

198
litellm-rust/crates/redis-store/src/router_usage.rs Normal file
View file

@ -0,0 +1,198 @@
use std::time::Duration;
use litellm_cache::{
Clock, DualCache, Error as MemoryError, OnBreakerOpen, TierPolicy, system_clock,
};
use litellm_cache_memory::InMemoryCache;
use litellm_redis_client::{Expiry, RedisClient, Value};
use crate::{DecodeError, Number, codec::python_value};
/// `main`'s usage tier: `DualCache`'s 10 s default refresh, `InMemoryCache`'s 600 s default for
/// backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(10),
backfill_ttl: Duration::from_secs(600),
on_breaker_open: OnBreakerOpen::ServeMemory,
tracked_reads: 1000,
};
const MEMORY_ENTRIES: usize = 200;
/// `RoutingArgs.ttl`, the window of a per-minute counter.
pub const COUNTER_TTL: Duration = Duration::from_secs(60);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Kind {
Tpm,
Rpm,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Counter {
pub deployment_id: String,
pub model: Option<String>,
pub minute: String,
pub kind: Kind,
}
impl Counter {
/// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`.
fn key(&self) -> String {
let kind = match self.kind {
Kind::Tpm => "tpm",
Kind::Rpm => "rpm",
};
let model = self.model.as_deref().unwrap_or("None");
format!(
"global_router:{}:{model}:{kind}:{}",
self.deployment_id, self.minute
)
}
}
fn decode(raw: &[u8]) -> Result<Number, DecodeError> {
Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape("a usage counter is a number"))
}
fn whole_seconds(ttl: Duration) -> Duration {
Duration::from_secs(ttl.as_secs())
}
/// One router's global TPM/RPM counters: memory per router, Redis shared by every pod.
pub struct RouterUsage(DualCache<InMemoryCache<Number>, RedisClient>);
impl RouterUsage {
pub fn new(client: Option<RedisClient>, policy: TierPolicy) -> Self {
Self::with_clock(client, policy, system_clock())
}
pub fn with_clock(client: Option<RedisClient>, policy: TierPolicy, now: Clock) -> Self {
let clock = now.clone();
let memory =
InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || {
clock()
});
Self(DualCache::read_through(memory, client, policy, now))
}
/// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL.
/// A Redis failure answers with the memory count.
pub async fn increment(
&self,
counter: &Counter,
value: Number,
ttl: Duration,
) -> Result<Number, MemoryError> {
let key = counter.key();
let local = self.add_local(&key, value, ttl)?;
let Some(remote) = self.0.remote() else {
return Ok(local);
};
match remote
.incr_by_float(
&key,
value.as_f64(),
Expiry::KeepExisting(whole_seconds(ttl)),
)
.await
{
Ok(total) => Ok(Number::Float(total)),
Err(error) => {
tracing::warn!(%error, "Redis usage increment failed, falling back to the in-memory count");
Ok(local)
}
}
}
/// One deployment's tokens and requests: memory, then one pipeline that resets each
/// counter's 60 s window, mirrored back into memory as `int` totals.
pub async fn increment_deployment(
&self,
deployment_id: &str,
model: Option<&str>,
minute: &str,
tokens: Number,
requests: i64,
) -> Result<(), MemoryError> {
let counter = |kind| Counter {
deployment_id: deployment_id.to_owned(),
model: model.map(str::to_owned),
minute: minute.to_owned(),
kind,
};
let increments = [
(counter(Kind::Tpm).key(), tokens),
(counter(Kind::Rpm).key(), Number::Int(requests)),
];
let local = increments
.iter()
.map(|(key, value)| self.add_local(key, *value, COUNTER_TTL))
.collect::<Result<Vec<_>, _>>()?;
let Some(remote) = self.0.remote() else {
return Ok(());
};
let pipeline = increments
.iter()
.fold(remote.pipeline(), |pipeline, (key, value)| {
pipeline.incr_by_float(key, value.as_f64(), Expiry::Refresh(COUNTER_TTL))
});
let totals = match remote
.run(&pipeline)
.await
.into_iter()
.map(total)
.collect::<Result<Vec<_>, _>>()
{
Ok(totals) => totals,
Err(error) => {
tracing::warn!(%error, "Redis usage pipeline failed, keeping the in-memory counts");
local
}
};
for ((key, _), total) in increments.iter().zip(totals) {
self.0
.set_memory(key, Number::Int(total.truncate()), COUNTER_TTL)?;
}
Ok(())
}
pub async fn get_many(&self, counters: &[Counter]) -> Result<Vec<Option<Number>>, MemoryError> {
let keys = counters.iter().map(Counter::key).collect::<Vec<_>>();
self.0.get_many(&keys, decode).await
}
pub fn flush_local_counters(&self) -> Result<(), MemoryError> {
self.0.flush_memory()
}
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
/// in-memory router state.
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, MemoryError> {
self.0.memory().entries()
}
fn add_local(&self, key: &str, value: Number, ttl: Duration) -> Result<Number, MemoryError> {
self.0.update_memory(key, ttl, |current| {
current.unwrap_or(Number::Int(0)) + value
})
}
}
#[derive(Debug, thiserror::Error)]
enum TotalError {
#[error(transparent)]
Redis(#[from] litellm_redis_client::Error),
#[error("INCRBYFLOAT answered {0:?}")]
Unexpected(Value),
}
fn total(reply: Result<Value, litellm_redis_client::Error>) -> Result<Number, TotalError> {
let reply = reply?;
let text = match &reply {
Value::BulkString(bytes) => std::str::from_utf8(bytes).ok(),
Value::SimpleString(text) => Some(text.as_str()),
_ => None,
};
text.and_then(|text| text.parse::<f64>().ok())
.map(Number::Float)
.ok_or(TotalError::Unexpected(reply))
}

294
litellm-rust/crates/redis-store/tests/router_cooldowns.rs Normal file
View file

@ -0,0 +1,294 @@
mod support;
use std::time::Duration;
use litellm_redis_client::{Expiry, Ttl};
use litellm_redis_store::{
Number,
router_cooldowns::{POLICY, RouterCooldowns},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{ALLOWED_FAILS, connect, cooldowns, within};
/// A clock fixed just after the fixture timestamps, so every fixture cooldown is still active.
fn fixture_clock() -> litellm_cache::Clock {
std::sync::Arc::new(|| Duration::from_secs_f64(1_730_000_000.75))
}
#[rstest]
#[tokio::test]
async fn put_writes_mains_str_dict_bytes_with_whole_second_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
for case in cooldowns() {
let ttl = Duration::from_secs_f64(case.ttl_seconds);
store.put(case.model_id, &case.cooldown, ttl).await.unwrap();
let key = case.key;
let written: Option<String> = raw.get(key).await.unwrap();
let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else {
panic!("{key} has no TTL");
};
assert_eq!(written.as_deref(), Some(case.str_dict));
assert!(within(remaining, case.redis_ttl_seconds), "{remaining:?}");
}
}
#[rstest]
#[tokio::test]
async fn reads_cooldowns_python_wrote_as_str_dict_or_json(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
#[values("str_dict", "json")] format: &str,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock());
for case in cooldowns() {
let bytes = match format {
"json" => case.json,
_ => case.str_dict,
};
raw.set(case.key, bytes, Duration::from_secs(60))
.await
.unwrap();
assert_eq!(
store.active(&[case.model_id.to_owned()]).await.unwrap(),
vec![(case.model_id.to_owned(), case.cooldown)]
);
}
}
#[rstest]
#[tokio::test]
async fn an_ended_cooldown_is_not_active_but_still_counts_toward_the_minimum(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let late = std::sync::Arc::new(|| Duration::from_secs(1_730_000_100));
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, late);
let models = cooldowns()
.iter()
.map(|case| case.model_id.to_owned())
.collect::<Vec<_>>();
for case in cooldowns() {
raw.set(case.key, case.str_dict, Duration::from_secs(60))
.await
.unwrap();
}
assert_eq!(
store.minimum(&models).await.unwrap(),
Some(Number::Float(0.5))
);
assert_eq!(store.active(&models).await.unwrap(), vec![]);
}
#[rstest]
#[tokio::test]
async fn allowed_fails_count_across_routers_and_keep_the_first_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let first = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
let second = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
for case in ALLOWED_FAILS {
let key = case.key;
raw.del(&[key]).await.unwrap();
let suffix = case.suffix;
let counts = [
first
.increment_allowed_fails("dep-a", suffix, Duration::from_secs_f64(30.9))
.await
.unwrap(),
second
.increment_allowed_fails("dep-a", suffix, Duration::from_secs(120))
.await
.unwrap(),
];
let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else {
panic!("{key} has no TTL");
};
assert_eq!(counts, [1, 2], "{key}");
assert!(within(remaining, 30), "{remaining:?}");
}
}
#[rstest]
#[tokio::test]
async fn python_written_fail_counts_continue_in_rust(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
raw.incr_by(
"deployment:dep-a:allowed_fails",
4,
Expiry::KeepExisting(Duration::from_secs(60)),
)
.await
.unwrap();
assert_eq!(
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
5
);
}
#[rstest]
#[tokio::test]
async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() {
let store = RouterCooldowns::with_clock(None, POLICY, fixture_clock());
let [case, _] = cooldowns();
store
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
.await
.unwrap();
let fails = [
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
];
store.flush_local_counters().unwrap();
assert_eq!(
store.active(&[case.model_id.to_owned()]).await.unwrap(),
vec![(case.model_id.to_owned(), case.cooldown)]
);
assert_eq!(fails, [1, 2]);
assert_eq!(
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
1
);
}
#[rstest]
#[tokio::test]
async fn a_redis_outage_serves_memory_and_counts_fails_locally() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock());
let [case, _] = cooldowns();
store
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
.await
.unwrap();
deployment.stop().await.unwrap();
let active = store
.active(&[case.model_id.to_owned(), "dep-missing".to_owned()])
.await
.unwrap();
let fails = store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap();
assert_eq!(active, vec![(case.model_id.to_owned(), case.cooldown)]);
assert_eq!(fails, 1);
}
struct StepClock(std::sync::Mutex<Duration>);
impl StepClock {
fn at(seconds: f64) -> std::sync::Arc<Self> {
std::sync::Arc::new(Self(std::sync::Mutex::new(Duration::from_secs_f64(
seconds,
))))
}
fn advance(&self, by: Duration) {
*self.0.lock().unwrap() += by;
}
fn clock(self: &std::sync::Arc<Self>) -> litellm_cache::Clock {
let this = self.clone();
std::sync::Arc::new(move || *this.0.lock().unwrap())
}
}
#[rstest]
#[tokio::test]
async fn a_miss_is_read_from_redis_again_only_after_the_one_second_refresh(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let time = StepClock::at(1_730_000_000.75);
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, time.clock());
let [case, _] = cooldowns();
let before = store.active(&[case.model_id.to_owned()]).await.unwrap();
raw.set(case.key, case.str_dict, Duration::from_secs(60))
.await
.unwrap();
let throttled = store.active(&[case.model_id.to_owned()]).await.unwrap();
time.advance(Duration::from_secs(1));
let refreshed = store.active(&[case.model_id.to_owned()]).await.unwrap();
assert_eq!(before, vec![]);
assert_eq!(throttled, vec![]);
assert_eq!(
refreshed,
vec![(case.model_id.to_owned(), case.cooldown.clone())]
);
}
#[rstest]
#[tokio::test]
async fn a_breaker_refusal_does_not_hold_back_the_next_read() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let breaker = litellm_circuit_breaker::Breaker::new(
litellm_circuit_breaker::Config {
failure_threshold: 1,
recovery_timeout: Duration::from_millis(200),
..litellm_circuit_breaker::Config::default()
},
Box::new(|_| {}),
);
let time = StepClock::at(1_730_000_000.75);
let guarded = support::connect_guarded(&deployment, breaker);
let store = RouterCooldowns::with_clock(Some(guarded.clone()), POLICY, time.clock());
let [case, _] = cooldowns();
deployment.stop().await.unwrap();
store.active(&[case.model_id.to_owned()]).await.unwrap();
time.advance(Duration::from_secs(1));
let refused = store.active(&[case.model_id.to_owned()]).await.unwrap();
deployment.resume().await.unwrap();
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
raw.set(case.key, case.str_dict, Duration::from_secs(60))
.await
.unwrap();
let recovered = store.active(&[case.model_id.to_owned()]).await.unwrap();
assert_eq!(refused, vec![]);
assert_eq!(
recovered,
vec![(case.model_id.to_owned(), case.cooldown.clone())]
);
}

156
litellm-rust/crates/redis-store/tests/router_usage.rs Normal file
View file

@ -0,0 +1,156 @@
mod support;
use std::time::Duration;
use litellm_redis_client::{Expiry, Ttl};
use litellm_redis_store::{
Number,
router_usage::{COUNTER_TTL, Counter, Kind, POLICY, RouterUsage},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, usage, within};
fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] {
[Kind::Tpm, Kind::Rpm].map(|kind| Counter {
deployment_id: "dep".to_owned(),
model: model.map(str::to_owned),
minute: "12-34".to_owned(),
kind,
})
}
#[rstest]
#[tokio::test]
async fn reads_counters_python_incremented_under_mains_keys(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(connect(&deployment)), POLICY);
let cases = usage();
for (case, delta) in cases.iter().zip([7.75, 3.0]) {
raw.incr_by_float(case.key, delta, Expiry::Refresh(COUNTER_TTL))
.await
.unwrap();
}
assert_eq!(
store
.get_many(&cases.map(|case| case.counter))
.await
.unwrap(),
vec![Some(Number::Float(7.75)), Some(Number::Int(3))]
);
}
#[rstest]
#[tokio::test]
async fn deployment_increments_write_incrbyfloat_totals_with_a_refreshed_minute_window(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let writer = RouterUsage::new(Some(connect(&deployment)), POLICY);
let reader = RouterUsage::new(Some(connect(&deployment)), POLICY);
let [case, _] = usage();
let model = case.counter.model.as_deref();
let [tpm_key, rpm_key] =
["tpm", "rpm"].map(|kind| case.key.replace(":tpm:", &format!(":{kind}:")));
writer
.increment_deployment("dep", model, "12-34", Number::Float(5.5), 1)
.await
.unwrap();
raw.expire(&tpm_key, Duration::from_secs(5)).await.unwrap();
writer
.increment_deployment("dep", model, "12-34", Number::Int(2), 1)
.await
.unwrap();
let written: [Option<String>; 2] = [
raw.get(&tpm_key).await.unwrap(),
raw.get(&rpm_key).await.unwrap(),
];
let Ttl::Expires(remaining) = raw.ttl(&tpm_key).await.unwrap() else {
panic!("{tpm_key} has no TTL");
};
assert_eq!(written, [Some("7.5".to_owned()), Some("2".to_owned())]);
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
reader.get_many(&tpm_and_rpm(model)).await.unwrap(),
vec![Some(Number::Float(7.5)), Some(Number::Int(2))]
);
assert_eq!(
writer.get_many(&tpm_and_rpm(model)).await.unwrap(),
vec![Some(Number::Int(7)), Some(Number::Int(2))]
);
}
#[rstest]
#[tokio::test]
async fn a_request_increment_keeps_the_window_python_started(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(connect(&deployment)), POLICY);
let [_, case] = usage();
let key = case.key;
raw.incr_by_float(key, 1.0, Expiry::KeepExisting(Duration::from_secs(20)))
.await
.unwrap();
let total = store
.increment(&case.counter, Number::Int(1), COUNTER_TTL)
.await
.unwrap();
let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else {
panic!("{key} has no TTL");
};
assert_eq!(total, Number::Float(2.0));
assert!(within(remaining, 20), "{remaining:?}");
}
#[rstest]
#[tokio::test]
async fn without_redis_counters_accumulate_in_memory_until_flushed() {
let store = RouterUsage::new(None, POLICY);
let counters = tpm_and_rpm(Some("model"));
store
.increment_deployment("dep", Some("model"), "12-34", Number::Int(2), 1)
.await
.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.5), 1)
.await
.unwrap();
let counted = store.get_many(&counters).await.unwrap();
store.flush_local_counters().unwrap();
assert_eq!(
counted,
vec![Some(Number::Float(7.5)), Some(Number::Int(2))]
);
assert_eq!(store.get_many(&counters).await.unwrap(), vec![None, None]);
}
#[rstest]
#[tokio::test]
async fn a_failed_pipeline_keeps_the_memory_counts_as_integers() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = RouterUsage::new(Some(connect(&deployment)), POLICY);
deployment.stop().await.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.75), 1)
.await
.unwrap();
assert_eq!(
store.get_many(&tpm_and_rpm(Some("model"))).await.unwrap(),
vec![Some(Number::Int(5)), Some(Number::Int(1))]
);
}

160
litellm-rust/crates/redis-store/tests/support/mod.rs Normal file
View file

@ -0,0 +1,160 @@
#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
use std::time::Duration;
use litellm_redis_client::{
ClusterEndpoint, ClusterNode, Credentials, Endpoint, Namespace, RedisClient, RedisConfig,
Timeouts, TlsOptions,
};
use litellm_redis_store::{
Number,
router_cooldowns::Cooldown,
router_usage::{Counter, Kind},
};
use litellm_redis_testing::RedisDeployment;
pub fn connect(deployment: &RedisDeployment) -> RedisClient {
connect_with(deployment, None)
}
pub fn connect_guarded(
deployment: &RedisDeployment,
breaker: litellm_circuit_breaker::Breaker,
) -> RedisClient {
connect_with(deployment, Some(breaker))
}
fn connect_with(
deployment: &RedisDeployment,
breaker: Option<litellm_circuit_breaker::Breaker>,
) -> RedisClient {
let endpoint = match deployment {
RedisDeployment::Standalone(server) => Endpoint::Url {
url: litellm_secret_value::SecretValue::new(server.url().to_owned()),
tls: TlsOptions::verified(),
},
RedisDeployment::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint {
nodes: cluster
.nodes()
.map(|(host, port)| ClusterNode {
host: host.to_owned(),
port,
})
.collect(),
tls: None,
}),
};
RedisClient::new(RedisConfig {
endpoint,
credentials: Credentials::NONE,
namespace: Namespace::default(),
timeouts: Timeouts::new(Duration::from_millis(500), None),
breaker,
events: None,
health_check: None,
})
.unwrap()
}
/// What main's `RedisCache` writes for a cooldown: `str(value)` with `EX int(ttl)`, read back with
/// `json.loads` then `ast.literal_eval` (`litellm/caching/redis_cache.py` `set_cache` and
/// `_get_cache_logic`). Keys come from main's `CooldownCache`. The same literals sit in
/// `tests/unit/redis/stores/test_router_cooldowns.py`, checked against main's codec.
pub struct CooldownCase {
pub model_id: &'static str,
pub cooldown: Cooldown,
pub ttl_seconds: f64,
pub key: &'static str,
pub redis_ttl_seconds: u64,
pub str_dict: &'static str,
pub json: &'static str,
}
pub fn cooldowns() -> [CooldownCase; 2] {
[
CooldownCase {
model_id: "dep-a",
cooldown: Cooldown {
exception_received: "RateLimitError: limited é 'quoted'".to_owned(),
status_code: "429".to_owned(),
timestamp: Number::Float(1_730_000_000.25),
cooldown_time: Number::Int(60),
},
ttl_seconds: 90.5,
key: "deployment:dep-a:cooldown",
redis_ttl_seconds: 90,
str_dict: "{'exception_received': \"RateLimitError: limited é 'quoted'\", 'status_code': '429', 'timestamp': 1730000000.25, 'cooldown_time': 60}",
json: "{\"exception_received\": \"RateLimitError: limited \\u00e9 'quoted'\", \"status_code\": \"429\", \"timestamp\": 1730000000.25, \"cooldown_time\": 60}",
},
CooldownCase {
model_id: "dep-b",
cooldown: Cooldown {
exception_received: "Timeout".to_owned(),
status_code: "408".to_owned(),
timestamp: Number::Float(1_730_000_000.5),
cooldown_time: Number::Float(0.5),
},
ttl_seconds: 30.0,
key: "deployment:dep-b:cooldown",
redis_ttl_seconds: 30,
str_dict: "{'exception_received': 'Timeout', 'status_code': '408', 'timestamp': 1730000000.5, 'cooldown_time': 0.5}",
json: "{\"exception_received\": \"Timeout\", \"status_code\": \"408\", \"timestamp\": 1730000000.5, \"cooldown_time\": 0.5}",
},
]
}
/// main's allowed-fails keys for `dep-a`: an empty suffix is the same key as none.
pub struct AllowedFailsCase {
pub suffix: Option<&'static str>,
pub key: &'static str,
}
pub const ALLOWED_FAILS: [AllowedFailsCase; 3] = [
AllowedFailsCase {
suffix: None,
key: "deployment:dep-a:allowed_fails",
},
AllowedFailsCase {
suffix: Some(""),
key: "deployment:dep-a:allowed_fails",
},
AllowedFailsCase {
suffix: Some("timeout"),
key: "deployment:dep-a:allowed_fails:timeout",
},
];
/// main's `RouterCacheEnum` usage keys, including the `None` a missing model prints into the key.
pub struct UsageCase {
pub counter: Counter,
pub key: &'static str,
}
pub fn usage() -> [UsageCase; 2] {
[
UsageCase {
counter: Counter {
deployment_id: "dep".to_owned(),
model: Some("anthropic/claude-sonnet-5-5".to_owned()),
minute: "12-34".to_owned(),
kind: Kind::Tpm,
},
key: "global_router:dep:anthropic/claude-sonnet-5-5:tpm:12-34",
},
UsageCase {
counter: Counter {
deployment_id: "dep".to_owned(),
model: None,
minute: "09-05".to_owned(),
kind: Kind::Rpm,
},
key: "global_router:dep:None:rpm:09-05",
},
]
}
/// A TTL Redis reports for a key written `whole_seconds` ago at most a few seconds back.
pub fn within(ttl: Duration, whole_seconds: u64) -> bool {
ttl <= Duration::from_secs(whole_seconds)
&& ttl > Duration::from_secs(whole_seconds.saturating_sub(5))
}

9
litellm-rust/crates/redis-testing/src/cluster.rs
View file

@ -71,6 +71,15 @@ impl RedisCluster {
self.ports.iter().map(|&port| (HOST, port))
}
/// One redis-cli reply per node, in `nodes()` order.
pub async fn cli_outputs(&self, args: &[&str]) -> Result<Vec<String>, Error> {
let mut replies = Vec::with_capacity(self.ports.len());
for port in self.ports {
replies.push(crate::exec::run(&self.container, crate::exec::cli(port, args)).await?);
}
Ok(replies)
}
pub async fn cli_on_every_node(&self, args: &[&str]) -> Result<(), Error> {
for port in self.ports {
crate::exec::run(&self.container, crate::exec::cli(port, args)).await?;

18
litellm-rust/crates/redis-testing/src/sentinel.rs
View file

@ -7,7 +7,8 @@ use testcontainers_modules::testcontainers::{
};
use crate::{
Error,
ClientConnection, Error,
deployment::parse_client_list,
exec::{cli, run},
ports::{HOST, start_on_free_ports, wait_for_pong},
standalone::REDIS_TAG,
@ -71,6 +72,21 @@ impl RedisSentinel {
(HOST, self.ports[2])
}
/// Runs one redis-cli command on the primary and the replica, which keep their own ACLs.
pub async fn cli_on_data_nodes(&self, args: &[&str]) -> Result<(), Error> {
for port in [self.ports[0], self.ports[1]] {
run(&self.container, cli(port, args)).await?;
}
Ok(())
}
pub async fn clients_on(&self, port: u16) -> Result<Vec<ClientConnection>, Error> {
Ok(parse_client_list(
port,
&run(&self.container, cli(port, &["CLIENT", "LIST"])).await?,
))
}
pub fn primary_port(&self) -> u16 {
self.ports[0]
}

7
litellm-rust/crates/redis-types/src/lib.rs
View file

@ -16,6 +16,9 @@ pub struct HostEndpoint {
pub port: u16,
pub tls: Option<TlsOptions>,
pub db: i64,
/// Connect over this unix socket instead of host, port and TLS, like redis-py's
/// `UnixDomainSocketConnection`
pub unix_socket_path: Option<PathBuf>,
}
#[derive(Clone, Debug, PartialEq)]
@ -42,6 +45,10 @@ pub struct SentinelEndpoint {
#[derive(Clone, Debug, Default, PartialEq)]
pub struct TlsOptions {
pub ca_certs: Option<PathBuf>,
/// PEM text trusted alongside `ca_certs`, redis-py's `ssl_ca_data`
pub ca_data: Option<String>,
/// A directory whose files are PEM certificates to trust, redis-py's `ssl_ca_path`
pub ca_path: Option<PathBuf>,
pub certfile: Option<PathBuf>,
pub keyfile: Option<PathBuf>,
pub verify_certificate: bool,

10
litellm/caching/redis_cache.py
View file

@ -472,8 +472,14 @@ _LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet()
def redis_breaker_state_counts() -> Mapping[str, int]:
"""How many live Redis circuit breakers sit in each state, for the Prometheus services logger."""
return MappingProxyType(dict(Counter(breaker.state for breaker in _LIVE_BREAKERS)))
"""How many live Redis circuit breakers sit in each state, for the Prometheus services logger.
Counts the Python breakers of every live ``RedisCache`` and the breaker of every native Redis
client the router stores built, since each opens and closes on its own."""
from litellm.rust_bridge._native import native_redis_breaker_states
python_states: Final = Counter(breaker.state for breaker in _LIVE_BREAKERS)
return MappingProxyType(dict(python_states + Counter(native_redis_breaker_states())))
class RedisCircuitBreakerOpenError(Exception):

14
litellm/redis/stores/_connection.py Normal file
View file

@ -0,0 +1,14 @@
from collections.abc import Mapping
from typing import Final
from pydantic import TypeAdapter
from litellm.caching.redis_cache import RedisCache
_KWARGS: Final = TypeAdapter(dict[str, object])
def redis_kwargs(connection: RedisCache) -> Mapping[str, object]:
return _KWARGS.validate_python(
connection.redis_kwargs # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped legacy dict
)

149
litellm/redis/stores/_routing_prefetch.py
View file

@ -1,149 +0,0 @@
from __future__ import annotations
import asyncio
import itertools
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm._internal_context import service_target
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_batch import BatchResult, active_request_redis_batches
_PREFETCH_SLOT: Final = "routing_read"
@dataclass(frozen=True, slots=True)
class TierRead:
cache: DualCache
keys: tuple[str, ...]
def selector_read(cache: DualCache, keys: Sequence[str]) -> TierRead:
return TierRead(cache, tuple(keys))
def routing_target(reads: Sequence[TierRead]) -> str:
if len(reads) < 2 or not reads[1].keys:
return "router_cooldowns"
return "router_cooldowns_usage" if reads[0].keys else "router_usage"
async def _backfill(read: TierRead, values: Mapping[str, object]) -> None:
pending: Final = await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # store owns tier
list(read.keys), local_only=True
)
missed: Final = {
key: values[key]
for key, local in zip(read.keys, pending.result)
if local is None and values.get(key) is not None
}
await read.cache._apply_batch_get(pending, missed) # pyright: ignore[reportPrivateUsage] # store owns tier
@dataclass(frozen=True, slots=True)
class ArmedRead:
keys: frozenset[str]
fetched: frozenset[str]
result: BatchResult[Mapping[str, object]]
reservations: tuple[tuple[TierRead, Mapping[str, float | None]], ...]
def release(self) -> None:
for read, previous in self.reservations:
read.cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # store owns tier
dict(previous)
)
async def settle(self, future: asyncio.Future[Mapping[str, object]]) -> None:
if future.cancelled() or future.exception() is not None:
self.release()
return
try:
for read, _ in self.reservations:
await _backfill(read, future.result())
except Exception:
self.release()
raise
def arm(reads: Sequence[TierRead]) -> None:
request: Final = active_request_redis_batches()
redis: Final = reads[0].cache.redis_cache if reads else None
if request is None or redis is None or _PREFETCH_SLOT in request.prefetched:
return
reserved: Final = tuple((read, read.cache.reserve_redis_batch_reads(read.keys)) for read in reads)
due_reads: Final = tuple(TierRead(read.cache, tuple(due)) for read, (due, _) in reserved)
due: Final = tuple(itertools.chain.from_iterable(read.keys for read in due_reads))
if not due:
return
with service_target(routing_target(due_reads)):
result: Final = request.batch(redis).mget(due)
prefetch: Final = ArmedRead(
keys=frozenset(itertools.chain.from_iterable(read.keys for read in reads)),
fetched=frozenset(due),
result=result,
reservations=tuple((read, reservation[1]) for read, (_, reservation) in zip(due_reads, reserved)),
)
result.on_settled(prefetch.settle)
request.prefetched[_PREFETCH_SLOT] = prefetch
def armed() -> bool:
request: Final = active_request_redis_batches()
return request is not None and _PREFETCH_SLOT in request.prefetched
def _take(needed: Sequence[str]) -> ArmedRead | None:
request: Final = active_request_redis_batches()
if request is None:
return None
taken: Final = request.prefetched.pop(_PREFETCH_SLOT, None)
if isinstance(taken, ArmedRead) and taken.keys.issuperset(needed):
return taken
if isinstance(taken, ArmedRead):
taken.release()
return None
async def _read_prefetched(reads: Sequence[TierRead]) -> Sequence[Sequence[object | None] | None] | None:
prefetch: Final = _take(tuple(itertools.chain.from_iterable(read.keys for read in reads)))
if prefetch is None:
return None
try:
values: Final = await prefetch.result
except Exception as e: # noqa: BLE001 # shared read applies the tiers' existing Redis fallback
verbose_router_logger.debug("routing prefetch failed, reading again: %s", e)
return None
pending_reads: Final = tuple(
[
await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # store owns tier
list(read.keys), local_only=True
)
for read in reads
]
)
if any(
any(key not in prefetch.fetched for key, value in zip(read.keys, pending.result) if value is None)
for read, pending in zip(reads, pending_reads)
):
return None
return [
await read.cache._apply_batch_get( # pyright: ignore[reportPrivateUsage] # store owns tier
pending, {key: values.get(key) for key, value in zip(read.keys, pending.result) if value is None}
)
for read, pending in zip(reads, pending_reads)
]
async def read(reads: Sequence[TierRead], parent_otel_span: Span | None) -> Sequence[Sequence[object | None] | None]:
with service_target(routing_target(reads)):
prefetched: Final = await _read_prefetched(reads)
if prefetched is not None:
return prefetched
return await DualCache.async_batch_get_cache_shared(
tuple((read.cache, list(read.keys)) for read in reads), parent_otel_span=parent_otel_span
)

175
litellm/redis/stores/router_cooldowns.py
View file

@ -1,9 +1,6 @@
from __future__ import annotations
import functools
import time
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
if TYPE_CHECKING:
@ -11,14 +8,10 @@ if TYPE_CHECKING:
from typing_extensions import ReadOnly, TypedDict
from litellm._internal_context import service_target
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
from litellm.rust_bridge._native import NativeRouterCooldowns
from . import _routing_prefetch
from ._routing_prefetch import TierRead
from ._connection import redis_kwargs
class CooldownCacheValue(TypedDict):
@ -28,156 +21,80 @@ class CooldownCacheValue(TypedDict):
cooldown_time: ReadOnly[float]
@dataclass(frozen=True, slots=True)
class CooldownPolicy:
redis_read_interval: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS
backfill_ttl: int = 600
max_memory_entries: int = 200
max_batch_entries: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
max_corrected_ttl: float = 60.0
POLICY: Final = CooldownPolicy()
ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns"
ActiveCooldowns: TypeAlias = Sequence[tuple[str, CooldownCacheValue]]
@dataclass(frozen=True, slots=True)
class RoutingRead:
cooldowns: ActiveCooldowns
selector_values: Sequence[object | None] | None
@functools.lru_cache(maxsize=1024)
def _cooldown_key(model_id: str) -> str:
return "deployment:" + model_id + ":cooldown"
class RouterCooldowns(Protocol):
def attach_redis(self, connection: RedisCache | None) -> None: ...
def flush_local_counters(self) -> None: ...
def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None: ...
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int: ...
def arm_routing_read(self, model_ids: Sequence[str], selector: TierRead | None) -> None: ...
async def routing_read(
self, model_ids: Sequence[str], selector: TierRead | None, parent_otel_span: Span | None
) -> RoutingRead: ...
async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ...
def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ...
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float: ...
def evict_local(self, model_id: str) -> None: ...
def snapshot(self) -> Mapping[str, CooldownCacheValue]: ...
class CooldownsHandle(Protocol):
def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: ...
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: ...
def active(self, model_ids: Sequence[str]) -> Sequence[tuple[str, CooldownCacheValue]]: ...
def aactive(self, model_ids: Sequence[str]) -> Awaitable[Sequence[tuple[str, CooldownCacheValue]]]: ...
def minimum(self, model_ids: Sequence[str]) -> float | None: ...
def evict_local(self, model_id: str) -> None: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> Mapping[str, CooldownCacheValue]: ...
HandleFactory: TypeAlias = Callable[[RedisCache | None], CooldownsHandle]
def native_handle(connection: RedisCache | None) -> CooldownsHandle:
if connection is None:
return NativeRouterCooldowns(None)
return NativeRouterCooldowns(redis_kwargs(connection), namespace=connection.namespace)
class RouterCooldownStore:
def __init__(self, connection: RedisCache | None, *, clock: Callable[[], float] | None = None) -> None:
self._clock: Final = clock if clock is not None else lambda: time.time()
self._memory: Final = InMemoryCache(
max_size_in_memory=POLICY.max_memory_entries, default_ttl=POLICY.backfill_ttl, clock=self._clock
)
self._cache: Final = DualCache(
redis_cache=connection,
in_memory_cache=self._memory,
default_redis_batch_cache_expiry=POLICY.redis_read_interval,
default_max_redis_batch_cache_size=POLICY.max_batch_entries,
clock=self._clock,
)
self._fails: Final = DualCache(redis_cache=connection, in_memory_cache=InMemoryCache(clock=self._clock))
def __init__(self, connection: RedisCache | None, *, handle: HandleFactory = native_handle) -> None:
self._connected = connection is not None
self._build: Final = handle
self._handle = handle(connection)
def attach_redis(self, connection: RedisCache | None) -> None:
self._cache.attach_redis_cache(connection)
self._fails.attach_redis_cache(connection)
if connection is None or self._connected:
return
self._connected = True
self._handle = self._build(connection)
def flush_local_counters(self) -> None:
self._fails.in_memory_cache.flush_cache()
self._handle.flush_local_counters()
def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None:
with service_target(ROUTER_COOLDOWNS_TARGET):
self._cache.set_cache(key=_cooldown_key(model_id), value=value, ttl=ttl)
self._handle.put(model_id, value, ttl)
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int:
base: Final = f"deployment:{model_id}:allowed_fails"
key: Final = f"{base}:{suffix}" if suffix else base
with service_target(ROUTER_COOLDOWNS_TARGET):
try:
return self._fails.increment_cache(key=key, value=1, ttl=ttl)
except Exception as e: # noqa: BLE001 # preserve worker-local counting during a Redis outage
verbose_router_logger.warning("allowed_fails counter fell back to this worker's in-memory count: %s", e)
local: Final[object] = self._fails.get_cache(key=key, local_only=True)
return local if isinstance(local, int) else 0
def _tier_read(self, model_ids: Sequence[str]) -> TierRead:
return TierRead(self._cache, tuple(_cooldown_key(model_id) for model_id in model_ids))
def arm_routing_read(self, model_ids: Sequence[str], selector: TierRead | None) -> None:
_routing_prefetch.arm((self._tier_read(model_ids), *(() if selector is None else (selector,))))
async def routing_read(
self, model_ids: Sequence[str], selector: TierRead | None, parent_otel_span: Span | None
) -> RoutingRead:
reads: Final = (self._tier_read(model_ids), *(() if selector is None else (selector,)))
results: Final = await _routing_prefetch.read(reads, parent_otel_span)
return RoutingRead(
cooldowns=self._active_from_results(model_ids, results[0]),
selector_values=None if selector is None else results[1],
)
def _corrected_active_cooldown(
self, model_id: str, result: Mapping[str, object], current_time: float
) -> CooldownCacheValue | None:
value: Final = CooldownCacheValue(
**result # pyright: ignore[reportArgumentType] # preserve main's unvalidated dict codec
)
remaining: Final = value["timestamp"] + value["cooldown_time"] - current_time
key: Final = _cooldown_key(model_id)
if remaining <= 0:
self._memory.delete_cache(key)
return None
expiry: Final = self._memory.ttl_dict.get(key)
if expiry is not None and expiry > current_time + remaining + 5:
self._memory.delete_cache(key)
self._memory.set_cache(key, result, ttl=min(remaining, POLICY.max_corrected_ttl))
return value
def _active_from_results(self, model_ids: Sequence[str], results: Sequence[object] | None) -> ActiveCooldowns:
if results is None:
return ()
now: Final = self._clock()
return tuple(
(model_id, active)
for model_id, result in zip(model_ids, results)
if isinstance(result, dict)
and result
and (active := self._corrected_active_cooldown(model_id, result, now)) is not None
)
return self._handle.increment_allowed_fails(model_id, suffix, ttl)
async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns:
with service_target(ROUTER_COOLDOWNS_TARGET):
results: Final = await self._cache.async_batch_get_cache(
keys=[_cooldown_key(model_id) for model_id in model_ids], parent_otel_span=parent_otel_span
)
return self._active_from_results(model_ids, results)
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float:
with service_target(ROUTER_COOLDOWNS_TARGET):
results: Final = (
self._cache.batch_get_cache(
keys=[_cooldown_key(model_id) for model_id in model_ids], parent_otel_span=parent_otel_span
)
or ()
)
durations: Final = tuple(
CooldownCacheValue(**result)["cooldown_time"] # pyright: ignore[reportArgumentType] # preserve main's dict codec
for result in results
if isinstance(result, dict) and result
)
return (min(durations) if durations else None) or default
pending: Final = self._handle.aactive(model_ids)
return tuple(await pending)
def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns:
with service_target(ROUTER_COOLDOWNS_TARGET):
results: Final = self._cache.batch_get_cache(
keys=[_cooldown_key(model_id) for model_id in model_ids], parent_otel_span=parent_otel_span
)
return self._active_from_results(model_ids, results)
return tuple(self._handle.active(model_ids))
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float:
with service_target(ROUTER_COOLDOWNS_TARGET):
return self._handle.minimum(model_ids) or default
def evict_local(self, model_id: str) -> None:
self._memory.delete_cache(_cooldown_key(model_id))
self._handle.evict_local(model_id)
def snapshot(self) -> Mapping[str, CooldownCacheValue]:
return self._handle.snapshot()

98
litellm/redis/stores/router_usage.py
View file

@ -1,6 +1,6 @@
from __future__ import annotations
from collections.abc import Sequence
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias
@ -8,27 +8,11 @@ if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm._internal_context import with_service_target
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.rust_bridge._native import NativeRouterUsage
from ._keys import ROUTER_RPM_KEY, ROUTER_TPM_KEY
from ._connection import redis_kwargs
@dataclass(frozen=True, slots=True)
class UsagePolicy:
redis_read_interval: float | None = None
memory_ttl: float | None = None
redis_ttl: float | None = None
backfill_ttl: int = 600
max_memory_entries: int = 200
max_batch_entries: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
counter_ttl: int = 60
POLICY: Final = UsagePolicy()
ROUTER_USAGE_TARGET: Final = "router_usage"
CounterKind: TypeAlias = Literal["tpm", "rpm"]
@ -41,11 +25,6 @@ class UsageCounter:
kind: CounterKind
def _key(counter: UsageCounter) -> str:
template: Final = ROUTER_TPM_KEY if counter.kind == "tpm" else ROUTER_RPM_KEY
return template.format(id=counter.deployment_id, model=counter.model, current_minute=counter.minute)
class RouterUsage(Protocol):
def attach_redis(self, connection: RedisCache | None) -> None: ...
def flush_local_counters(self) -> None: ...
@ -62,34 +41,55 @@ class RouterUsage(Protocol):
parent_otel_span: Span | None,
) -> None: ...
async def async_batch_get(self, counters: Sequence[UsageCounter]) -> Sequence[object] | None: ...
def snapshot(self) -> Mapping[str, float]: ...
NativeCounter: TypeAlias = tuple[str, str | None, str, CounterKind]
class UsageHandle(Protocol):
def increment(self, counter: NativeCounter, value: float, ttl_seconds: int) -> Awaitable[float]: ...
def increment_deployment(
self, deployment_id: str, model: str | None, minute: str, tokens: float, requests: int
) -> Awaitable[None]: ...
def get_many(self, counters: Sequence[NativeCounter]) -> Awaitable[Sequence[float | None]]: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> Mapping[str, float]: ...
HandleFactory: TypeAlias = Callable[[RedisCache | None], UsageHandle]
def native_handle(connection: RedisCache | None) -> UsageHandle:
if connection is None:
return NativeRouterUsage(None)
return NativeRouterUsage(redis_kwargs(connection), namespace=connection.namespace)
def _native(counter: UsageCounter) -> NativeCounter:
return (counter.deployment_id, counter.model, counter.minute, counter.kind)
class RouterUsageStore:
def __init__(self, connection: RedisCache | None) -> None:
self._cache: Final = DualCache(
redis_cache=connection,
in_memory_cache=InMemoryCache(
max_size_in_memory=POLICY.max_memory_entries, default_ttl=POLICY.backfill_ttl
),
default_in_memory_ttl=POLICY.memory_ttl,
default_redis_ttl=POLICY.redis_ttl,
default_redis_batch_cache_expiry=POLICY.redis_read_interval,
default_max_redis_batch_cache_size=POLICY.max_batch_entries,
)
def __init__(self, connection: RedisCache | None, *, handle: HandleFactory = native_handle) -> None:
self._connected = connection is not None
self._build: Final = handle
self._handle = handle(connection)
def attach_redis(self, connection: RedisCache | None) -> None:
self._cache.attach_redis_cache(connection)
if connection is None or self._connected:
return
self._connected = True
self._handle = self._build(connection)
def flush_local_counters(self) -> None:
self._cache.in_memory_cache.flush_cache()
self._handle.flush_local_counters()
@with_service_target(ROUTER_USAGE_TARGET)
async def async_increment(
self, counter: UsageCounter, value: float, ttl: int, parent_otel_span: Span | None = None
) -> float | None:
return await self._cache.async_increment_cache(
key=_key(counter), value=value, ttl=ttl, parent_otel_span=parent_otel_span
)
return await self._handle.increment(_native(counter), value, ttl)
@with_service_target(ROUTER_USAGE_TARGET)
async def increment_deployment(
@ -101,18 +101,10 @@ class RouterUsageStore:
requests: int,
parent_otel_span: Span | None,
) -> None:
tpm: Final = UsageCounter(deployment_id, model, minute, "tpm")
rpm: Final = UsageCounter(deployment_id, model, minute, "rpm")
operations: Final = [
RedisPipelineIncrementOperation(key=_key(counter), increment_value=value, ttl=POLICY.counter_ttl)
for counter, value in ((tpm, tokens), (rpm, requests))
]
values: Final = await self._cache.async_increment_cache_pipeline(
increment_list=operations, parent_otel_span=parent_otel_span
)
if values is not None and self._cache.redis_cache is not None:
for operation, value in zip(operations, values):
await self._cache.async_set_cache(operation["key"], int(value), local_only=True, ttl=POLICY.counter_ttl)
await self._handle.increment_deployment(deployment_id, model, minute, tokens, requests)
async def async_batch_get(self, counters: Sequence[UsageCounter]) -> Sequence[object] | None:
return await self._cache.async_batch_get_cache(keys=[_key(counter) for counter in counters])
return await self._handle.get_many(tuple(map(_native, counters)))
def snapshot(self) -> Mapping[str, float]:
return self._handle.snapshot()

121
litellm/router_utils/routing_prefetch.py Normal file
View file

@ -0,0 +1,121 @@
from __future__ import annotations
import asyncio
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm._internal_context import service_target
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_batch import BatchResult, active_request_redis_batches
_PREFETCH_SLOT: Final = "routing_read"
ROUTER_USAGE_TARGET: Final = "router_usage"
@dataclass(frozen=True, slots=True)
class TierRead:
cache: DualCache
keys: tuple[str, ...]
def selector_read(cache: DualCache, keys: Sequence[str]) -> TierRead:
return TierRead(cache, tuple(keys))
@dataclass(frozen=True, slots=True)
class ArmedRead:
read: TierRead
fetched: frozenset[str]
result: BatchResult[Mapping[str, object]]
previous: Mapping[str, float | None]
def release(self) -> None:
self.read.cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # the selector owns this tier
dict(self.previous)
)
async def settle(self, future: asyncio.Future[Mapping[str, object]]) -> None:
if future.cancelled() or future.exception() is not None:
self.release()
return
try:
pending: Final = await self.read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier
list(self.read.keys), local_only=True
)
values: Final = future.result()
missed: Final = {
key: values[key]
for key, local in zip(self.read.keys, pending.result)
if local is None and values.get(key) is not None
}
await self.read.cache._apply_batch_get(pending, missed) # pyright: ignore[reportPrivateUsage] # the selector owns this tier
except Exception:
self.release()
raise
def arm(read: TierRead) -> None:
request: Final = active_request_redis_batches()
redis: Final = read.cache.redis_cache
if request is None or redis is None or _PREFETCH_SLOT in request.prefetched:
return
due, previous = read.cache.reserve_redis_batch_reads(read.keys)
if not due:
return
with service_target(ROUTER_USAGE_TARGET):
result: Final = request.batch(redis).mget(tuple(due))
armed_read: Final = ArmedRead(read=read, fetched=frozenset(due), result=result, previous=previous)
result.on_settled(armed_read.settle)
request.prefetched[_PREFETCH_SLOT] = armed_read
def armed() -> bool:
request: Final = active_request_redis_batches()
return request is not None and _PREFETCH_SLOT in request.prefetched
def _take(read: TierRead) -> ArmedRead | None:
request: Final = active_request_redis_batches()
if request is None:
return None
taken: Final = request.prefetched.pop(_PREFETCH_SLOT, None)
if isinstance(taken, ArmedRead) and set(taken.read.keys).issuperset(read.keys):
return taken
if isinstance(taken, ArmedRead):
taken.release()
return None
async def _read_prefetched(read: TierRead) -> Sequence[object | None] | None:
prefetch: Final = _take(read)
if prefetch is None:
return None
try:
values: Final = await prefetch.result
except Exception as e: # noqa: BLE001 # the read falls back to the tier's own Redis read
verbose_router_logger.debug("routing prefetch failed, reading again: %s", e)
return None
pending: Final = await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier
list(read.keys), local_only=True
)
if any(key not in prefetch.fetched for key, value in zip(read.keys, pending.result) if value is None):
return None
return await read.cache._apply_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier
pending, {key: values.get(key) for key, value in zip(read.keys, pending.result) if value is None}
)
async def read(read: TierRead, parent_otel_span: Span | None) -> Sequence[object | None] | None:
with service_target(ROUTER_USAGE_TARGET):
prefetched: Final = await _read_prefetched(read)
if prefetched is not None:
return prefetched
results: Final = await DualCache.async_batch_get_cache_shared(
((read.cache, list(read.keys)),), parent_otel_span=parent_otel_span
)
return results[0]

33
litellm/router_utils/routing_read_batch.py
View file

@ -1,21 +1,21 @@
from collections.abc import Iterator
import asyncio
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from contextvars import ContextVar
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_router_logger
from litellm.redis.stores import _routing_prefetch
from litellm.redis.stores._routing_prefetch import TierRead, selector_read
from litellm.router_deprecated import renamed_router_param
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage
from litellm.router_utils import routing_prefetch
from litellm.router_utils.routing_prefetch import ROUTER_USAGE_TARGET, TierRead, selector_read
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm.router_api import RoutingContext
ROUTER_USAGE_TARGET: Final = "router_usage"
__all__ = ("ROUTER_USAGE_TARGET", "RoutingPrefetch", "RoutingReadBatch")
def _selector_keys(selector: LowestTPMLoggingHandler_v2, deployments: list) -> list[str]:
@ -27,17 +27,21 @@ def _selector_read(selector: LowestTPMLoggingHandler_v2 | None, deployments: lis
return None if selector is None else selector_read(selector.router_cache, _selector_keys(selector, deployments))
async def _nothing() -> Sequence[object | None] | None:
return None
class RoutingPrefetch:
@staticmethod
@renamed_router_param("litellm_router_instance")
def arm(ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: list) -> None:
ctx.health().cooldowns.cooldown_store.arm_routing_read(
ctx.registry().names(kind="ids"), _selector_read(usage_selector, deployments)
)
read: Final = _selector_read(usage_selector, deployments)
if read is not None:
routing_prefetch.arm(read)
@staticmethod
def armed() -> bool:
return _routing_prefetch.armed()
return routing_prefetch.armed()
_active_routing_read_batch: Final[ContextVar["RoutingReadBatch | None"]] = ContextVar(
@ -74,13 +78,14 @@ class RoutingReadBatch:
self, ctx: "RoutingContext", healthy_deployments: list, parent_otel_span: "Span | None"
) -> list[str]:
usage_read: Final = _selector_read(self.usage_selector, healthy_deployments)
read: Final = await ctx.health().cooldowns.cooldown_store.routing_read(
ctx.registry().names(kind="ids"), usage_read, parent_otel_span
cooldowns, selector_values = await asyncio.gather(
ctx.health().cooldowns.cooldown_store.active(ctx.registry().names(kind="ids"), parent_otel_span),
_nothing() if usage_read is None else routing_prefetch.read(usage_read, parent_otel_span),
)
if usage_read is not None:
self.prefetched_usage = PrefetchedUsage(
keys=frozenset(usage_read.keys),
values=None if read.selector_values is None else dict(zip(usage_read.keys, read.selector_values)),
values=None if selector_values is None else dict(zip(usage_read.keys, selector_values)),
)
verbose_router_logger.debug("retrieve cooldown models: %s", read.cooldowns)
return [model_id for model_id, _ in read.cooldowns]
verbose_router_logger.debug("retrieve cooldown models: %s", cooldowns)
return [model_id for model_id, _ in cooldowns]

43
litellm/rust_bridge/_native.pyi
View file

@ -1,6 +1,6 @@
from asyncio import Future
from collections.abc import Mapping, Sequence
from typing import Never, final
from typing import Literal, Never, TypedDict, final
import httpx
from pydantic import JsonValue
@ -15,6 +15,7 @@ class ProcessReservedForForking(RuntimeError): ...
def trace_encode_error(message: str) -> bytes: ...
def bundled_model_cost_map() -> bytes: ...
def native_redis_breaker_states() -> list[str]: ...
@final
class NativeClickHouseSpendConfig:
@ -26,6 +27,43 @@ class NativeClickHouseSpendStorage:
def ensure_schema(self) -> Future[None]: ...
def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ...
class _NativeCooldown(TypedDict):
exception_received: str
status_code: str
timestamp: float
cooldown_time: float
@final
class NativeRouterCooldowns:
def __new__(
cls, redis_kwargs: Mapping[str, object] | None, *, namespace: str | None = None
) -> NativeRouterCooldowns: ...
def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: ...
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: ...
def active(self, model_ids: Sequence[str]) -> list[tuple[str, _NativeCooldown]]: ...
def aactive(self, model_ids: Sequence[str]) -> Future[list[tuple[str, _NativeCooldown]]]: ...
def minimum(self, model_ids: Sequence[str]) -> float | None: ...
def evict_local(self, model_id: str) -> None: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> dict[str, _NativeCooldown]: ...
@final
class NativeRouterUsage:
def __new__(
cls, redis_kwargs: Mapping[str, object] | None, *, namespace: str | None = None
) -> NativeRouterUsage: ...
def increment(
self, counter: tuple[str, str | None, str, Literal["tpm", "rpm"]], value: float, ttl_seconds: int
) -> Future[float]: ...
def increment_deployment(
self, deployment_id: str, model: str | None, minute: str, tokens: float, requests: int
) -> Future[None]: ...
def get_many(
self, counters: Sequence[tuple[str, str | None, str, Literal["tpm", "rpm"]]]
) -> Future[list[float | None]]: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> dict[str, float]: ...
@final
class NativeDatabase:
@staticmethod
@ -216,6 +254,8 @@ __all__ = [
"NativeClickHouseSpendStorage",
"NativeDatabase",
"NativeDiagnosticProcessor",
"NativeRouterCooldowns",
"NativeRouterUsage",
"ProcessReservedForForking",
"ResponsesWebSocketConnection",
"RustBridgeDeclined",
@ -233,6 +273,7 @@ __all__ = [
"embedding",
"gil_stats",
"messages",
"native_redis_breaker_states",
"ocr",
"ocr_health_check_document",
"ocr_passthrough_response",

2
litellm/rust_bridge/host/AGENTS.md
View file

@ -2,7 +2,7 @@
Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds)
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`)
- A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here
- Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers
- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green

46
litellm/rust_bridge/host/redis_events.py Normal file
View file

@ -0,0 +1,46 @@
"""Redis call and circuit breaker events from the native Redis client, reported the way ``RedisCache`` reports them."""
from __future__ import annotations
from functools import lru_cache
from typing import Final, Literal, TypeAlias
from litellm._service_logger import ServiceLogging
from litellm.types.services import ServiceTypes
Failure: TypeAlias = Literal["timeout", "transport", "server", "decode", "breaker_open"]
BreakerState: TypeAlias = Literal["closed", "open", "half_open"]
_HEALTH_FAILURE_CLASSES: Final = {"timeout": "timeout", "transport": "connectivity"}
class RedisCallError(Exception):
"""The failure class of a native Redis call, the only detail the Rust client hands over."""
@lru_cache(maxsize=1)
def _service_logging() -> ServiceLogging:
return ServiceLogging()
def record_redis_call(call_type: str, duration: float, failure: Failure | None, caller: str | None) -> None:
if failure == "breaker_open":
return
logger: Final = _service_logging()
if failure is None:
logger.service_success_hook(service=ServiceTypes.REDIS, duration=duration, call_type=call_type, caller=caller)
return
logger.service_failure_hook(
service=ServiceTypes.REDIS,
duration=duration,
error=RedisCallError(failure),
call_type=call_type,
caller=caller,
)
failure_class: Final = _HEALTH_FAILURE_CLASSES.get(failure)
if failure_class is not None:
logger.redis_breaker_hook(transitions_to=None, failure_class=failure_class)
def record_redis_breaker(state: BreakerState) -> None:
_service_logging().redis_breaker_hook(transitions_to=state, failure_class=None)

20
litellm/rust_bridge/host/settings.py
View file

@ -31,6 +31,15 @@ class ProviderDefaults:
enable_azure_ad_token_refresh: object
@dataclass(frozen=True, slots=True)
class RouterTierSettings:
"""The memory-over-Redis tier values `main`'s router reads from the environment and `litellm` globals."""
cooldown_redis_read_interval: float
usage_redis_read_interval: float
max_redis_batch_cache_size: int
@dataclass(frozen=True, slots=True)
class SecretManager:
readable: bool
@ -108,6 +117,17 @@ def secret_manager_binding() -> SecretManagerBinding:
)
def router_tiers() -> RouterTierSettings:
import litellm
from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
return RouterTierSettings(
cooldown_redis_read_interval=DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS,
usage_redis_read_interval=litellm.default_redis_batch_cache_expiry or 10,
max_redis_batch_cache_size=DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE,
)
def provider_defaults() -> ProviderDefaults:
import litellm

1
tests/unit/caching/test_redis_batch.py
View file

@ -93,6 +93,7 @@ class FakeRedisCache(RedisCache):
def __init__(self, client: FakeClient, namespace: str | None = None) -> None: # super().__init__ needs a server
self.client = client
self.namespace = namespace
self.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2, "db": id(self) % 1_000_000}
self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30)
self.service_logger_obj = ServiceLogging()
self.default_ttl = None

632
tests/unit/caching/test_request_redis_batch_pre_call.py
View file

@ -1,13 +1,13 @@
"""One Redis pipeline per backend for the pre-call reads a request makes: rate limiter Lua groups, the
router's cooldown and usage read, auth identity and spend counters all join the request batch."""
router's usage counter read, auth identity and spend counters all join the request batch."""
from __future__ import annotations
import asyncio
import hashlib
import json
from collections.abc import Callable
from datetime import datetime, timezone
from itertools import chain
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock
@ -24,15 +24,13 @@ from litellm.proxy.auth.auth_object_prefetch import _CacheEntry, _write_back, pr
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
CHECK_AND_INCREMENT_BY_N_SCRIPT,
PROXY_MaxParallelRequestsHandler_v3,
RateLimitDescriptor,
RateLimitUnverifiableError,
PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.router_api import RoutingContext
from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET
from litellm.redis.stores.router_cooldowns import POLICY
from litellm.redis.stores._routing_prefetch import ArmedRead, TierRead, routing_target
from litellm.router_utils.routing_prefetch import ArmedRead, TierRead
from litellm.router_utils.routing_read_batch import (
ROUTER_USAGE_TARGET,
RoutingPrefetch,
@ -294,12 +292,76 @@ def _router(redis_cache: FakeRedisCache, routing_strategy: str = "usage-based-ro
return router
@pytest.mark.asyncio
async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issues_no_read_of_its_own(monkeypatch):
monkeypatch.setattr(
"litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime",
lambda: datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc),
_PINNED_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
_START: Final = 1_000_000.0
_PING: Final = [{"role": "user", "content": "ping"}]
@pytest.fixture
def pinned_clock(monkeypatch: pytest.MonkeyPatch) -> list[float]:
"""Pins the usage counters' minute and the router cache's read throttle clock; tests move `now[0]`."""
now: Final = [_START] # mutable-ok: the test advances the pinned clock
monkeypatch.setattr("litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", lambda: _PINNED_MINUTE)
monkeypatch.setattr(dual_cache_module.time, "time", lambda: now[0])
return now
def _counter_keys(router: Router, deployment_ids: tuple[str, ...] = ("dep-a", "dep-b")) -> frozenset[str]:
deployments: Final = [d for d in router.model_list if d["model_info"]["id"] in deployment_ids]
tpm, rpm = router.lowesttpm_logger_v2.usage_counter_keys(deployments)
return frozenset(tpm + rpm)
def _tpm_key(router: Router, deployment_id: str) -> str:
deployment: Final = next(d for d in router.model_list if d["model_info"]["id"] == deployment_id)
return router.lowesttpm_logger_v2.usage_counter_keys([deployment])[0][0]
def _counter_mgets(router: Router, client: FakeClient, redis_cache: FakeRedisCache) -> tuple[frozenset[str], ...]:
counters: Final = _counter_keys(router)
pipelined: Final = tuple(
frozenset(command[1:])
for pipeline in client.pipelines
for command in pipeline.commands
if command[0] == "MGET" and counters.intersection(command[1:])
)
alone: Final = tuple(
frozenset(keys) for command, keys in redis_cache.alone if command == "MGET" and counters.intersection(keys)
)
return pipelined + alone
def _python_cooldown_reads(client: FakeClient, redis_cache: FakeRedisCache) -> tuple[str, ...]:
pipelined: Final = tuple(
key
for pipeline in client.pipelines
for command in pipeline.commands
if command[0] == "MGET"
for key in command[1:]
)
alone: Final = tuple(key for command, keys in redis_cache.alone if command == "MGET" for key in keys)
return tuple(key for key in pipelined + alone if key.endswith(":cooldown"))
def _replies_with(values: dict[str, object]) -> Callable[[tuple[Any, ...]], Any]:
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "MGET":
return [None if values.get(key) is None else json.dumps(values[key]) for key in command[1:]]
return _lua_ok_replies(command)
return replies
async def _route(router: Router) -> str:
deployment: Final = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=_PING, request_kwargs={}
)
return deployment["model_info"]["id"]
@pytest.mark.asyncio
@pytest.mark.usefixtures("pinned_clock")
async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issues_no_read_of_its_own():
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache)
@ -311,55 +373,40 @@ async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issue
descriptors=[_descriptor("api_key", "k1", 10), _descriptor("team", "t1", 20)],
increments=[{"requests": 1}, {"requests": 1}],
)
deployment = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
deployment = await _route(router)
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
assert deployment in {"dep-a", "dep-b"}
assert len(client.pipelines) == 1
commands = client.pipelines[0].commands
assert [c[0] for c in commands] == ["MGET", "EVALSHA", "EVALSHA"]
mget_keys = set(commands[0][1:])
assert {_cooldown_key("dep-a"), _cooldown_key("dep-b")} <= mget_keys
assert any(":tpm:" in key for key in mget_keys) and any(":rpm:" in key for key in mget_keys)
assert frozenset(commands[0][1:]) == _counter_keys(router)
assert redis_cache.alone == []
@pytest.mark.asyncio
async def test_a_cooldown_recorded_locally_after_the_prefetch_left_still_excludes_its_deployment():
expired = {"exception_received": "429", "status_code": "429", "timestamp": 0.0, "cooldown_time": 60}
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "MGET": # Redis holds a stale cooldown for dep-b and nothing for dep-a
return [json.dumps(expired) if key == _cooldown_key("dep-b") else None for key in command[1:]]
return _lua_ok_replies(command)
client = FakeClient(replies)
@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing-v2"])
async def test_a_cooldown_recorded_through_the_store_excludes_its_deployment_without_a_python_cooldown_read(
routing_strategy: str,
):
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache)
cooldown_store = router.cooldown_cache.cooldown_store
router = _router(redis_cache, routing_strategy=routing_strategy)
router.cooldown_cache.cooldown_store.put(
"dep-a",
{"exception_received": "429", "status_code": "429", "timestamp": _FAR_FUTURE, "cooldown_time": 60},
60,
)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
cooldown_store.put(
"dep-a",
{"exception_received": "429", "status_code": "429", "timestamp": _FAR_FUTURE, "cooldown_time": 60},
60,
)
picks = {
(
await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
)["model_info"]["id"]
for _ in range(5)
}
picks = {await _route(router) for _ in range(5)}
assert picks == {"dep-b"}
assert len(client.pipelines) == 1
assert _python_cooldown_reads(client, redis_cache) == ()
@pytest.mark.asyncio
@pytest.mark.usefixtures("pinned_clock")
async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_routing_reads_itself():
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
@ -370,79 +417,56 @@ async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_ro
armed = request.prefetched["routing_read"]
assert isinstance(armed, ArmedRead)
request.prefetched["routing_read"] = ArmedRead(
keys=frozenset({"other"}),
read=TierRead(armed.read.cache, ("other",)),
fetched=armed.fetched,
result=armed.result,
reservations=armed.reservations,
)
deployment = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
previous=armed.previous,
)
deployment = await _route(router)
assert request.prefetched == {}
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
assert len(redis_cache.alone) == 1 # the shared cooldown+usage read, one round trip as in P1
assert deployment in {"dep-a", "dep-b"}
assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)]
@pytest.mark.asyncio
async def test_a_prefetch_with_incomplete_usage_keys_releases_cooldown_reservations():
@pytest.mark.usefixtures("pinned_clock")
async def test_a_prefetch_with_incomplete_usage_keys_releases_its_reservations():
client: Final = FakeClient(_lua_ok_replies)
redis_cache: Final = FakeRedisCache(client)
router: Final = _router(redis_cache)
cooldown_keys: Final = frozenset(
{
_cooldown_key("dep-a"),
_cooldown_key("dep-b"),
}
)
with request_redis_batch_scope():
RoutingPrefetch.arm(RoutingContext(router), router.lowesttpm_logger_v2, router.model_list[:1])
deployment: Final = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
deployment: Final = await _route(router)
fallback_cooldown_mgets: Final = tuple(
keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_keys.issubset(keys)
)
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
assert len(fallback_cooldown_mgets) == 1
assert deployment in {"dep-a", "dep-b"}
assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)]
@pytest.mark.asyncio
async def test_a_failed_prefetch_falls_back_to_the_shared_read():
@pytest.mark.usefixtures("pinned_clock")
async def test_a_failed_prefetch_falls_back_to_the_selectors_own_read():
client = FakeClient(_lua_ok_replies, fail=ConnectionError("redis down"))
redis_cache = FakeRedisCache(client)
router = _router(redis_cache)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
deployment = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
deployment = await _route(router)
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
assert len(redis_cache.alone) == 1
cooldown_keys: Final = frozenset(
{
_cooldown_key("dep-a"),
_cooldown_key("dep-b"),
}
)
fallback_cooldown_mgets: Final = tuple(
keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_keys.issubset(keys)
)
assert len(fallback_cooldown_mgets) == 1
assert deployment in {"dep-a", "dep-b"}
assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)]
@pytest.mark.asyncio
async def test_the_armed_routing_read_is_declared_under_the_router_cooldowns_family():
@pytest.mark.usefixtures("pinned_clock")
async def test_the_armed_routing_read_is_declared_under_the_router_usage_family():
"""The prefetch is declared before routing runs under a target of its own, so the pipeline that
carries it renders ``redis.pipeline router_cooldowns`` instead of a bare ``redis.pipeline``."""
carries it renders ``redis.pipeline router_usage`` instead of a bare ``redis.pipeline``."""
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache, routing_strategy="simple-shuffle")
router = _router(redis_cache)
pipeline_targets: list[str | None] = [] # mutable-ok: filled by the recording hook
async def record(**kwargs: object) -> None:
@ -452,35 +476,21 @@ async def test_the_armed_routing_read_is_declared_under_the_router_cooldowns_fam
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task()))
await _route(router)
async with asyncio.timeout(5):
while not pipeline_targets:
await asyncio.sleep(0.01)
assert pipeline_targets == [ROUTER_COOLDOWNS_TARGET]
assert pipeline_targets == [ROUTER_USAGE_TARGET]
@pytest.mark.asyncio
async def test_an_abandoned_prefetch_still_backfills_the_cooldown_it_read(monkeypatch):
clock: Final = 1_000_000.0
monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
cooldown_key: Final = _cooldown_key("dep-a")
active_cooldown: Final = {
"exception_received": "429",
"status_code": "429",
"timestamp": _FAR_FUTURE,
"cooldown_time": 60,
}
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "MGET":
return [json.dumps(active_cooldown) if key == cooldown_key else None for key in command[1:]]
return _lua_ok_replies(command)
client: Final = FakeClient(replies)
@pytest.mark.usefixtures("pinned_clock")
async def test_an_abandoned_prefetch_still_backfills_the_counters_it_read():
client: Final = FakeClient(_lua_ok_replies)
redis_cache: Final = FakeRedisCache(client)
redis_cache.store[cooldown_key] = active_cooldown
router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
router: Final = _router(redis_cache)
client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100})
limiter: Final = _limiter(redis_cache)
with request_redis_batch_scope():
@ -489,176 +499,66 @@ async def test_an_abandoned_prefetch_still_backfills_the_cooldown_it_read(monkey
descriptors=[_descriptor("api_key", "k1", 10)],
increments=[{"requests": 1}],
)
after_abandoned: Final = len(_counter_mgets(router, client, redis_cache))
pipeline_count: Final = len(client.pipelines)
first_cooldown_mgets: Final = tuple(
command
for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[:pipeline_count])
if command[0] == "MGET" and cooldown_key in command[1:]
)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
deployment: Final = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
cooldowns: Final = await router.cooldown_cache.async_get_active_cooldowns(["dep-a"], parent_otel_span=None)
deployment: Final = await _route(router)
second_cooldown_mgets: Final = tuple(
command
for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[pipeline_count:])
if command[0] == "MGET" and cooldown_key in command[1:]
)
assert deployment["model_info"]["id"] == "dep-b"
assert [model_id for model_id, _ in cooldowns] == ["dep-a"]
assert len(first_cooldown_mgets) == 1
assert second_cooldown_mgets == ()
assert deployment == "dep-b"
assert after_abandoned == 1
assert len(_counter_mgets(router, client, redis_cache)) == 1
assert redis_cache.alone == []
@pytest.mark.asyncio
async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misses(monkeypatch):
clock: Final = 1_000_000.0
monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
dep_a_key: Final = _cooldown_key("dep-a")
dep_b_key: Final = _cooldown_key("dep-b")
old_cooldown: Final = {
"exception_received": "429",
"status_code": "429",
"timestamp": _FAR_FUTURE,
"cooldown_time": 60,
}
newer_memory_cooldown: Final = {
"exception_received": "429",
"status_code": "429",
"timestamp": _FAR_FUTURE + 1,
"cooldown_time": 60,
}
redis_only_cooldown: Final = {
"exception_received": "429",
"status_code": "429",
"timestamp": _FAR_FUTURE + 2,
"cooldown_time": 60,
}
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "MGET":
return [json.dumps(redis_cache.store[key]) if key in redis_cache.store else None for key in command[1:]]
return _lua_ok_replies(command)
client: Final = FakeClient(replies)
redis_cache: Final = FakeRedisCache(client)
redis_cache.store[dep_a_key] = old_cooldown
redis_cache.store[dep_b_key] = redis_only_cooldown
router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
cooldown_store: Final = router.cooldown_cache.cooldown_store
with request_redis_batch_scope() as request:
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
cooldown_store.put("dep-a", newer_memory_cooldown, 60)
redis_cache.store[dep_a_key] = old_cooldown
await request.flush_all()
prefetched_mgets: Final = tuple(command for command in client.pipelines[0].commands if command[0] == "MGET")
assert len(prefetched_mgets) == 1
assert frozenset(prefetched_mgets[0][1:]) == frozenset({dep_a_key, dep_b_key})
assert dict(await cooldown_store.active(("dep-a", "dep-b"), None)) == {
"dep-a": newer_memory_cooldown,
"dep-b": redis_only_cooldown,
}
@pytest.mark.asyncio
async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(monkeypatch):
clock: Final = 1_000_000.0
monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
cooldown_key: Final = _cooldown_key("dep-a")
active_cooldown: Final = {
"exception_received": "429",
"status_code": "429",
"timestamp": _FAR_FUTURE,
"cooldown_time": 60,
}
mget_replies: Final = iter((ConnectionError("redis down"), None))
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "MGET":
response: Final = next(mget_replies)
if isinstance(response, Exception):
return response
return [json.dumps(active_cooldown) if key == cooldown_key else None for key in command[1:]]
return _lua_ok_replies(command)
client: Final = FakeClient(replies)
redis_cache: Final = FakeRedisCache(client)
redis_cache.store[cooldown_key] = active_cooldown
router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
with request_redis_batch_scope() as request:
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await request.flush_all()
pipeline_count: Final = len(client.pipelines)
first_cooldown_mgets: Final = tuple(
command
for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[:pipeline_count])
if command[0] == "MGET" and cooldown_key in command[1:]
)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
deployment: Final = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
cooldowns: Final = await router.cooldown_cache.async_get_active_cooldowns(["dep-a"], parent_otel_span=None)
second_cooldown_mgets: Final = tuple(
command
for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[pipeline_count:])
if command[0] == "MGET" and cooldown_key in command[1:]
)
assert deployment["model_info"]["id"] == "dep-b"
assert [model_id for model_id, _ in cooldowns] == ["dep-a"]
assert len(first_cooldown_mgets) == 1
assert len(second_cooldown_mgets) == 1
@pytest.mark.asyncio
async def test_a_cooldown_that_leaves_memory_before_routing_is_read_again(monkeypatch):
clock: Final = 1_000_000.0
monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
cooldown_key: Final = _cooldown_key("dep-a")
active_cooldown: Final = {
"exception_received": "429",
"status_code": "429",
"timestamp": _FAR_FUTURE,
"cooldown_time": 60,
}
@pytest.mark.usefixtures("pinned_clock")
async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misses():
client: Final = FakeClient(_lua_ok_replies)
redis_cache: Final = FakeRedisCache(client)
redis_cache.store[cooldown_key] = active_cooldown
router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
cooldown_store: Final = router.cooldown_cache.cooldown_store
cooldown_store.put("dep-a", active_cooldown, 60)
router: Final = _router(redis_cache)
dep_a: Final = _tpm_key(router, "dep-a")
dep_b: Final = _tpm_key(router, "dep-b")
client.reply_for = _replies_with({dep_a: 5, dep_b: 7})
usage_cache: Final = router.lowesttpm_logger_v2.router_cache
with request_redis_batch_scope() as request:
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
usage_cache.in_memory_cache.set_cache(dep_a, 50)
await request.flush_all()
assert await usage_cache.async_batch_get_cache([dep_a, dep_b], local_only=True) == [50, 7]
@pytest.mark.asyncio
async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(pinned_clock: list[float]):
mget_replies: Final = iter((ConnectionError("redis down"), None))
client: Final = FakeClient(_lua_ok_replies)
redis_cache: Final = FakeRedisCache(client)
router: Final = _router(redis_cache)
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "MGET":
failure: Final = next(mget_replies)
if isinstance(failure, Exception):
return failure
return ["100" if key == _tpm_key(router, "dep-a") else None for key in command[1:]]
return _lua_ok_replies(command)
client.reply_for = replies
with request_redis_batch_scope() as request:
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await request.flush_all()
cooldown_store.evict_local("dep-a")
deployment: Final = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
after_failure: Final = len(_counter_mgets(router, client, redis_cache))
prefetched_mgets: Final = tuple(command for command in client.pipelines[0].commands if command[0] == "MGET")
fallback_cooldown_mgets: Final = tuple(
keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_key in keys
)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
deployment: Final = await _route(router)
assert len(prefetched_mgets) == 1
assert prefetched_mgets[0][1:] == (_cooldown_key("dep-b"),)
assert deployment["model_info"]["id"] == "dep-b"
assert fallback_cooldown_mgets == ((cooldown_key,),)
assert after_failure == 1
assert len(_counter_mgets(router, client, redis_cache)) == 2
assert deployment == "dep-b"
@pytest.mark.asyncio
@ -670,52 +570,35 @@ async def test_arming_outside_a_request_scope_is_a_no_op():
@pytest.mark.asyncio
async def test_simple_shuffle_prefetches_only_its_cooldown_read_into_the_admission_pipeline():
async def test_simple_shuffle_arms_no_python_prefetch_and_reads_no_cooldown_key():
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache, routing_strategy="simple-shuffle")
limiter = _limiter(redis_cache)
with request_redis_batch_scope():
with request_redis_batch_scope() as request:
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
assert request.prefetched == {}
await limiter.atomic_check_and_increment_by_n(
descriptors=[_descriptor("api_key", "k1", 10)],
increments=[{"requests": 1}],
)
deployment = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
deployment = await _route(router)
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
assert len(client.pipelines) == 1
commands = client.pipelines[0].commands
assert [c[0] for c in commands] == ["MGET", "EVALSHA"]
assert set(commands[0][1:]) == {
_cooldown_key("dep-a"),
_cooldown_key("dep-b"),
}
assert deployment in {"dep-a", "dep-b"}
assert [[c[0] for c in pipeline.commands] for pipeline in client.pipelines] == [["EVALSHA"]]
assert redis_cache.alone == []
shuffle = Router(model_list=[_deployment("dep-a")], routing_strategy="simple-shuffle")
shuffle.update_redis_cache(cache=redis_cache)
with request_redis_batch_scope() as request:
shuffle.arm_routing_read_prefetch(_MODEL_GROUP, {})
armed = request.prefetched["routing_read"]
assert isinstance(armed, ArmedRead)
assert armed.keys == {_cooldown_key("dep-a")} # no usage counters for shuffle
@pytest.mark.asyncio
@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing-v2"])
@pytest.mark.usefixtures("pinned_clock")
@pytest.mark.parametrize("with_limiter", [True, False])
async def test_requests_within_the_cooldown_read_interval_read_cooldowns_from_redis_once(
routing_strategy: str, with_limiter: bool
):
async def test_requests_within_the_read_interval_read_the_counters_from_redis_once(with_limiter: bool):
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache, routing_strategy=routing_strategy)
router = _router(redis_cache)
limiter = _limiter(redis_cache)
request_round_trips: list[tuple[int, int]] = []
request_round_trips: list[tuple[int, int]] = [] # mutable-ok: one entry per simulated request
for _ in range(3):
pipeline_count = len(client.pipelines)
@ -727,35 +610,24 @@ async def test_requests_within_the_cooldown_read_interval_read_cooldowns_from_re
descriptors=[_descriptor("api_key", "k1", 10)],
increments=[{"requests": 1}],
)
await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
await _route(router)
request_round_trips.append((len(client.pipelines) - pipeline_count, len(redis_cache.alone) - alone_count))
pipeline_mgets = [command for pipeline in client.pipelines for command in pipeline.commands if command[0] == "MGET"]
alone_mgets = [keys for command, keys in redis_cache.alone if command == "MGET"]
cooldown_keys = {
_cooldown_key("dep-a"),
_cooldown_key("dep-b"),
}
cooldown_mgets = [command[1:] for command in pipeline_mgets if cooldown_keys.intersection(command[1:])] + [
keys for keys in alone_mgets if cooldown_keys.intersection(keys)
]
assert len(cooldown_mgets) == 1
assert len(_counter_mgets(router, client, redis_cache)) == 1
if not with_limiter:
assert request_round_trips[1:] == [(0, 0), (0, 0)]
@pytest.mark.asyncio
async def test_concurrent_requests_share_one_cooldown_read_per_interval():
@pytest.mark.usefixtures("pinned_clock")
async def test_concurrent_requests_share_one_counter_read_per_interval():
client: Final = FakeClient(_lua_ok_replies)
redis_cache: Final = FakeRedisCache(client)
router: Final = _router(redis_cache)
first_armed: Final = asyncio.Event()
both_armed: Final = asyncio.Event()
async def route_after_both_requests_arm():
async def route_after_both_requests_arm() -> str:
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
if first_armed.is_set():
@ -763,126 +635,57 @@ async def test_concurrent_requests_share_one_cooldown_read_per_interval():
else:
first_armed.set()
await both_armed.wait()
return await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
return await _route(router)
deployments: Final = await asyncio.gather(route_after_both_requests_arm(), route_after_both_requests_arm())
cooldown_keys: Final = frozenset(
{
_cooldown_key("dep-a"),
_cooldown_key("dep-b"),
}
)
cooldown_mgets: Final = tuple(
command
for pipeline in client.pipelines
for command in pipeline.commands
if command[0] == "MGET" and cooldown_keys.intersection(command[1:])
)
assert all(deployment["model_info"]["id"] in {"dep-a", "dep-b"} for deployment in deployments)
assert len(cooldown_mgets) == 1
assert set(deployments) <= {"dep-a", "dep-b"}
assert len(_counter_mgets(router, client, redis_cache)) == 1
@pytest.mark.asyncio
async def test_the_prefetch_reads_cooldowns_again_once_the_read_interval_elapses(monkeypatch):
first_time: Final = 1_000_000.0
monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time)
active_cooldown = {
"exception_received": "429",
"status_code": "429",
"timestamp": _FAR_FUTURE,
"cooldown_time": 60,
}
mget_results = iter((None, active_cooldown))
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "MGET":
result = next(mget_results)
return [
None if result is None or key != _cooldown_key("dep-a") else json.dumps(result) for key in command[1:]
]
return _lua_ok_replies(command)
client = FakeClient(replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache, routing_strategy="simple-shuffle")
cooldown_store = router.cooldown_cache.cooldown_store
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
monkeypatch.setattr(
dual_cache_module.time,
"time",
lambda: first_time + POLICY.redis_read_interval + 1,
)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
deployment = await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
cooldown_keys = {
_cooldown_key("dep-a"),
_cooldown_key("dep-b"),
}
cooldown_mgets = [
command
for pipeline in client.pipelines
for command in pipeline.commands
if command[0] == "MGET" and cooldown_keys.intersection(command[1:])
]
assert len(cooldown_mgets) == 2
assert deployment["model_info"]["id"] == "dep-b"
@pytest.mark.asyncio
async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(monkeypatch):
first_time: Final = 1_000_000.0
pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time)
monkeypatch.setattr(
"litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime",
lambda: pinned,
)
async def test_the_prefetch_reads_the_counters_again_once_the_read_interval_elapses(pinned_clock: list[float]):
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache, routing_strategy="usage-based-routing-v2")
cooldown_store = router.cooldown_cache.cooldown_store
router = _router(redis_cache)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await _route(router)
client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100})
pinned_clock[0] = _START + router.lowesttpm_logger_v2.router_cache.redis_batch_cache_expiry + 1
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
deployment = await _route(router)
assert len(_counter_mgets(router, client, redis_cache)) == 2
assert deployment == "dep-b"
@pytest.mark.asyncio
async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(pinned_clock: list[float]):
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache)
usage_cache = router.lowesttpm_logger_v2.router_cache
time_offset = POLICY.redis_read_interval + 0.5
assert time_offset < usage_cache.redis_batch_cache_expiry
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
await _route(router)
monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time + time_offset)
for key in _counter_keys(router, ("dep-a",)):
usage_cache.in_memory_cache.set_cache(key, 1)
pinned_clock[0] = _START + usage_cache.redis_batch_cache_expiry + 1
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
await _route(router)
cooldown_keys = frozenset(
{
_cooldown_key("dep-a"),
_cooldown_key("dep-b"),
}
)
second_pipeline_mgets = tuple(command for command in client.pipelines[1].commands if command[0] == "MGET")
assert len(client.pipelines) == 2
assert len(second_pipeline_mgets) == 1
assert frozenset(second_pipeline_mgets[0][1:]) == cooldown_keys
assert frozenset(second_pipeline_mgets[0][1:]) == _counter_keys(router, ("dep-b",))
@pytest.mark.asyncio
@ -897,6 +700,7 @@ async def test_two_backends_flush_concurrently_one_pipeline_each():
@pytest.mark.asyncio
@pytest.mark.usefixtures("pinned_clock")
async def test_a_single_lua_group_rides_the_pipeline_with_the_armed_routing_read():
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
@ -1079,19 +883,3 @@ async def test_identity_prefetch_is_one_mget_after_which_hits_and_misses_alike_c
assert await cache.async_get_cache("end_user_id:eu-miss") is None
assert len(client.pipelines) == 1 and redis_cache.alone == []
assert cache.in_memory_cache.get_cache("end_user_id:eu-miss") is None
@pytest.mark.parametrize(
("cooldown_keys", "usage_keys", "expected"),
[
(("cooldown:a",), (), ROUTER_COOLDOWNS_TARGET),
((), ("usage:a",), ROUTER_USAGE_TARGET),
(("cooldown:a",), ("usage:a",), "router_cooldowns_usage"),
],
)
def test_the_routing_read_family_follows_the_keys_that_are_actually_due(
cooldown_keys: tuple[str, ...], usage_keys: tuple[str, ...], expected: str
):
"""A routing MGET is ``router_cooldowns`` when only cooldown keys go out, ``router_usage`` when the
cooldowns were already in memory and only usage counters go out, and the combined family otherwise."""
assert routing_target((TierRead(DualCache(), cooldown_keys), TierRead(DualCache(), usage_keys))) == expected

15
tests/unit/conftest.py
View file

@ -16,12 +16,27 @@ import pytest
from pytest_socket import enable_socket, socket_allow_hosts
if TYPE_CHECKING:
from unittest.mock import MagicMock
from litellm.caching.redis_cache import RedisCache
RouterRedis: TypeAlias = tuple["RedisCache", asyncio.Queue[tuple[tuple[object, ...], str | None, str | None]]]
RedisWriteGate: TypeAlias = tuple[asyncio.Event, asyncio.Event, asyncio.Event]
def redis_cache_double() -> "MagicMock":
"""A `RedisCache` double for router tests. The native router stores connect from its settings, which name
a closed local port, so they serve their memory tier while the double answers the Python Redis calls."""
from unittest.mock import MagicMock
from litellm.caching.redis_cache import RedisCache
redis: Final = MagicMock(spec=RedisCache)
redis.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2}
redis.namespace = None
return redis
@pytest.fixture
def redis_write_gate() -> RedisWriteGate:
released: Final = asyncio.Event()

289
tests/unit/redis/stores/test_router_cooldowns.py
View file

@ -1,112 +1,231 @@
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from unittest.mock import AsyncMock, MagicMock
import ast
import asyncio
import json
import logging
import time
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field, replace
from typing import Final
import pytest
import litellm
from litellm._internal_context import current_service_target
from litellm.caching.redis_batch import request_redis_batch_scope
from litellm.caching.redis_cache import RedisCache
from litellm.caching.redis_cache import RedisCache, redis_breaker_state_counts
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.redis.stores._connection import redis_kwargs
from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore
from litellm.rust_bridge._native import NativeRouterCooldowns
from litellm.rust_bridge.host import settings
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
if TYPE_CHECKING:
from tests.unit.conftest import RouterRedis
@dataclass(frozen=True, slots=True)
class _WireCooldown:
"""One cooldown as main's ``RedisCache`` writes it (``str(value)``) and reads it (``json.loads`` then
``ast.literal_eval``). The same literals sit in ``litellm-rust/crates/redis-store/tests/support/mod.rs``,
where real Redis is checked to hold them after the Rust store writes."""
cooldown: CooldownCacheValue
str_dict: str
json: str
ttl_seconds: float
redis_ttl_seconds: int
_WIRE_COOLDOWNS: Final = (
_WireCooldown(
cooldown={
"exception_received": "RateLimitError: limited é 'quoted'",
"status_code": "429",
"timestamp": 1730000000.25,
"cooldown_time": 60,
},
str_dict="{'exception_received': \"RateLimitError: limited é 'quoted'\", 'status_code': '429', 'timestamp': 1730000000.25, 'cooldown_time': 60}",
json='{"exception_received": "RateLimitError: limited \\u00e9 \'quoted\'", "status_code": "429", "timestamp": 1730000000.25, "cooldown_time": 60}',
ttl_seconds=90.5,
redis_ttl_seconds=90,
),
_WireCooldown(
cooldown={
"exception_received": "Timeout",
"status_code": "408",
"timestamp": 1730000000.5,
"cooldown_time": 0.5,
},
str_dict="{'exception_received': 'Timeout', 'status_code': '408', 'timestamp': 1730000000.5, 'cooldown_time': 0.5}",
json='{"exception_received": "Timeout", "status_code": "408", "timestamp": 1730000000.5, "cooldown_time": 0.5}',
ttl_seconds=30.0,
redis_ttl_seconds=30,
),
)
def _entry(model_id: str = "dep") -> CooldownCacheValue:
return {"exception_received": model_id, "status_code": "429", "timestamp": 1000.0, "cooldown_time": 120.0}
@dataclass
class Clock:
now: float = 1000.0
failure: RuntimeError | None = None
class FakeCooldowns:
redis_kwargs: Mapping[str, object] | None
namespace: str | None
stored: dict[str, CooldownCacheValue] = field(default_factory=dict)
fails: dict[str, int] = field(default_factory=dict)
targets: list[str | None] = field(default_factory=list)
smallest: float | None = None
def __call__(self) -> float:
if self.failure is not None:
raise self.failure
return self.now
def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None:
self.targets.append(current_service_target())
self.stored[model_id] = CooldownCacheValue(**cooldown) # pyright: ignore[reportArgumentType] # test double
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int:
self.targets.append(current_service_target())
self.fails[model_id] = self.fails.get(model_id, 0) + 1
return self.fails[model_id]
def active(self, model_ids: Sequence[str]) -> list[tuple[str, CooldownCacheValue]]:
self.targets.append(current_service_target())
return [(model_id, self.stored[model_id]) for model_id in model_ids if model_id in self.stored]
def aactive(self, model_ids: Sequence[str]) -> "asyncio.Future[list[tuple[str, CooldownCacheValue]]]":
future: Final[asyncio.Future[list[tuple[str, CooldownCacheValue]]]] = asyncio.get_running_loop().create_future()
future.set_result(self.active(model_ids))
return future
def minimum(self, model_ids: Sequence[str]) -> float | None:
return self.smallest
def evict_local(self, model_id: str) -> None:
self.stored.pop(model_id, None)
def flush_local_counters(self) -> None:
self.fails.clear()
def snapshot(self) -> dict[str, CooldownCacheValue]:
return dict(self.stored)
def _entry(timestamp: float = 1000.0, ttl: float = 120.0) -> CooldownCacheValue:
return {"exception_received": "limited", "status_code": "429", "timestamp": timestamp, "cooldown_time": ttl}
@dataclass
class Handles:
built: list[FakeCooldowns] = field(default_factory=list)
def __call__(self, connection: RedisCache | None) -> FakeCooldowns:
handle: Final = FakeCooldowns(
None if connection is None else redis_kwargs(connection),
None if connection is None else connection.namespace,
)
self.built.append(handle)
return handle
@pytest.mark.parametrize("suffix", (None, "", "timeout"))
def test_allowed_fails_preserves_local_count_during_redis_failure(suffix: str | None) -> None:
redis: Final = MagicMock(spec=RedisCache)
redis.increment_cache.side_effect = ConnectionError("offline")
clock: Final = Clock()
store: Final = RouterCooldownStore(redis, clock=clock)
assert store.increment_allowed_fails("dep", suffix, 30.0) == 1
assert store.increment_allowed_fails("dep", suffix, 30.0) == 2
clock.now += 31.0
assert store.increment_allowed_fails("dep", suffix, 30.0) == 1
expected: Final = "deployment:dep:allowed_fails" + (f":{suffix}" if suffix else "")
redis.increment_cache.assert_called_with(expected, 1, ttl=30.0)
def _redis(namespace: str | None = None) -> RedisCache:
return RedisCache(host="redis.internal", port=6380, namespace=namespace)
def test_sync_wire_write_and_reads_keep_service_target_and_float_ttl() -> None:
redis: Final = MagicMock(spec=RedisCache)
seen: Final[list[str | None]] = []
redis.set_cache.side_effect = lambda *args, **kw: seen.append(current_service_target())
redis.batch_get_cache.side_effect = lambda **kw: {key: _entry() for key in kw["key_list"]}
store: Final = RouterCooldownStore(redis, clock=Clock())
store.put("dep", _entry(ttl=0.5), 0.5)
redis.set_cache.assert_called_once_with("deployment:dep:cooldown", _entry(ttl=0.5), ttl=0.5)
reader: Final = RouterCooldownStore(redis, clock=Clock())
assert reader.active_sync(["dep"], None) == (("dep", _entry()),)
assert reader.minimum(["dep"], 60.0, None) == 120.0
assert seen == ["router_cooldowns"]
def test_the_store_hands_the_native_handle_redis_kwargs_and_namespace_not_a_client() -> None:
handles: Final = Handles()
RouterCooldownStore(_redis("tenant"), handle=handles)
assert handles.built[0].redis_kwargs == {"host": "redis.internal", "port": 6380, "socket_timeout": 5.0}
assert handles.built[0].namespace == "tenant"
def test_late_attach_rebuilds_once_and_later_attaches_are_ignored() -> None:
handles: Final = Handles()
store: Final = RouterCooldownStore(None, handle=handles)
store.put("dep", _entry(), 60.0)
store.attach_redis(None)
store.attach_redis(_redis())
store.attach_redis(_redis("other"))
assert [handle.redis_kwargs is not None for handle in handles.built] == [False, True]
assert handles.built[1].namespace is None
assert store.active_sync(["dep"], None) == ()
def test_writes_and_reads_carry_the_router_cooldowns_service_target() -> None:
handles: Final = Handles()
store: Final = RouterCooldownStore(_redis(), handle=handles)
store.put("dep", _entry(), 60.0)
assert store.increment_allowed_fails("dep", "timeout", 30.0) == 1
assert store.active_sync(["dep", "other"], None) == (("dep", _entry()),)
assert handles.built[0].targets == ["router_cooldowns"] * 3
assert current_service_target() is None
@pytest.mark.asyncio
async def test_remote_backfill_expires_at_corrected_sixty_second_cap() -> None:
clock: Final = Clock()
redis: Final = MagicMock(spec=RedisCache)
redis.async_batch_get_cache = AsyncMock(return_value={"deployment:dep:cooldown": _entry()})
store: Final = RouterCooldownStore(redis, clock=clock)
assert await store.active(["dep"], None) == (("dep", _entry()),)
clock.now += 61.0
assert await store.active(["dep"], None) == (("dep", _entry()),)
assert redis.async_batch_get_cache.await_count == 2
clock.now += 60.0
assert await store.active(["dep"], None) == ()
@pytest.mark.parametrize(("smallest", "expected"), ((None, 60.0), (0, 60.0), (5, 5)))
def test_minimum_falls_back_to_the_default_when_nothing_or_zero_is_stored(
smallest: float | None, expected: float
) -> None:
handles: Final = Handles()
store: Final = RouterCooldownStore(None, handle=handles)
handles.built[0].smallest = smallest
assert store.minimum(["dep"], 60.0, None) == expected
@pytest.mark.asyncio
async def test_lazy_attach_and_failure_retry_do_not_throttle_next_read() -> None:
store: Final = RouterCooldownStore(None, clock=Clock())
assert await store.active(["dep"], None) == ()
redis: Final = MagicMock(spec=RedisCache)
redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("offline"))
store.attach_redis(redis)
assert await store.active(["dep"], None) == ()
redis.async_batch_get_cache.side_effect = None
redis.async_batch_get_cache.return_value = {"deployment:dep:cooldown": _entry()}
assert await store.active(["dep"], None) == (("dep", _entry()),)
assert redis.async_batch_get_cache.await_count == 2
@pytest.mark.parametrize("case", _WIRE_COOLDOWNS, ids=("rate-limit", "timeout"))
def test_the_wire_cases_are_what_mains_codec_writes_and_reads(case: _WireCooldown) -> None:
assert str(case.cooldown) == case.str_dict
assert ast.literal_eval(case.str_dict) == case.cooldown
assert json.loads(case.json) == case.cooldown
assert int(case.ttl_seconds) == case.redis_ttl_seconds
def test_minimum_preserves_zero_default_and_does_not_correct_expired_entries() -> None:
store: Final = RouterCooldownStore(None, clock=Clock())
store.put("zero", _entry(ttl=0), 600)
store.put("expired", _entry(timestamp=0, ttl=20), 600)
assert store.minimum(["zero"], 60, None) == 60
assert store.minimum(["expired"], 60, None) == 20
assert store.active_sync(["expired"], None) == ()
@pytest.mark.parametrize(
"setting",
({"ssl_ciphers": "HIGH"}, {"ssl": True, "ssl_validate_ocsp": True}, {"credential_provider": object()}),
)
def test_a_redis_setting_the_native_client_cannot_honor_never_fails_construction(
setting: dict[str, object], caplog: pytest.LogCaptureFixture
) -> None:
redis: Final = RedisCache(host="127.0.0.1", port=1, **setting)
with caplog.at_level(logging.WARNING):
store: Final = RouterCooldownStore(redis)
cooldown: Final = CooldownCacheValue(
exception_received="model", status_code="429", timestamp=time.time(), cooldown_time=120.0
)
store.put("model", cooldown, 60.0)
assert store.active_sync(["model"], None) == (("model", cooldown),)
assert next(iter(setting)) in caplog.text
@pytest.mark.asyncio
async def test_prefetch_backfill_failure_allows_an_immediate_read_retry(router_redis: "RouterRedis") -> None:
redis, _ = router_redis
clock: Final = Clock()
value: Final = _entry()
writer: Final = RouterCooldownStore(redis, clock=clock)
reader: Final = RouterCooldownStore(redis, clock=clock)
writer.put("dep", value, 120.0)
class _RecordingSecrets(CustomSecretManager):
def __init__(self, values: Mapping[str, str]) -> None:
super().__init__(secret_manager_name="router_cooldowns_test") # pyright: ignore[reportUnknownMemberType] # CustomSecretManager takes untyped **kwargs
self.values: Final = values
self.asked: tuple[str, ...] = ()
with request_redis_batch_scope() as request:
reader.arm_routing_read(("dep",), None)
clock.failure = RuntimeError("backfill clock unavailable")
await request.flush_all()
async def async_read_secret(
self, secret_name: str, optional_params: dict[str, object] | None = None, timeout: object = None
) -> str | None:
raise AssertionError("get_secret reads custom managers synchronously")
clock.failure = None
assert await reader.active(("dep",), None) == (("dep", value),)
def sync_read_secret(
self, secret_name: str, optional_params: dict[str, object] | None = None, timeout: object = None
) -> str | None:
self.asked = (*self.asked, secret_name)
return self.values.get(secret_name)
def test_the_native_store_reads_redis_secrets_through_the_secret_manager(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("REDIS_PASSWORD", raising=False)
manager: Final = _RecordingSecrets({"REDIS_PASSWORD": "vault-password"})
monkeypatch.setattr(litellm, "secret_manager_client", manager)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
configured: Final = settings.secret_manager
monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=False))
NativeRouterCooldowns({"host": "127.0.0.1", "port": 1}, namespace="secret-manager-test")
assert "REDIS_PASSWORD" in manager.asked
def test_the_native_breaker_counts_in_the_redis_breaker_state_gauge(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true")
monkeypatch.setenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", "1")
open_before: Final = redis_breaker_state_counts().get("open", 0)
store: Final = NativeRouterCooldowns({"host": "127.0.0.1", "port": 1}, namespace="breaker-gauge-test")
store.put("model", _entry("model"), 60.0)
assert redis_breaker_state_counts().get("open", 0) == open_before + 1

142
tests/unit/redis/stores/test_router_usage.py
View file

@ -1,12 +1,85 @@
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import asyncio
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Final, TypeVar
import pytest
from litellm import Router
from litellm._internal_context import current_service_target
from litellm.caching.redis_cache import RedisCache
from litellm.redis.stores.router_usage import RouterUsageStore, UsageCounter
from litellm.redis.stores.router_usage import NativeCounter, RouterUsageStore, UsageCounter
T = TypeVar("T")
def _done(value: T) -> "asyncio.Future[T]":
future: Final[asyncio.Future[T]] = asyncio.get_running_loop().create_future()
future.set_result(value)
return future
@dataclass
class FakeUsage:
connected: bool
counts: dict[NativeCounter, float] = field(default_factory=dict)
calls: list[tuple[str, object, str | None]] = field(default_factory=list)
def increment(self, counter: NativeCounter, value: float, ttl_seconds: int) -> "asyncio.Future[float]":
self.calls.append(("increment", (counter, value, ttl_seconds), current_service_target()))
self.counts[counter] = self.counts.get(counter, 0) + value
return _done(self.counts[counter])
def increment_deployment(
self, deployment_id: str, model: str | None, minute: str, tokens: float, requests: int
) -> "asyncio.Future[None]":
self.calls.append(("deployment", (deployment_id, model, minute, tokens, requests), current_service_target()))
for kind, value in (("tpm", tokens), ("rpm", requests)):
counter: NativeCounter = (deployment_id, model, minute, kind) # pyright: ignore[reportAssignmentType] # literal kinds
self.counts[counter] = self.counts.get(counter, 0) + value
return _done(None)
def get_many(self, counters: Sequence[NativeCounter]) -> "asyncio.Future[list[float | None]]":
return _done([self.counts.get(counter) for counter in counters])
def flush_local_counters(self) -> None:
self.counts.clear()
def snapshot(self) -> dict[str, float]:
return {}
@dataclass
class Handles:
built: list[FakeUsage] = field(default_factory=list)
def __call__(self, connection: RedisCache | None) -> FakeUsage:
self.built.append(FakeUsage(connection is not None))
return self.built[-1]
@pytest.mark.asyncio
async def test_counters_cross_the_bridge_as_domain_tuples_under_the_router_usage_target() -> None:
handles: Final = Handles()
store: Final = RouterUsageStore(None, handle=handles)
counter: Final = UsageCounter("dep", None, "12-34", "rpm")
await store.async_increment(counter, 1, 60)
await store.increment_deployment("dep", "model", "12-34", 7.5, 1, None)
assert handles.built[0].calls == [
("increment", (("dep", None, "12-34", "rpm"), 1, 60), "router_usage"),
("deployment", ("dep", "model", "12-34", 7.5, 1), "router_usage"),
]
assert await store.async_batch_get((counter, UsageCounter("dep", "model", "12-34", "tpm"))) == [1, 7.5]
def test_late_attach_rebuilds_the_handle_once() -> None:
handles: Final = Handles()
store: Final = RouterUsageStore(None, handle=handles)
store.attach_redis(RedisCache(host="redis.internal"))
store.attach_redis(RedisCache(host="other"))
assert [handle.connected for handle in handles.built] == [False, True]
@pytest.mark.asyncio
@ -22,66 +95,3 @@ async def test_router_reset_clears_usage_and_failure_counters(operation: str) ->
getattr(router, operation)()
assert await usage.async_batch_get(counters) == [None, None]
assert cooldowns.increment_allowed_fails("dep", None, 30.0) == 1
@pytest.mark.asyncio
async def test_deployment_increment_preserves_pipeline_ttl_and_integer_mirror() -> None:
redis: Final = MagicMock(spec=RedisCache)
redis.async_increment_pipeline = AsyncMock(return_value=[7.75, 2.0])
store: Final = RouterUsageStore(redis)
await store.increment_deployment("dep", "anthropic/model", "12-34", 7.75, 2, None)
redis.async_increment_pipeline.assert_awaited_once_with(
increment_list=[
{"key": "global_router:dep:anthropic/model:tpm:12-34", "increment_value": 7.75, "ttl": 60},
{"key": "global_router:dep:anthropic/model:rpm:12-34", "increment_value": 2, "ttl": 60},
],
parent_otel_span=None,
)
values: Final = await store.async_batch_get(
(UsageCounter("dep", "anthropic/model", "12-34", "tpm"), UsageCounter("dep", "anthropic/model", "12-34", "rpm"))
)
assert values == [7, 2]
assert all(type(value) is int for value in values)
redis.async_batch_get_cache.assert_not_called()
assert current_service_target() is None
@pytest.mark.asyncio
async def test_memory_only_usage_increments_accumulate_without_redis() -> None:
store: Final = RouterUsageStore(None)
await store.increment_deployment("dep", "model", "12-34", 2, 1, None)
await store.increment_deployment("dep", "model", "12-34", 5, 1, None)
assert await store.async_batch_get(
(UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm"))
) == [7, 2]
@pytest.mark.asyncio
async def test_pre_call_rpm_increment_uses_same_namespace_and_service_target() -> None:
redis: Final = MagicMock(spec=RedisCache)
seen: Final[list[str | None]] = []
async def increment(*args: object, **kwargs: object) -> float:
seen.append(current_service_target())
return 4.0
redis.async_increment = AsyncMock(side_effect=increment)
store: Final = RouterUsageStore(None)
store.attach_redis(redis)
counter: Final = UsageCounter("dep", "model", "12-34", "rpm")
assert await store.async_increment(counter, value=1, ttl=60) == 4.0
assert seen == ["router_usage"]
redis.async_increment.assert_awaited_once_with(
"global_router:dep:model:rpm:12-34", 1, parent_otel_span=None, ttl=60, refresh_ttl=False
)
@pytest.mark.asyncio
async def test_failed_usage_pipeline_keeps_the_memory_increment_and_is_readable() -> None:
redis: Final = MagicMock(spec=RedisCache)
redis.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("offline"))
store: Final = RouterUsageStore(redis)
await store.increment_deployment("dep", "model", "12-34", 5, 1, None)
assert await store.async_batch_get(
(UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm"))
) == [5, 1]

30
tests/unit/router_strategy/test_router_tag_routing.py
View file

@ -4,7 +4,6 @@
import logging
from typing import Final, Literal
from unittest.mock import patch
import pytest
@ -1846,23 +1845,20 @@ async def test_tag_fail_open_survives_deployment_narrowing(
) -> None:
router: Final = _fail_open_router(route=route)
requested_model: Final = "chat" if route == "model" else "chat-test" if route == "wildcard" else "public-chat"
with patch("time.time", return_value=1_700_000_000.0):
if narrowing == "cooldown":
router.cooldown_cache.add_deployment_to_cooldown(
model_id="opted-in",
original_exception=RuntimeError("unavailable"),
exception_status=500,
cooldown_time=60,
)
candidates: Final = await router.async_get_healthy_deployments(
model=requested_model,
request_kwargs={
"litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"},
**(
{"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {}
),
},
if narrowing == "cooldown":
router.cooldown_cache.add_deployment_to_cooldown(
model_id="opted-in",
original_exception=RuntimeError("unavailable"),
exception_status=500,
cooldown_time=60,
)
candidates: Final = await router.async_get_healthy_deployments(
model=requested_model,
request_kwargs={
"litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"},
**({"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {}),
},
)
assert isinstance(candidates, (list, tuple))
assert tuple(deployment["model_info"]["id"] for deployment in candidates) == ("remaining",)

38
tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py
View file

@ -1,7 +1,6 @@
import asyncio
import json
from itertools import chain
from typing import TYPE_CHECKING, Final
from typing import Final
from collections.abc import AsyncIterator, Mapping
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
@ -22,9 +21,6 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
)
if TYPE_CHECKING:
from tests.unit.conftest import RouterRedis
def _repinning_router(enabled: bool | None = True) -> litellm.Router:
return litellm.Router(
@ -325,38 +321,6 @@ async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: M
router.discard()
@pytest.mark.parametrize("case", ["disabled", "no-session", "generated-session", "group-disabled", "active"])
async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str, router_redis: "RouterRedis") -> None:
router: Final = _repinning_router(case != "disabled")
redis, commands = router_redis
try:
router.update_redis_cache(redis)
callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck))
if case == "group-disabled":
callback.model_group_affinity_config = {"repin-group": []}
await router._repin_session_after_success(
{
"model_info": {"id": "replacement"},
"metadata": {
"deployment_model_name": "repin-group",
**({"session_id": "repin-session"} if case != "no-session" else {}),
SESSION_ID_GENERATED_METADATA_KEY: case == "generated-session",
},
},
{"status": "completed"},
)
expected: Final = (
{f"deployment:{model_id}:cooldown" for model_id in ("original", "replacement", "other")}
if case == "active"
else set()
)
observed: Final = tuple(commands.get_nowait() for _ in range(commands.qsize()))
reads: Final = frozenset(chain.from_iterable(args[1:] for args, _, _ in observed if args[0] == "MGET"))
assert reads == expected
finally:
router.discard()
class MockResponse:
def __init__(self, json_data, status_code):
self._json_data = json_data

20
tests/unit/router_utils/test_cooldown_cache.py
View file

@ -1,3 +1,4 @@
import time
from typing import Final
import pytest
@ -5,14 +6,12 @@ import pytest
from litellm.redis.stores.router_cooldowns import RouterCooldownStore
from litellm.router_utils.cooldown_cache import CooldownCache
_NOW: Final = 1000.0
@pytest.mark.parametrize("message", ("Short error", "Rate limit " + "secret" * 100))
@pytest.mark.parametrize("ttl", (None, 0.5, 30.0))
@pytest.mark.parametrize("ttl", (None, 5.5, 30.0))
def test_cooldown_masks_exception_and_preserves_status_timestamp_and_ttl(message: str, ttl: float | None) -> None:
store: Final = RouterCooldownStore(None, clock=lambda: _NOW)
cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: _NOW)
now: Final = time.time()
store: Final = RouterCooldownStore(None)
cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: now)
cache.add_deployment_to_cooldown("dep", Exception(message), 429, ttl)
expected_ttl: Final = 60.0 if ttl is None else ttl
expected_message: Final = message if len(message) <= 50 else message[:50] + "*" * (len(message) - 50)
@ -22,7 +21,7 @@ def test_cooldown_masks_exception_and_preserves_status_timestamp_and_ttl(message
{
"exception_received": expected_message,
"status_code": "429",
"timestamp": _NOW,
"timestamp": now,
"cooldown_time": expected_ttl,
},
)
@ -31,11 +30,12 @@ def test_cooldown_masks_exception_and_preserves_status_timestamp_and_ttl(message
@pytest.mark.asyncio
async def test_sync_and_async_router_cooldown_surfaces_share_active_state() -> None:
store: Final = RouterCooldownStore(None, clock=lambda: _NOW)
cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: _NOW)
now: Final = time.time()
store: Final = RouterCooldownStore(None)
cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: now)
cache.add_deployment_to_cooldown("dep", Exception("unavailable"), 503, 30.0)
expected: Final = [
("dep", {"exception_received": "unavailable", "status_code": "503", "timestamp": _NOW, "cooldown_time": 30.0})
("dep", {"exception_received": "unavailable", "status_code": "503", "timestamp": now, "cooldown_time": 30.0})
]
assert cache.get_active_cooldowns(["absent", "dep"], None) == expected
assert await cache.async_get_active_cooldowns(["absent", "dep"], None) == expected

79
tests/unit/router_utils/test_cooldown_handlers.py
View file

@ -16,12 +16,11 @@ import respx
import litellm
from litellm import Router
from litellm._internal_context import current_service_target
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_cache import RedisCache
from litellm.redis.stores.router_cooldowns import RouterCooldowns, RouterCooldownStore
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
from litellm.redis.stores.router_cooldowns import RouterCooldowns, RouterCooldownStore
from litellm.router_api import RoutingContext
from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue
from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback
@ -378,9 +377,8 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy:
store.increment_allowed_fails.assert_called_once_with("dep-1", None, 0.0)
def test_fail_counter_is_shared_across_router_instances(self):
redis: Final = MagicMock(spec=RedisCache)
redis.increment_cache.side_effect = itertools.count(1)
workers: Final = tuple(self._make_router(store=RouterCooldownStore(redis)) for _ in range(2))
shared: Final = RouterCooldownStore(None)
workers: Final = tuple(self._make_router(store=shared) for _ in range(2))
exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4")
results: Final = [
should_cooldown_based_on_allowed_fails_policy(
@ -389,30 +387,16 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy:
for i in range(6)
]
assert results == [False, False, False, False, False, True]
assert redis.increment_cache.call_count == 6
def test_fleet_wide_count_from_redis_decides_cooldown(self):
redis: Final = MagicMock(spec=RedisCache)
redis.increment_cache.return_value = 6
ctx: Final = self._make_router(store=RouterCooldownStore(redis))
def test_fleet_wide_count_from_the_store_decides_cooldown(self):
store: Final = MagicMock(spec=RouterCooldowns)
store.increment_allowed_fails.return_value = 6
ctx: Final = self._make_router(store=store)
exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4")
assert should_cooldown_based_on_allowed_fails_policy(
ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5
)
redis.increment_cache.assert_called_once_with("deployment:dep-1:allowed_fails", 1, ttl=60.0)
def test_redis_outage_falls_back_to_this_workers_count(self):
redis: Final = MagicMock(spec=RedisCache)
redis.increment_cache.side_effect = ConnectionError("redis down")
ctx: Final = self._make_router(store=RouterCooldownStore(redis))
exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4")
assert [
should_cooldown_based_on_allowed_fails_policy(
ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5
)
for _ in range(6)
] == [False, False, False, False, False, True]
assert redis.increment_cache.call_count == 6
store.increment_allowed_fails.assert_called_once_with("dep-1", None, 60.0)
class TestRoutingGroupCooldownAlternatives:
@ -560,38 +544,6 @@ class TestTeamModelCooldownAlternatives:
)
class TestIncrementAllowedFailsServiceTarget:
def test_fail_counter_bump_declares_the_router_cooldowns_key_family(self):
"""The allowed_fails INCR is cooldown bookkeeping, so its service span must read
``redis.incr router_cooldowns`` rather than a bare ``redis.incr``."""
seen: list[str | None] = []
cache = MagicMock(spec=RedisCache)
def _increment(*_args, **_kwargs):
seen.append(current_service_target())
return 2
cache.increment_cache.side_effect = _increment
assert RouterCooldownStore(cache).increment_allowed_fails("dep-1", None, 60.0) == 2
assert seen == ["router_cooldowns"]
assert current_service_target() is None
def test_in_memory_fallback_reads_under_the_same_target(self):
seen: list[str | None] = []
cache = MagicMock(spec=RedisCache)
def _increment(*_args, **_kwargs):
seen.append(current_service_target())
raise ConnectionError("redis down")
cache.increment_cache.side_effect = _increment
store = RouterCooldownStore(cache)
assert [store.increment_allowed_fails("dep-1", None, 60.0) for _ in range(4)] == [1, 2, 3, 4]
assert seen == ["router_cooldowns"] * 4
@pytest.fixture()
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
@ -869,14 +821,14 @@ class TestExceptionTypeCountersTrackedIndependently:
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
class TestCooldownCacheTTLCorrection:
@pytest.mark.parametrize(("timestamp", "active"), ((880.0, False), (1000.0, True)))
@pytest.mark.parametrize(("age", "active"), ((120.0, False), (0.0, True)))
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.asyncio
async def test_only_unexpired_cooldowns_are_returned(self, timestamp: float, active: bool, asynchronous: bool):
store: Final = RouterCooldownStore(None, clock=lambda: 1000.0)
async def test_only_unexpired_cooldowns_are_returned(self, age: float, active: bool, asynchronous: bool):
store: Final = RouterCooldownStore(None)
cc: Final = CooldownCache(store, default_cooldown_time=60.0)
value: Final = CooldownCacheValue(
exception_received="Rate limit", status_code="429", timestamp=timestamp, cooldown_time=60.0
exception_received="Rate limit", status_code="429", timestamp=time.time() - age, cooldown_time=60.0
)
store.put("deployment", value, 600)
result: Final = (
@ -2097,7 +2049,6 @@ def test_increment_deployment_successes_for_current_minute_does_not_write_to_red
"""
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
increment_deployment_successes_for_current_minute,
)
@ -2595,7 +2546,6 @@ async def test_high_traffic_cooldowns_all_healthy_deployments():
await asyncio.sleep(0.0001)
except litellm.InternalServerError:
model_stats[model_id]["failures"] += 1
pass
except Exception as e:
print("Failed test model stats=", model_stats)
raise e
@ -2689,7 +2639,6 @@ async def test_high_traffic_cooldowns_one_bad_deployment():
await asyncio.sleep(0.0001)
except litellm.InternalServerError:
model_stats[model_id]["failures"] += 1
pass
except Exception as e:
print("Failed test model stats=", model_stats)
raise e
@ -2783,10 +2732,8 @@ async def test_high_traffic_cooldowns_one_rate_limited_deployment():
await asyncio.sleep(0.0001)
except litellm.InternalServerError:
model_stats[model_id]["failures"] += 1
pass
except litellm.RateLimitError:
model_stats[bad_deployment_id]["failures"] += 1
pass
except Exception as e:
print("Failed test model stats=", model_stats)
raise e

54
tests/unit/router_utils/test_routing_read_batch.py
View file

@ -1,8 +1,6 @@
"""
One Redis round trip per request for the router's pre-call reads.
Before `RoutingReadBatch`, `async_get_available_deployment` issued one MGET for the cooldown keys
(`CooldownCache`) and a second one for the tpm/rpm counters (`LowestTPMLoggingHandler_v2`).
The router's pre-call reads. Cooldowns go through the native store, and the tpm/rpm selector
counters (`LowestTPMLoggingHandler_v2`) ride one MGET on the Python Redis.
"""
import time
@ -15,6 +13,7 @@ import litellm
from litellm import Router
from litellm.caching.redis_cache import RedisCache
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
from tests.unit.conftest import redis_cache_double
_MODEL_GROUP = "claude"
_MESSAGES = [{"role": "user", "content": "ping"}]
@ -34,7 +33,7 @@ def _redis_answering(values_by_key_prefix: dict[str, object]) -> MagicMock:
def _mget(key_list, parent_otel_span=None):
return {key: values_by_key_prefix.get(key.rsplit(":", 1)[0], values_by_key_prefix.get(key)) for key in key_list}
redis = MagicMock(spec=RedisCache)
redis = redis_cache_double()
redis.async_batch_get_cache = AsyncMock(side_effect=_mget)
return redis
@ -55,12 +54,14 @@ def _redis_key_families(redis: MagicMock) -> list[list[str]]:
]
def _cooldown(seconds: float) -> dict:
return {"exception_received": "429", "status_code": "429", "timestamp": time.time(), "cooldown_time": seconds}
def _cool_down(router: Router, model_id: str, started: float) -> None:
router.cooldown_cache.cooldown_store.put(
model_id, {"exception_received": "429", "status_code": "429", "timestamp": started, "cooldown_time": 60}, 60
)
@pytest.mark.asyncio
async def test_usage_based_routing_reads_cooldowns_and_counters_in_one_redis_round_trip():
async def test_usage_based_routing_reads_the_counters_in_one_python_redis_round_trip():
redis = _redis_answering({})
router = _router(redis, "usage-based-routing-v2")
@ -73,10 +74,8 @@ async def test_usage_based_routing_reads_cooldowns_and_counters_in_one_redis_rou
"dep-a:anthropic/claude-x:tpm",
"dep-b:anthropic/claude-x:rpm",
"dep-b:anthropic/claude-x:tpm",
"deployment:dep-a:cooldown",
"deployment:dep-b:cooldown",
]
], "cooldown state and usage counters must arrive in one MGET"
], "the usage counters must arrive in one MGET, without cooldown keys"
@pytest.mark.asyncio
@ -112,21 +111,19 @@ async def test_usage_based_routing_still_batches_when_the_strategy_is_a_fixed_si
"dep-a:anthropic/claude-x:tpm",
"dep-b:anthropic/claude-x:rpm",
"dep-b:anthropic/claude-x:tpm",
"deployment:dep-a:cooldown",
"deployment:dep-b:cooldown",
]
], "the subclassed strategy must still get the batched read, not a second MGET"
router.cache.async_batch_get_cache.assert_not_awaited()
@pytest.mark.asyncio
async def test_simple_shuffle_still_reads_only_cooldowns():
async def test_simple_shuffle_reads_nothing_from_the_python_redis():
redis = _redis_answering({})
router = _router(redis, "simple-shuffle")
await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
assert _redis_key_families(redis) == [["deployment:dep-a:cooldown", "deployment:dep-b:cooldown"]]
assert _redis_key_families(redis) == []
@pytest.mark.asyncio
@ -151,14 +148,9 @@ async def test_batched_counters_pick_the_deployment_the_strategy_picks_reading_a
@pytest.mark.asyncio
async def test_batched_read_still_excludes_a_cooled_down_deployment():
redis = _redis_answering(
{
"dep-a:anthropic/claude-x:tpm": 100,
"dep-b:anthropic/claude-x:tpm": 10,
"deployment:dep-b:cooldown": _cooldown(seconds=60),
}
)
redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10})
router = _router(redis, "usage-based-routing-v2")
_cool_down(router, "dep-b", started=time.time())
deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
@ -168,14 +160,9 @@ async def test_batched_read_still_excludes_a_cooled_down_deployment():
@pytest.mark.asyncio
async def test_batched_read_ignores_an_expired_cooldown():
redis = _redis_answering(
{
"dep-a:anthropic/claude-x:tpm": 100,
"dep-b:anthropic/claude-x:tpm": 10,
"deployment:dep-b:cooldown": _cooldown(seconds=-1),
}
)
redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10})
router = _router(redis, "usage-based-routing-v2")
_cool_down(router, "dep-b", started=time.time() - 61)
deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
@ -184,7 +171,7 @@ async def test_batched_read_ignores_an_expired_cooldown():
@pytest.mark.asyncio
async def test_a_failed_batched_read_degrades_like_the_two_failed_reads_did():
redis = MagicMock(spec=RedisCache)
redis = redis_cache_double()
redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable"))
routed = _router(redis, "usage-based-routing-v2")
alone = _router(redis, "usage-based-routing-v2")
@ -198,16 +185,11 @@ async def test_a_failed_batched_read_degrades_like_the_two_failed_reads_did():
assert str(routed_error.value) == str(alone_error.value)
assert len(routed.cache.last_redis_batch_access_time) == 0, "a failed read must not throttle the next one"
before: Final = redis.async_batch_get_cache.await_count
redis.async_batch_get_cache.side_effect = None
redis.async_batch_get_cache.return_value = {}
assert await routed.cooldown_cache.async_get_active_cooldowns(["dep-a"], None) == []
assert redis.async_batch_get_cache.await_count == before + 1
@pytest.mark.asyncio
async def test_a_failed_batched_read_leaves_simple_shuffle_routing():
redis = MagicMock(spec=RedisCache)
redis = redis_cache_double()
redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable"))
router = _router(redis, "simple-shuffle")

81
tests/unit/rust_bridge/host/test_redis_events.py Normal file
View file

@ -0,0 +1,81 @@
from typing import Final
import pytest
from prometheus_client import REGISTRY
import litellm
from litellm.rust_bridge.host import redis_events
def _sample(name: str, labels: dict[str, str]) -> float:
return REGISTRY.get_sample_value(name, labels) or 0.0
@pytest.fixture
def prometheus(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"])
redis_events._service_logging.cache_clear()
yield
redis_events._service_logging.cache_clear()
@pytest.mark.parametrize(
("failure", "failure_class"),
[("timeout", "timeout"), ("transport", "connectivity")],
)
def test_a_health_failure_counts_as_a_failed_request_and_a_breaker_failure(
prometheus: None, failure: redis_events.Failure, failure_class: str
):
failed: Final = (
"litellm_redis_failed_requests_total",
{"redis": "redis", "error_class": "RedisCallError", "function_name": "async_increment"},
)
breaker: Final = ("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class})
before: Final = (_sample(*failed), _sample(*breaker))
redis_events.record_redis_call("async_increment", 0.25, failure, "caller.py")
assert (_sample(*failed), _sample(*breaker)) == (before[0] + 1, before[1] + 1)
@pytest.mark.parametrize("failure", ["server", "decode"])
def test_a_rejected_command_is_a_failed_request_but_not_a_breaker_failure(
prometheus: None, failure: redis_events.Failure
):
breaker: Final = {
failure_class: _sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class})
for failure_class in ("timeout", "connectivity")
}
failed: Final = (
"litellm_redis_failed_requests_total",
{"redis": "redis", "error_class": "RedisCallError", "function_name": "async_increment"},
)
before: Final = _sample(*failed)
redis_events.record_redis_call("async_increment", 0.25, failure, None)
assert _sample(*failed) == before + 1
assert {
failure_class: _sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class})
for failure_class in ("timeout", "connectivity")
} == breaker
def test_a_success_is_counted_and_a_breaker_rejection_is_not(prometheus: None):
total: Final = ("litellm_redis_total_requests_total", {"redis": "redis"})
before: Final = _sample(*total)
redis_events.record_redis_call("async_get_cache", 0.01, None, None)
redis_events.record_redis_call("async_get_cache", 0.0, "breaker_open", None)
assert _sample(*total) == before + 1
@pytest.mark.parametrize("state", ["open", "half_open", "closed"])
def test_breaker_transitions_are_counted_by_state(prometheus: None, state: redis_events.BreakerState):
labels: Final = {"state": state}
before: Final = _sample("litellm_redis_circuit_breaker_transitions_total", labels)
redis_events.record_redis_breaker(state)
assert _sample("litellm_redis_circuit_breaker_transitions_total", labels) == before + 1

25
tests/unit/rust_bridge/host/test_settings.py
View file

@ -4,6 +4,7 @@ import httpx
import pytest
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.llms.custom_httpx.http_handler import default_user_agent
from litellm.rust_bridge.host import settings
@ -46,6 +47,26 @@ def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch
)
def test_router_tiers_read_the_env_constants_and_the_litellm_global(monkeypatch: pytest.MonkeyPatch) -> None:
import litellm.constants
monkeypatch.setattr(litellm.constants, "DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", 2.5)
monkeypatch.setattr(litellm.constants, "DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 42)
monkeypatch.setattr(litellm, "default_redis_batch_cache_expiry", 30.0)
assert settings.router_tiers() == settings.RouterTierSettings(
cooldown_redis_read_interval=2.5,
usage_redis_read_interval=30.0,
max_redis_batch_cache_size=42,
)
def test_router_tiers_fall_back_to_dual_caches_ten_second_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "default_redis_batch_cache_expiry", None)
assert settings.router_tiers().usage_redis_read_interval == DualCache().redis_batch_cache_expiry
def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_USER_AGENT", "operator/1")
monkeypatch.setenv("SSL_VERIFY", "false")
@ -114,9 +135,7 @@ def test_secret_manager_is_native_only_when_a_readable_system_is_configured(
monkeypatch.setattr(litellm, "_key_management_system", system)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode))
assert settings.secret_manager() == settings.SecretManager(
readable=access_mode != "write_only", native=native
)
assert settings.secret_manager() == settings.SecretManager(readable=access_mode != "write_only", native=native)
def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None:

2
tests/unit/test_internal_context.py
View file

@ -32,7 +32,7 @@ _CACHE_PARAMETER_HELPERS: Final = frozenset(
"litellm/llms/anthropic/prompt_cache_prediction.py",
"litellm/proxy/common_utils/cache_coordinator.py",
"litellm/proxy/common_utils/user_api_key_cache.py",
"litellm/redis/stores/_routing_prefetch.py",
"litellm/router_utils/routing_prefetch.py",
"litellm/utils.py",
}
)

77
tests/unit/test_router/test_router.py
View file

@ -105,12 +105,12 @@ from litellm.utils import _invalidate_model_cost_lowercase_map
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
from tests.large_text import text
from tests.unit.conftest import redis_cache_double
import traceback
import inspect
from typing import List, Optional
if TYPE_CHECKING:
from tests.unit.conftest import RedisWriteGate, RouterRedis
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from litellm.integrations.otel.logger import OpenTelemetryV2
@ -15461,11 +15461,10 @@ class TestClaudeCodeSubagentSessionRouterBinding:
@pytest.mark.asyncio
async def test_redis_cleanup_failure_does_not_reject_a_subagent_request(self):
from litellm.caching.caching import RedisCache
router = self._router()
del router.complexity_routers["smart-router"]
redis_cache = MagicMock(spec=RedisCache)
redis_cache = redis_cache_double()
redis_cache.async_get_cache = AsyncMock(return_value="smart-router")
redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable"))
router.update_redis_cache(cache=redis_cache)
@ -15480,7 +15479,6 @@ class TestClaudeCodeSubagentSessionRouterBinding:
@pytest.mark.asyncio
async def test_redis_read_failure_does_not_reject_a_subagent_request(self):
from litellm.caching.caching import RedisCache
router = self._router()
request_kwargs = self._request_kwargs(agent_id="agent-1234")
@ -15490,7 +15488,7 @@ class TestClaudeCodeSubagentSessionRouterBinding:
cache_key,
"smart-router",
)
redis_cache = MagicMock(spec=RedisCache)
redis_cache = redis_cache_double()
redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open"))
router.update_redis_cache(cache=redis_cache)
@ -15505,10 +15503,9 @@ class TestClaudeCodeSubagentSessionRouterBinding:
@pytest.mark.asyncio
async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self):
from litellm.caching.caching import RedisCache
router = self._router()
redis_cache = MagicMock(spec=RedisCache)
redis_cache = redis_cache_double()
redis_cache.async_get_cache = AsyncMock(return_value="smart-router")
redis_cache.async_set_cache = AsyncMock(side_effect=Exception("redis unavailable"))
router.update_redis_cache(cache=redis_cache)
@ -15532,10 +15529,9 @@ class TestClaudeCodeSubagentSessionRouterBinding:
async def test_subagents_follow_the_main_threads_latest_router_across_workers(self):
from types import SimpleNamespace
from litellm.caching.caching import RedisCache
shared_binding = SimpleNamespace(value=None)
shared_redis = MagicMock(spec=RedisCache)
shared_redis = redis_cache_double()
shared_redis.async_get_cache = AsyncMock(side_effect=lambda key, **_: shared_binding.value)
shared_redis.async_set_cache = AsyncMock(
side_effect=lambda key, value, **_: setattr(shared_binding, "value", value)
@ -15563,11 +15559,10 @@ class TestClaudeCodeSubagentSessionRouterBinding:
@pytest.mark.asyncio
async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self):
from litellm.caching.caching import RedisCache
router = self._router()
router.complexity_routers.clear()
redis_cache = MagicMock(spec=RedisCache)
redis_cache = redis_cache_double()
redis_cache.async_get_cache = AsyncMock(return_value=None)
redis_cache.async_set_cache = AsyncMock()
redis_cache.async_delete_cache = AsyncMock()
@ -24154,11 +24149,7 @@ def test_router_keeps_wif_secret_pointers_unresolved(monkeypatch):
@pytest.mark.asyncio
async def test_failure_rpm_increment_declares_the_router_usage_key_family(router_redis: "RouterRedis"):
"""The RPM bump a failed call still earns is router usage bookkeeping, so its Redis span
reads ``redis.incr router_usage`` rather than a bare ``redis.incr``."""
from litellm._internal_context import current_service_target
async def test_a_failed_call_still_counts_one_rpm():
router = Router(
model_list=[
{
@ -24168,8 +24159,6 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family(router
}
]
)
redis, commands = router_redis
router.update_redis_cache(redis)
await router.async_deployment_callback_on_failure(
kwargs={
"call_type": "acompletion",
@ -24183,11 +24172,18 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family(router
end_time=None,
)
observed: Final = tuple(commands.get_nowait() for _ in range(commands.qsize()))
assert tuple(
target for args, target, _ in observed if isinstance(args[0], bytes) and b"INCRBYFLOAT" in args[0]
) == ("router_usage",)
assert current_service_target() is None
assert _usage_by_kind(router) == {"rpm": 1}
def _usage_by_kind(router: Router) -> dict[str, float]:
return {key.rsplit(":", 2)[-2]: value for key, value in router._core_usage_store.snapshot().items()}
async def _usage_settles(router: Router, kinds: int) -> dict[str, float]:
async with asyncio.timeout(5):
while len(_usage_by_kind(router)) < kinds:
await asyncio.sleep(0.01)
return _usage_by_kind(router)
@pytest.fixture
@ -24207,9 +24203,7 @@ def v2_span_exporter(monkeypatch):
@pytest.mark.asyncio
async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_model_group(
v2_span_exporter, router_redis: "RouterRedis"
):
async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_model_group(v2_span_exporter):
"""Picking a deployment opens ``route {model_group}`` (the requested group, not the deployment
it picks) under the server span, and the cooldown reads it issues run inside it, so their Redis
spans nest there instead of lying flat under the request."""
@ -24229,9 +24223,6 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo
},
]
)
redis, commands = router_redis
router.update_redis_cache(redis)
with TracerProvider().get_tracer("test").start_as_current_span("POST /v1/chat/completions") as server_span:
deployment = await router.async_get_available_deployment(model="gpt-group", request_kwargs={})
@ -24240,8 +24231,6 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo
assert route_span.name == "route gpt-group"
assert route_span.parent is not None and route_span.parent.span_id == server_span.get_span_context().span_id
assert route_span.end_time is not None
observed: Final = tuple(commands.get_nowait() for _ in range(commands.qsize()))
assert tuple(span for args, _, span in observed if args[0] == "MGET") == ("route gpt-group",)
def _record_phase_events(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, dict[str, str | int]]]:
@ -24996,7 +24985,7 @@ async def test_router_get_model_group_usage_wildcard_routes():
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_call_router_callbacks_on_success(router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"):
async def test_call_router_callbacks_on_success():
router = Router(
model_list=[
{
@ -25007,26 +24996,18 @@ async def test_call_router_callbacks_on_success(router_redis: "RouterRedis", red
]
)
redis, _ = router_redis
started, _, _ = redis_write_gate
router.update_redis_cache(redis)
response: Final = await router.acompletion(
model="gemini/gemini-2.5-flash",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response="Hello, I'm good.",
)
await asyncio.wait_for(started.wait(), timeout=5)
client: Final = redis.init_async_client()
keys: Final = tuple(await client.keys("global_router:*"))
values: Final = {key.rsplit(b":", 2)[-2]: float(await client.get(key)) for key in keys}
assert len(keys) == 2
assert values == {b"tpm": response.usage.total_tokens, b"rpm": 1}
assert await _usage_settles(router, kinds=2) == {"tpm": response.usage.total_tokens, "rpm": 1}
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.serial
@pytest.mark.asyncio
async def test_call_router_callbacks_on_failure(router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"):
async def test_call_router_callbacks_on_failure():
router = Router(
model_list=[
{
@ -25037,9 +25018,6 @@ async def test_call_router_callbacks_on_failure(router_redis: "RouterRedis", red
]
)
redis, _ = router_redis
started, _, _ = redis_write_gate
router.update_redis_cache(redis)
with pytest.raises(litellm.RateLimitError):
await router.acompletion(
model="gemini/gemini-2.5-flash",
@ -25047,12 +25025,7 @@ async def test_call_router_callbacks_on_failure(router_redis: "RouterRedis", red
mock_response="litellm.RateLimitError",
num_retries=0,
)
await asyncio.wait_for(started.wait(), timeout=5)
client: Final = redis.init_async_client()
keys: Final = tuple(await client.keys("global_router:*"))
values: Final = {key.rsplit(b":", 2)[-2]: float(await client.get(key)) for key in keys}
assert len(keys) == 1
assert values == {b"rpm": 1}
assert await _usage_settles(router, kinds=1) == {"rpm": 1}
@pytest.mark.asyncio
@ -25207,7 +25180,7 @@ def test_router_get_model_access_groups(potential_access_group, expected_result)
def test_router_redis_cache():
router = Router(model_list=[{"model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}}])
redis_cache = MagicMock()
redis_cache = redis_cache_double()
router.update_redis_cache(cache=redis_cache)

50
tests/unit/test_router/test_router_helper_utils.py
View file

@ -501,7 +501,7 @@ async def test_deployment_callback_on_success(sync_mode):
@pytest.mark.asyncio
async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(router_redis: "RouterRedis"):
async def test_deployment_callback_on_success_counts_a_responses_websocket_turn():
import time
model_list = [
@ -512,8 +512,6 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(
}
]
router = Router(model_list=model_list)
redis, _ = router_redis
router.update_redis_cache(redis)
standard_logging_payload = create_standard_logging_payload()
standard_logging_payload["total_tokens"] = 17
@ -534,8 +532,7 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(
end_time=time.time(),
)
assert counter is not None
tpm_key = f"global_router:ws-deployment:None:tpm:{counter.minute}"
assert await redis.async_get_cache(key=tpm_key) == 17
assert await router._core_usage_store.async_batch_get((counter,)) == [17]
realtime_key = await router.deployment_callback_on_success(
kwargs=turn_kwargs("_arealtime"),
@ -544,7 +541,7 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(
end_time=time.time(),
)
assert realtime_key is None
assert await redis.async_get_cache(key=tpm_key) == 17
assert await router._core_usage_store.async_batch_get((counter,)) == [17]
async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]:
@ -1762,49 +1759,28 @@ def test_deployment_tracks_usage_for_every_limit_source(
assert router._deployment_tracks_usage(deployment_id, "my-unmapped-model") is tracks
def _shared_redis_stub(store: dict) -> MagicMock:
from litellm.caching.redis_cache import RedisCache
async def increment_pipeline(increment_list, **kwargs):
for op in increment_list:
store[op["key"]] = store.get(op["key"], 0.0) + op["increment_value"]
return [store[op["key"]] for op in increment_list]
async def batch_get(keys, **kwargs):
return {key: store.get(key) for key in keys}
redis_stub = MagicMock(spec=RedisCache)
redis_stub.async_increment_pipeline = increment_pipeline
redis_stub.async_batch_get_cache = batch_get
return redis_stub
@pytest.mark.asyncio
@pytest.mark.usefixtures("router_minute_pinned")
async def test_headers_on_fresh_worker_reflect_shared_redis_usage():
store: dict = {}
worker_a = _rpm_tpm_router("lit-3058-workers")
worker_b = _rpm_tpm_router("lit-3058-workers")
worker_a.update_redis_cache(_shared_redis_stub(store))
worker_b.update_redis_cache(_shared_redis_stub(store))
async def test_headers_reflect_usage_counted_by_earlier_requests():
router = _rpm_tpm_router("lit-3058-workers")
messages = [{"role": "user", "content": "hi"}]
tokens_on_a = 0
earlier_tokens = 0
for _ in range(3):
response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
tokens_on_a += response.usage.total_tokens
response = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
earlier_tokens += response.usage.total_tokens
response = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
response = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
headers = _ratelimit_headers(response)
assert headers["x-ratelimit-remaining-requests"] == 96
assert headers["x-ratelimit-remaining-tokens"] == 1000 - tokens_on_a - response.usage.total_tokens
assert headers["x-ratelimit-remaining-tokens"] == 1000 - earlier_tokens - response.usage.total_tokens
counted_tokens = tokens_on_a + response.usage.total_tokens
counted_tokens = earlier_tokens + response.usage.total_tokens
for _ in range(2):
response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
response = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
counted_tokens += response.usage.total_tokens
stream = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True)
stream = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True)
stream_headers = _ratelimit_headers(stream)
assert stream_headers["x-ratelimit-remaining-requests"] == 93
assert stream_headers["x-ratelimit-remaining-tokens"] == 1000 - counted_tokens

29
tests/unit/test_router/test_router_provider_endpoints.py
View file

@ -5,7 +5,7 @@ import json
import uuid
from collections.abc import Callable, Iterator
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final
from typing import Final
import httpx
import litellm
@ -24,8 +24,8 @@ from litellm.types.utils import ImageResponse, ModelResponse, RerankResponse
from litellm.types.videos.utils import encode_video_id_with_provider
from openai import AsyncAzureOpenAI, AsyncOpenAI
if TYPE_CHECKING:
from tests.unit.conftest import RedisWriteGate, RouterRedis
from tests.unit.conftest import redis_cache_double
EVENT_TIMEOUT_SECONDS: Final = 5
ROUTING_SELECTIONS: Final = 40
@ -536,15 +536,9 @@ def test_router_track_deployment_metrics_updates_observable_usage() -> None:
@pytest.mark.asyncio
@pytest.mark.usefixtures("router_minute_pinned")
async def test_router_success_callback_during_pre_header_increment_does_not_double_count(
router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"
) -> None:
async def test_router_success_callback_during_pre_header_increment_does_not_double_count() -> None:
capture: Final = _RouterLoggingCapture("lit-3058-race")
router: Final = _rpm_tpm_router("lit-3058-race")
redis, _ = router_redis
started, released, _ = redis_write_gate
router.update_redis_cache(redis)
released.clear()
request: Final = asyncio.create_task(
router.acompletion(
model="gpt-5-mini",
@ -554,10 +548,8 @@ async def test_router_success_callback_during_pre_header_increment_does_not_doub
)
)
await asyncio.wait_for(started.wait(), timeout=EVENT_TIMEOUT_SECONDS)
await capture.next_event()
assert get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1
released.set()
response: Final = await request
assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1)
@ -565,16 +557,10 @@ async def test_router_success_callback_during_pre_header_increment_does_not_doub
@pytest.mark.asyncio
@pytest.mark.usefixtures("router_minute_pinned")
async def test_router_failed_redis_increment_keeps_local_usage_without_double_counting(
router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"
) -> None:
async def test_router_failed_redis_increment_keeps_local_usage_without_double_counting() -> None:
capture: Final = _RouterLoggingCapture("lit-3058-fail")
router: Final = _rpm_tpm_router("lit-3058-fail")
redis, _ = router_redis
started, released, unavailable = redis_write_gate
router.update_redis_cache(redis)
released.clear()
unavailable.set()
router.update_redis_cache(redis_cache_double())
metadata: Final[dict[str, object]] = {}
request: Final = asyncio.create_task(
router.acompletion(
@ -586,11 +572,8 @@ async def test_router_failed_redis_increment_keeps_local_usage_without_double_co
)
)
await asyncio.wait_for(started.wait(), timeout=EVENT_TIMEOUT_SECONDS)
await capture.next_event()
assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30
assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1
released.set()
response: Final = await request
assert response.usage.total_tokens == 30