mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(redis): move router cooldown and usage stores into Rust
Some checks failed
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
Some checks failed
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
This commit is contained in:
parent
bad662428d
commit
84373ee188
84 changed files with 5474 additions and 2767 deletions
27
litellm-rust/Cargo.lock
generated
27
litellm-rust/Cargo.lock
generated
|
|
@ -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]]
|
||||
|
|
@ -4358,6 +4360,8 @@ dependencies = [
|
|||
"litellm-cache-redis",
|
||||
"litellm-cache-response",
|
||||
"litellm-callbacks-legacy-python",
|
||||
"litellm-circuit-breaker",
|
||||
"litellm-config",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-host",
|
||||
|
|
@ -4372,6 +4376,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",
|
||||
|
|
@ -4433,6 +4440,7 @@ dependencies = [
|
|||
"socket2 0.6.5",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4465,6 +4473,25 @@ dependencies = [
|
|||
"veil",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-redis-store"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"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"
|
||||
|
|
|
|||
|
|
@ -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" }
|
||||
|
|
|
|||
|
|
@ -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
2
litellm-rust/crates/cache/Cargo.toml
vendored
|
|
@ -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
|
||||
|
|
|
|||
744
litellm-rust/crates/cache/src/dual.rs
vendored
744
litellm-rust/crates/cache/src/dual.rs
vendored
|
|
@ -1,519 +1,319 @@
|
|||
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::{Declared, Error as RedisError, RedisBatch, 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
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
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 round trip 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],
|
||||
decode: impl Fn(&[u8]) -> Result<V, E>,
|
||||
) -> Result<Vec<Option<V>>, Error> {
|
||||
let declared = self.declare_many(RedisClient::batch, keys)?;
|
||||
self.resolve_many(declared, decode).await
|
||||
}
|
||||
|
||||
/// The first half of `get_many`: reserves the due refreshes and declares their GETs on the
|
||||
/// batch `batch_for` picks for this tier's client, without sending anything.
|
||||
pub fn declare_many(
|
||||
&self,
|
||||
batch_for: impl FnOnce(&RedisClient) -> RedisBatch,
|
||||
keys: &[String],
|
||||
) -> Result<DeclaredRead, Error> {
|
||||
let fetch = match &self.remote {
|
||||
None => None,
|
||||
Some(remote) => {
|
||||
let reservation = self.reserve(keys, &self.read_memory(keys)?)?;
|
||||
(!reservation.keys.is_empty()).then(|| {
|
||||
let batch = batch_for(remote);
|
||||
let replies = reservation
|
||||
.keys
|
||||
.iter()
|
||||
.map(|key| batch.get::<Vec<u8>>(key))
|
||||
.collect();
|
||||
Fetch {
|
||||
reservation,
|
||||
replies,
|
||||
}
|
||||
})
|
||||
}
|
||||
};
|
||||
Ok(DeclaredRead {
|
||||
keys: keys.to_vec(),
|
||||
fetch,
|
||||
})
|
||||
}
|
||||
|
||||
/// The second half of `get_many`. Memory is read again here, so a value written to memory
|
||||
/// since the declaration wins over what Redis answered.
|
||||
pub async fn resolve_many<E: std::fmt::Display>(
|
||||
&self,
|
||||
declared: DeclaredRead,
|
||||
decode: impl Fn(&[u8]) -> Result<V, E>,
|
||||
) -> Result<Vec<Option<V>>, Error> {
|
||||
let DeclaredRead { keys, fetch } = declared;
|
||||
let Some(Fetch {
|
||||
reservation,
|
||||
replies,
|
||||
}) = fetch
|
||||
else {
|
||||
return self.read_memory(&keys);
|
||||
};
|
||||
let fetched = settle(replies).await;
|
||||
let local = self.read_memory(&keys)?;
|
||||
let fetched = match fetched {
|
||||
Ok(fetched) => fetched,
|
||||
Err(RedisError::BreakerOpen) => {
|
||||
self.release(reservation)?;
|
||||
return Ok(local);
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "Redis read failed, serving the memory tier");
|
||||
return Ok(local);
|
||||
}
|
||||
};
|
||||
let decoded = match fetched
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, entry)| (!matches!(entry, BatchEntry::Hit(_))).then_some(index))
|
||||
.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 = reservation
|
||||
.keys
|
||||
.into_iter()
|
||||
.zip(decoded)
|
||||
.collect::<HashMap<_, _>>();
|
||||
keys.iter()
|
||||
.zip(local)
|
||||
.map(|(key, local)| match local {
|
||||
Some(value) => Ok(Some(value)),
|
||||
None => self.backfill(key, read.get(key).cloned().flatten()),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn merge_batch(
|
||||
&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)
|
||||
pub fn delete_memory(&self, key: &str) -> Result<(), Error> {
|
||||
self.memory.delete_cache(key)
|
||||
}
|
||||
|
||||
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)
|
||||
pub fn flush_memory(&self) -> Result<(), Error> {
|
||||
self.memory.flush_cache()
|
||||
}
|
||||
|
||||
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();
|
||||
fn read_memory(&self, keys: &[String]) -> Result<Vec<Option<V>>, Error> {
|
||||
let context = ExactCacheContext::default();
|
||||
keys.iter()
|
||||
.map(|key| self.memory.get_cache(key, &context))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn backfill(&self, key: &str, value: Option<V>) -> Result<Option<V>, Error> {
|
||||
if let Some(value) = &value {
|
||||
let promotion_context = self.promotion_context(context);
|
||||
self.l1.set_cache(key, value.clone(), &promotion_context)?;
|
||||
self.set_memory(key, value.clone(), self.policy.backfill_ttl)?;
|
||||
}
|
||||
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
|
||||
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()
|
||||
.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),
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
|
||||
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)
|
||||
.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,
|
||||
}
|
||||
async fn settle(
|
||||
replies: Vec<Declared<Option<Vec<u8>>>>,
|
||||
) -> Result<Vec<Option<Vec<u8>>>, RedisError> {
|
||||
let mut fetched = Vec::with_capacity(replies.len());
|
||||
for reply in replies {
|
||||
fetched.push(reply.resolve().await?);
|
||||
}
|
||||
Ok(fetched)
|
||||
}
|
||||
|
||||
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),
|
||||
/// The GETs `declare_many` put on a batch, with the refreshes they reserved.
|
||||
pub struct DeclaredRead {
|
||||
keys: Vec<String>,
|
||||
fetch: Option<Fetch>,
|
||||
}
|
||||
|
||||
struct Fetch {
|
||||
reservation: Reservation,
|
||||
replies: Vec<Declared<Option<Vec<u8>>>>,
|
||||
}
|
||||
|
||||
struct Reservation {
|
||||
keys: Vec<String>,
|
||||
previous: Vec<Option<Duration>>,
|
||||
}
|
||||
|
||||
/// 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
4
litellm-rust/crates/cache/src/lib.rs
vendored
|
|
@ -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, DeclaredRead, DualCache, OnBreakerOpen, TierPolicy, system_clock};
|
||||
pub use error::{Error, ShortBatch};
|
||||
|
|
|
|||
839
litellm-rust/crates/cache/tests/dual.rs
vendored
839
litellm-rust/crates/cache/tests/dual.rs
vendored
|
|
@ -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())));
|
||||
}
|
||||
|
|
@ -1,9 +1,9 @@
|
|||
## Redis settings
|
||||
|
||||
- `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
|
||||
- That contract is `resolves_like_python_redis` in `tests/redis.rs` and its mirror `test_redis_settings_resolve_like_the_rust_resolver_expects` in `tests/unit/test_redis.py`. Both hold the same named cases with the expected view written next to the input, the Python side reading the effective redis-py 5.3.1 connection objects without connecting. A new resolution rule is a new case in both files, never a generated fixture. A case where the two sides expect different outcomes is a divergence and is listed below
|
||||
- That contract is `resolves_like_python_redis` in `tests/redis.rs` and its mirror `test_redis_settings_resolve_like_the_rust_resolver_expects` in `tests/unit/test_redis.py`. Both hold the same named cases with the expected view written next to the input, the Python side reading the effective redis-py 5.3.1 connection objects without connecting. A new resolution rule is a new case in both files, never a generated fixture
|
||||
- `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
|
||||
- Numbers are read with `litellm_python_compat::number`, so whitespace and digit underscores parse like Python's `int()` and `float()`. A value they 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 both contract tests with a different expected outcome on each side
|
||||
- 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 cases in both contract tests with a different expected outcome on each side. `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=...)`
|
||||
|
|
|
|||
|
|
@ -73,6 +73,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)]
|
||||
|
|
@ -132,7 +134,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") {
|
||||
|
|
@ -197,6 +198,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),
|
||||
}),
|
||||
},
|
||||
};
|
||||
|
|
@ -213,6 +215,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())),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -279,15 +284,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(),
|
||||
|
|
@ -339,17 +335,17 @@ fn coerced<N, 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()
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ use std::{
|
|||
collections::HashMap,
|
||||
fmt::Write,
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_config::{
|
||||
|
|
@ -136,14 +137,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 {
|
||||
|
|
@ -258,22 +276,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(
|
||||
|
|
@ -522,19 +540,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]
|
||||
|
|
@ -649,6 +671,8 @@ 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,
|
||||
|
|
@ -668,24 +692,11 @@ fn expose(value: Option<&litellm_secret_value::SecretValue>) -> Option<&str> {
|
|||
value.map(litellm_secret_value::SecretValue::expose)
|
||||
}
|
||||
|
||||
fn expected_view(fields: serde_json::Value) -> serde_json::Value {
|
||||
let serde_json::Value::Object(fields) = fields else {
|
||||
unreachable!("expected views are objects")
|
||||
};
|
||||
let defaults = json!({
|
||||
"auth": {"kind": "static", "username": null, "password": null},
|
||||
"tls": null,
|
||||
"socket_timeout": 5.0,
|
||||
"socket_connect_timeout": 5.0,
|
||||
});
|
||||
let serde_json::Value::Object(defaults) = defaults else {
|
||||
unreachable!("the defaults are an object")
|
||||
};
|
||||
serde_json::Value::Object(defaults.into_iter().chain(fields).collect())
|
||||
}
|
||||
|
||||
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,
|
||||
|
|
@ -740,6 +751,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")
|
||||
|
|
@ -747,51 +762,68 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
|
|||
serde_json::Value::Object(target.into_iter().chain(rest).collect())
|
||||
}
|
||||
|
||||
fn expected_view(fields: serde_json::Value) -> serde_json::Value {
|
||||
let serde_json::Value::Object(fields) = fields else {
|
||||
unreachable!("expected views are objects")
|
||||
};
|
||||
let defaults = json!({
|
||||
"auth": {"kind": "static", "username": null, "password": null},
|
||||
"tls": null,
|
||||
"socket_timeout": 5.0,
|
||||
"socket_connect_timeout": 5.0,
|
||||
"health_check_interval": null,
|
||||
});
|
||||
let serde_json::Value::Object(defaults) = defaults else {
|
||||
unreachable!("the defaults are an object")
|
||||
};
|
||||
serde_json::Value::Object(defaults.into_iter().chain(fields).collect())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::host_defaults(
|
||||
json!({"host": "redis"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::empty_cluster_nodes_use_the_environment(
|
||||
json!({"startup_nodes": []}),
|
||||
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": 7000}]"#)],
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))),
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::empty_sentinel_nodes_use_the_environment(
|
||||
json!({"sentinel_nodes": [], "service_name": "primary"}),
|
||||
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#)],
|
||||
Ok(expected_view(json!({"db": 0, "mode": "sentinel", "nodes": ["sentinel-a:26379"], "sentinel_password": null, "service_name": "primary"}))),
|
||||
Ok(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))),
|
||||
)]
|
||||
#[case::empty_service_name_uses_the_environment(
|
||||
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}),
|
||||
&[("REDIS_SERVICE_NAME", "primary")],
|
||||
Ok(expected_view(json!({"db": 0, "mode": "sentinel", "nodes": ["sentinel-a:26379"], "sentinel_password": null, "service_name": "primary"}))),
|
||||
Ok(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))),
|
||||
)]
|
||||
#[case::plaintext_ignores_tls_requirements(
|
||||
json!({"host": "redis", "ssl_cert_reqs": "strict"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::ssl_ca_path_is_unsupported(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ca_path": "/ca"}),
|
||||
&[],
|
||||
Err("Redis host does not yet support configured setting ssl_ca_path"),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/ca", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::host_with_every_basic_setting(
|
||||
json!({"db": 2, "host": "redis", "password": "pass", "port": 6380, "socket_timeout": 1.5, "ssl": true, "username": "user"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"auth": {"kind": "static", "password": "pass", "username": "user"}, "db": 2, "host": "redis", "mode": "host", "port": 6380, "socket_connect_timeout": 1.5, "socket_timeout": 1.5, "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 2, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": "pass"}, "socket_timeout": 1.5, "socket_connect_timeout": 1.5}))),
|
||||
)]
|
||||
#[case::environment_host_port_and_password(
|
||||
json!({}),
|
||||
&[("REDIS_HOST", "environment"), ("REDIS_PASSWORD", "environment-secret"), ("REDIS_PORT", "6390"), ("REDIS_SSL", "True")],
|
||||
Ok(expected_view(json!({"auth": {"kind": "static", "password": "environment-secret", "username": null}, "db": 0, "host": "environment", "mode": "host", "port": 6390, "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "environment", "port": 6390, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": null, "password": "environment-secret"}}))),
|
||||
)]
|
||||
#[case::string_numbers_are_coerced(
|
||||
json!({"db": "3", "host": "redis", "port": "6380", "socket_timeout": "2.5"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 3, "host": "redis", "mode": "host", "port": 6380, "socket_connect_timeout": 2.5, "socket_timeout": 2.5}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 3, "socket_timeout": 2.5, "socket_connect_timeout": 2.5}))),
|
||||
)]
|
||||
#[case::padded_and_underscored_numbers_parse_like_python_int_and_float(
|
||||
json!({"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}),
|
||||
|
|
@ -806,7 +838,7 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
|
|||
#[case::socket_timeout_from_the_environment(
|
||||
json!({"host": "redis"}),
|
||||
&[("REDIS_SOCKET_TIMEOUT", "0.75")],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "socket_connect_timeout": 0.75, "socket_timeout": 0.75}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 0.75, "socket_connect_timeout": 0.75}))),
|
||||
)]
|
||||
#[case::environment_url(
|
||||
json!({}),
|
||||
|
|
@ -821,42 +853,42 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
|
|||
#[case::an_explicit_host_drops_the_environment_url(
|
||||
json!({"host": "configured"}),
|
||||
&[("REDIS_URL", "redis://environment:6379")],
|
||||
Ok(expected_view(json!({"db": 0, "host": "configured", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "configured", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::a_url_drops_host_db_and_password_and_keeps_the_username(
|
||||
json!({"db": 4, "host": "dropped", "password": "dropped", "url": "rediss://cache:6380", "username": "user"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"auth": {"kind": "static", "password": null, "username": "user"}, "mode": "url", "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}, "url": "rediss://cache:6380"}))),
|
||||
Ok(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": null}}))),
|
||||
)]
|
||||
#[case::a_url_ignores_tls_settings_outside_the_url(
|
||||
json!({"ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true, "url": "rediss://cache:6380"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "url", "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}, "url": "rediss://cache:6380"}))),
|
||||
Ok(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::startup_nodes_beat_a_url(
|
||||
json!({"password": "cluster-secret", "startup_nodes": [{"host": "node-a", "port": 7000}], "url": "redis://single:6379"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"auth": {"kind": "static", "password": "cluster-secret", "username": null}, "mode": "cluster", "nodes": ["node-a:7000"]}))),
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "auth": {"kind": "static", "username": null, "password": "cluster-secret"}, "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::environment_cluster_nodes(
|
||||
json!({}),
|
||||
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]"#)],
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"]}))),
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::cluster_tls_options(
|
||||
json!({"ssl": "true", "ssl_ca_certs": "/ca.pem", "ssl_check_hostname": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "tls": {"ca_certs": "/ca.pem", "certfile": null, "check_hostname": true, "keyfile": null, "verify_certificate": true}}))),
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": true}, "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::sentinel_from_settings(
|
||||
json!({"password": "data-secret", "sentinel_nodes": [["sentinel-a", 26379], ["sentinel-b", 26380]], "sentinel_password": "sentinel-secret", "service_name": "primary"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"auth": {"kind": "static", "password": "data-secret", "username": null}, "db": 0, "mode": "sentinel", "nodes": ["sentinel-a:26379", "sentinel-b:26380"], "sentinel_password": "sentinel-secret", "service_name": "primary"}))),
|
||||
Ok(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379", "sentinel-b:26380"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0, "auth": {"kind": "static", "username": null, "password": "data-secret"}}))),
|
||||
)]
|
||||
#[case::sentinel_from_the_environment(
|
||||
json!({}),
|
||||
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#), ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), ("REDIS_SERVICE_NAME", "primary")],
|
||||
Ok(expected_view(json!({"db": 0, "mode": "sentinel", "nodes": ["sentinel-a:26379"], "sentinel_password": "sentinel-secret", "service_name": "primary"}))),
|
||||
Ok(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0}))),
|
||||
)]
|
||||
#[case::a_host_breaks_sentinel(
|
||||
json!({"host": "ignored", "sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}),
|
||||
|
|
@ -871,47 +903,47 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
|
|||
#[case::an_explicit_credential_provider_beats_cloud_auth(
|
||||
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
|
||||
&[("REDIS_AZURE_AD_TOKEN", "true")],
|
||||
Ok(expected_view(json!({"auth": {"kind": "host"}, "db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "host"}}))),
|
||||
)]
|
||||
#[case::gcp_beats_azure_and_aws(
|
||||
json!({"aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
|
||||
&[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com")],
|
||||
Ok(expected_view(json!({"auth": {"kind": "gcp_iam", "service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}, "db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "gcp_iam", "service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}}))),
|
||||
)]
|
||||
#[case::gcp_ca_certs_replace_the_configured_ones_under_ssl(
|
||||
json!({"gcp_service_account": "account", "gcp_ssl_ca_certs": "/gcp-ca.pem", "host": "redis", "ssl": true, "ssl_ca_certs": "/configured-ca.pem"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"auth": {"kind": "gcp_iam", "service_account": "account"}, "db": 0, "host": "redis", "mode": "host", "port": 6379, "tls": {"ca_certs": "/gcp-ca.pem", "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/gcp-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "gcp_iam", "service_account": "account"}}))),
|
||||
)]
|
||||
#[case::azure_beats_aws_and_takes_the_environment_username(
|
||||
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
|
||||
&[("REDIS_USERNAME", "environment-user")],
|
||||
Ok(expected_view(json!({"auth": {"kind": "azure_ad", "username": "environment-user"}, "db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": "environment-user"}}))),
|
||||
)]
|
||||
#[case::azure_without_a_username_sends_only_the_token(
|
||||
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"auth": {"kind": "azure_ad", "username": null}, "db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))),
|
||||
)]
|
||||
#[case::azure_needs_the_literal_true(
|
||||
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": "1", "azure_tenant_id": "tenant", "host": "redis"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::aws_iam_over_tls_with_the_default_region(
|
||||
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true}),
|
||||
&[("AWS_DEFAULT_REGION", "us-east-1")],
|
||||
Ok(expected_view(json!({"auth": {"cache_name": "cache-1", "kind": "aws_iam", "region": "us-east-1", "serverless": false, "user_name": "iam-user"}, "db": 0, "host": "cache-1", "mode": "host", "port": 6379, "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "us-east-1", "serverless": false}}))),
|
||||
)]
|
||||
#[case::aws_region_beats_the_default_region(
|
||||
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}),
|
||||
&[("AWS_DEFAULT_REGION", "us-east-1"), ("AWS_REGION", "eu-central-1")],
|
||||
Ok(expected_view(json!({"auth": {"cache_name": "cache-1", "kind": "aws_iam", "region": "eu-central-1", "serverless": false, "user_name": "iam-user"}, "db": 0, "host": "cache-1", "mode": "host", "port": 6379, "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-central-1", "serverless": false}}))),
|
||||
)]
|
||||
#[case::aws_iam_serverless_over_a_rediss_url_drops_its_userinfo(
|
||||
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "eu-west-1", "aws_iam_serverless": true, "aws_iam_user_name": "iam-user", "url": "rediss://user:pass@cache-1:6379"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"auth": {"cache_name": "cache-1", "kind": "aws_iam", "region": "eu-west-1", "serverless": true, "user_name": "iam-user"}, "mode": "url", "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}, "url": "rediss://cache-1:6379"}))),
|
||||
Ok(expected_view(json!({"mode": "url", "url": "rediss://cache-1:6379", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-west-1", "serverless": true}}))),
|
||||
)]
|
||||
#[case::aws_iam_without_tls_fails(
|
||||
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "us-east-1", "aws_iam_user_name": "iam-user", "host": "cache-1"}),
|
||||
|
|
@ -926,22 +958,22 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
|
|||
#[case::tls_options_from_settings(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "tls": {"ca_certs": "/ca.pem", "certfile": null, "check_hostname": true, "keyfile": null, "verify_certificate": false}}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))),
|
||||
)]
|
||||
#[case::tls_options_from_the_environment(
|
||||
json!({"host": "redis"}),
|
||||
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "tls": {"ca_certs": "/env-ca.pem", "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/env-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::connect_timeout_is_separate_from_the_socket_timeout(
|
||||
json!({"host": "redis", "socket_connect_timeout": 0.25, "socket_timeout": 2.5}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "socket_connect_timeout": 0.25, "socket_timeout": 2.5}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 2.5, "socket_connect_timeout": 0.25}))),
|
||||
)]
|
||||
#[case::environment_references_resolve(
|
||||
json!({"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}),
|
||||
&[("COORD_HOST", "referenced"), ("COORD_PASSWORD", "referenced-secret")],
|
||||
Ok(expected_view(json!({"auth": {"kind": "static", "password": "referenced-secret", "username": null}, "db": 0, "host": "referenced", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "referenced", "port": 6379, "db": 0, "auth": {"kind": "static", "username": null, "password": "referenced-secret"}}))),
|
||||
)]
|
||||
#[case::no_target_fails(
|
||||
json!({"port": 6379}),
|
||||
|
|
@ -951,22 +983,22 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
|
|||
#[case::client_certificates(
|
||||
json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "tls": {"ca_certs": null, "certfile": "/client.pem", "check_hostname": false, "keyfile": "/client.key", "verify_certificate": true}}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": "/client.key", "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::a_rediss_url_takes_tls_options_from_its_query(
|
||||
json!({"url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "url", "tls": {"ca_certs": "/ca.pem", "certfile": null, "check_hostname": true, "keyfile": null, "verify_certificate": false}, "url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}))),
|
||||
Ok(expected_view(json!({"mode": "url", "url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes", "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))),
|
||||
)]
|
||||
#[case::a_url_query_setting_without_a_redis_rs_equivalent(
|
||||
json!({"url": "redis://cache:6379/0?socket_timeout=2"}),
|
||||
&[],
|
||||
Err("Redis host does not yet support configured setting url ?socket_timeout"),
|
||||
Ok(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2"}))),
|
||||
)]
|
||||
#[case::tls_settings_outside_a_url_never_fail_it(
|
||||
json!({"db": "99999999999999999999", "ssl_cert_reqs": "strict", "url": "rediss://cache:6380"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "url", "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}, "url": "rediss://cache:6380"}))),
|
||||
Ok(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::db_is_not_read_for_a_cluster(
|
||||
json!({"db": 99_999_999_999_i64, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
|
|
@ -981,7 +1013,7 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
|
|||
#[case::cluster_db_false(
|
||||
json!({"db": false, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))),
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::cluster_db_fraction(
|
||||
json!({"db": 0.5, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
|
|
@ -991,7 +1023,7 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
|
|||
#[case::cluster_db_float_zero(
|
||||
json!({"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))),
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::cluster_db_nonempty_list(
|
||||
json!({"db": [1], "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
|
|
@ -1001,52 +1033,87 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value {
|
|||
#[case::cluster_db_empty_list(
|
||||
json!({"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))),
|
||||
Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::ssl_ciphers(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ciphers": "HIGH"}),
|
||||
&[],
|
||||
Err("Redis host does not yet support configured setting ssl_ciphers"),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::retry_on_timeout(
|
||||
json!({"host": "redis", "retry_on_timeout": true}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::malformed_port(
|
||||
json!({"host": "redis", "port": "abc"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::malformed_db(
|
||||
json!({"db": "first", "host": "redis"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::socket_timeout_not_a_number(
|
||||
json!({"host": "redis", "socket_timeout": "slow"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::malformed_port_from_the_environment(
|
||||
json!({}),
|
||||
&[("REDIS_HOST", "redis"), ("REDIS_PORT", "abc")],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::protocol_3(
|
||||
json!({"host": "redis", "protocol": 3}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::retry_on_error_from_the_environment(
|
||||
json!({"host": "redis"}),
|
||||
&[("REDIS_RETRY_ON_ERROR", "true")],
|
||||
Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))),
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::unix_socket_path(
|
||||
json!({"host": "redis", "unix_socket_path": "/tmp/redis.sock"}),
|
||||
&[],
|
||||
Err("Redis host does not yet support configured setting unix_socket_path"),
|
||||
Ok(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))),
|
||||
)]
|
||||
#[case::unix_socket_path_keeps_tls_settings_out_of_the_target(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))),
|
||||
)]
|
||||
#[case::ssl_ca_data(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::ssl_ca_path_from_the_environment(
|
||||
json!({"host": "redis"}),
|
||||
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_PATH", "/etc/ssl/private-cas")],
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/etc/ssl/private-cas", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::health_check_interval_is_whole_seconds(
|
||||
json!({"health_check_interval": "7", "host": "redis"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "health_check_interval": 7.0}))),
|
||||
)]
|
||||
#[case::health_check_interval_zero_is_off(
|
||||
json!({"health_check_interval": 0, "host": "redis"}),
|
||||
&[],
|
||||
Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::a_cluster_rejects_a_nonzero_db(
|
||||
json!({"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
Err("db must be 0 in cluster mode"),
|
||||
)]
|
||||
#[case::unix_socket_path_without_a_host_fails_like_python(
|
||||
json!({"unix_socket_path": "/tmp/redis.sock"}),
|
||||
&[],
|
||||
Err("either host or url must be specified for redis"),
|
||||
)]
|
||||
fn resolves_like_python_redis(
|
||||
#[case] kwargs: serde_json::Value,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,14 @@ 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::request::NativeRequestRedis;
|
||||
#[pymodule_export]
|
||||
use crate::redis::router_cooldowns::{NativeDeclaredCooldowns, 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 +122,11 @@ mod tests {
|
|||
"aresponses",
|
||||
"ResponsesWebSocketConnection",
|
||||
"NativeDiagnosticProcessor",
|
||||
"NativeRequestRedis",
|
||||
"NativeDeclaredCooldowns",
|
||||
"NativeRouterCooldowns",
|
||||
"NativeRouterUsage",
|
||||
"native_redis_breaker_states",
|
||||
"NativeClickHouseSpendConfig",
|
||||
"NativeClickHouseSpendStorage",
|
||||
"NativeDatabase",
|
||||
|
|
|
|||
|
|
@ -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
145
litellm-rust/crates/python-bridge/src/redis/client.rs
Normal 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
251
litellm-rust/crates/python-bridge/src/redis/events.rs
Normal 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)]"
|
||||
);
|
||||
}
|
||||
}
|
||||
61
litellm-rust/crates/python-bridge/src/redis/mod.rs
Normal file
61
litellm-rust/crates/python-bridge/src/redis/mod.rs
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
mod client;
|
||||
mod events;
|
||||
pub(crate) use client::native_redis_breaker_states;
|
||||
mod policy;
|
||||
pub(crate) mod request;
|
||||
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
130
litellm-rust/crates/python-bridge/src/redis/policy.rs
Normal 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}");
|
||||
});
|
||||
}
|
||||
}
|
||||
40
litellm-rust/crates/python-bridge/src/redis/request.rs
Normal file
40
litellm-rust/crates/python-bridge/src/redis/request.rs
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_store::RequestRedis;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use super::call;
|
||||
use crate::execution::run_async_value;
|
||||
|
||||
/// One request's native Redis batches. Python builds one per request scope and passes it to each
|
||||
/// request-aware store call.
|
||||
#[pyclass(frozen, name = "NativeRequestRedis")]
|
||||
pub(crate) struct NativeRequestRedis(pub(crate) RequestRedis);
|
||||
|
||||
#[pymethods]
|
||||
impl NativeRequestRedis {
|
||||
#[new]
|
||||
fn new(post_call_deadline_seconds: f64) -> Self {
|
||||
Self(RequestRedis::new(Duration::from_secs_f64(
|
||||
post_call_deadline_seconds.max(0.0),
|
||||
)))
|
||||
}
|
||||
|
||||
fn flush_pre<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request = self.0.clone();
|
||||
let scope = call(py, "request_redis_batch")?;
|
||||
run_async_value(py, async move {
|
||||
scope.scope(request.flush_pre()).await;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn flush_post<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request = self.0.clone();
|
||||
let scope = call(py, "post_call_redis_batch")?;
|
||||
run_async_value(py, async move {
|
||||
scope.scope(request.flush_post()).await;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
244
litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs
Normal file
244
litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs
Normal file
|
|
@ -0,0 +1,244 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex, PoisonError},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_redis_store::router_cooldowns::{Cooldown, DeclaredCooldowns, RouterCooldowns};
|
||||
use pyo3::{
|
||||
exceptions::PyRuntimeError,
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use super::{
|
||||
PyNumber, call,
|
||||
client::{PythonRedisSettings, shared_client},
|
||||
memory_error,
|
||||
policy::router_tier_policies,
|
||||
request::NativeRequestRedis,
|
||||
};
|
||||
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 declare_active(
|
||||
&self,
|
||||
request: &NativeRequestRedis,
|
||||
model_ids: Vec<String>,
|
||||
) -> PyResult<NativeDeclaredCooldowns> {
|
||||
let declared = self
|
||||
.0
|
||||
.declare_active(&request.0, &model_ids)
|
||||
.map_err(memory_error)?;
|
||||
Ok(NativeDeclaredCooldowns {
|
||||
store: self.0.clone(),
|
||||
declared: Mutex::new(Some(declared)),
|
||||
})
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
/// An `active` read declared on a request's pre-call batch. `resolve` answers it once.
|
||||
#[pyclass(frozen, name = "NativeDeclaredCooldowns")]
|
||||
pub(crate) struct NativeDeclaredCooldowns {
|
||||
store: Arc<RouterCooldowns>,
|
||||
declared: Mutex<Option<DeclaredCooldowns>>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeDeclaredCooldowns {
|
||||
fn resolve<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let declared = self
|
||||
.declared
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.take()
|
||||
.ok_or_else(|| PyRuntimeError::new_err("this cooldown read was already resolved"))?;
|
||||
let store = self.store.clone();
|
||||
let scope = call(py, "async_batch_get_cache")?;
|
||||
run_async_value(py, async move {
|
||||
let active = scope
|
||||
.scope(store.resolve_active(declared))
|
||||
.await
|
||||
.map_err(memory_error)?;
|
||||
Python::attach(|py| {
|
||||
Ok(active_to_py(py, active)?
|
||||
.into_iter()
|
||||
.map(Bound::unbind)
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
160
litellm-rust/crates/python-bridge/src/redis/router_usage.rs
Normal file
160
litellm-rust/crates/python-bridge/src/redis/router_usage.rs
Normal file
|
|
@ -0,0 +1,160 @@
|
|||
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,
|
||||
request::NativeRequestRedis,
|
||||
};
|
||||
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)
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (deployment_id, model, minute, tokens, requests, *, request = None))]
|
||||
#[expect(
|
||||
clippy::too_many_arguments,
|
||||
reason = "preserves the Python increment_deployment signature"
|
||||
)]
|
||||
fn increment_deployment<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
deployment_id: String,
|
||||
model: Option<String>,
|
||||
minute: String,
|
||||
tokens: PyNumber,
|
||||
requests: i64,
|
||||
request: Option<PyRef<'_, NativeRequestRedis>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let request = request.map(|request| request.0.clone());
|
||||
let scope = call(py, "async_increment_pipeline")?;
|
||||
run_async_value(py, async move {
|
||||
let model = model.as_deref();
|
||||
match request {
|
||||
Some(request) => store.increment_deployment_post_call(
|
||||
&request,
|
||||
&deployment_id,
|
||||
model,
|
||||
&minute,
|
||||
tokens.0,
|
||||
requests,
|
||||
),
|
||||
None => {
|
||||
scope
|
||||
.scope(store.increment_deployment(
|
||||
&deployment_id,
|
||||
model,
|
||||
&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)
|
||||
}
|
||||
}
|
||||
|
|
@ -42,6 +42,7 @@
|
|||
- `RedisBatch` declares ops now and sends them later as one `Pipeline`, so it inherits every pipeline rule above. It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local
|
||||
- Each op returns a `Declared<T>` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once
|
||||
- An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis
|
||||
- `settled` waits for whichever flush carries the op and never sends one itself, for follow-ups that must not force a flush. `RedisClient::same_client` tells clones of one client apart from a second connection, so a host can keep one batch per client
|
||||
- Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes
|
||||
- Tests that only need the public API live in `tests/<module>.rs`, with client builders in `tests/support`. Inline `#[cfg(test)]` modules are only for private items (`error.rs`, `namespace.rs`)
|
||||
- Observe connections from the server side (`RedisDeployment::clients`, `cli_on_every_node`), never by reaching into `RedisClient`'s connection
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -135,6 +135,10 @@ impl<T> Declared<T> {
|
|||
Err(TryRecvError::Empty) => {}
|
||||
}
|
||||
self.batch.flush().await;
|
||||
self.settled().await
|
||||
}
|
||||
|
||||
pub async fn settled(self) -> Result<T, Error> {
|
||||
let reply = self.reply.await.unwrap_or_else(|_| Err(abandoned()));
|
||||
reply.and_then(self.decode)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -141,6 +150,10 @@ impl RedisClient {
|
|||
Pipeline::new(self.namespace.clone())
|
||||
}
|
||||
|
||||
pub fn same_client(&self, other: &Self) -> bool {
|
||||
Arc::ptr_eq(&self.last_used, &other.last_used)
|
||||
}
|
||||
|
||||
pub fn batch(&self) -> RedisBatch {
|
||||
RedisBatch::new(self.clone())
|
||||
}
|
||||
|
|
@ -163,6 +176,7 @@ impl RedisClient {
|
|||
.collect();
|
||||
}
|
||||
};
|
||||
self.check_health().await;
|
||||
let results = pipeline.run(&self.connection).await;
|
||||
let transport = results.iter().find_map(|result| {
|
||||
result
|
||||
|
|
@ -220,6 +234,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()));
|
||||
|
|
@ -233,6 +248,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,
|
||||
|
|
@ -242,6 +261,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()
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ fn observed(server: &RedisDeployment) -> (RedisClient, Arc<RoundTrips>) {
|
|||
timeouts: Timeouts::default(),
|
||||
breaker: None,
|
||||
events: Some(trips.clone()),
|
||||
health_check: None,
|
||||
})
|
||||
.unwrap();
|
||||
(client, trips)
|
||||
|
|
@ -145,3 +146,50 @@ async fn an_op_whose_flush_was_cancelled_resolves_as_transport(
|
|||
assert!(cancelled.is_err());
|
||||
assert!(matches!(counter.resolve().await, Err(Error::Transport(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_settled_op_waits_for_someone_else_to_flush_and_sends_nothing_itself(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = observed(&server);
|
||||
let batch = client.batch();
|
||||
let counter = tokio::spawn(batch.incr_by("counter", 4, KEEP).settled());
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
assert!(!counter.is_finished());
|
||||
assert!(trips.seen().is_empty());
|
||||
|
||||
batch.flush().await;
|
||||
|
||||
assert_eq!(counter.await.unwrap().unwrap(), 4);
|
||||
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 1 }]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_settled_op_whose_flush_was_cancelled_resolves_as_transport() {
|
||||
let server = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let client = connect(&server, "", Duration::from_secs(5));
|
||||
let batch = client.batch();
|
||||
let counter = batch.incr_by("counter", 1, KEEP);
|
||||
|
||||
server.freeze().await.unwrap();
|
||||
let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await;
|
||||
server.thaw().await.unwrap();
|
||||
|
||||
assert!(cancelled.is_err());
|
||||
assert!(matches!(counter.settled().await, Err(Error::Transport(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn clones_are_the_same_client_and_a_second_connection_is_not() {
|
||||
let server = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
let twin = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
|
||||
assert!(client.same_client(&client.clone()));
|
||||
assert!(!client.same_client(&twin));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -247,9 +247,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(
|
||||
|
|
@ -623,6 +625,7 @@ mod tests {
|
|||
timeouts: Timeouts::default(),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
})
|
||||
.unwrap();
|
||||
client.set("k", "1", Duration::from_secs(60)).await.unwrap();
|
||||
|
|
@ -658,4 +661,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")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
7
litellm-rust/crates/redis-store/AGENTS.md
Normal file
7
litellm-rust/crates/redis-store/AGENTS.md
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
- 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
|
||||
- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. A family exposes a request-aware variant that declares on it (`declare_active`, `increment_deployment_post_call`) next to its direct method, which runs the same declaration on a batch of its own
|
||||
- Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages
|
||||
24
litellm-rust/crates/redis-store/Cargo.toml
Normal file
24
litellm-rust/crates/redis-store/Cargo.toml
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
[package]
|
||||
name = "litellm-redis-store"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
futures-util.workspace = true
|
||||
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
|
||||
tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] }
|
||||
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
20
litellm-rust/crates/redis-store/src/codec.rs
Normal 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),
|
||||
}
|
||||
}
|
||||
9
litellm-rust/crates/redis-store/src/lib.rs
Normal file
9
litellm-rust/crates/redis-store/src/lib.rs
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
mod codec;
|
||||
mod number;
|
||||
mod request;
|
||||
pub mod router_cooldowns;
|
||||
pub mod router_usage;
|
||||
|
||||
pub use codec::DecodeError;
|
||||
pub use number::Number;
|
||||
pub use request::RequestRedis;
|
||||
68
litellm-rust/crates/redis-store/src/number.rs
Normal file
68
litellm-rust/crates/redis-store/src/number.rs
Normal 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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
99
litellm-rust/crates/redis-store/src/request.rs
Normal file
99
litellm-rust/crates/redis-store/src/request.rs
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex, MutexGuard, PoisonError},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use futures_util::future::join_all;
|
||||
use litellm_redis_client::{RedisBatch, RedisClient};
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
/// One request's Redis batches: a pre-call batch for the reads it needs before answering and a
|
||||
/// post-call batch for the writes nobody waits on, one of each per Redis client. The host that owns
|
||||
/// the request builds one and passes it down; it knows nothing else about the request.
|
||||
#[derive(Clone)]
|
||||
pub struct RequestRedis(Arc<Shared>);
|
||||
|
||||
struct Shared {
|
||||
post_call_deadline: Duration,
|
||||
batches: Mutex<Vec<ClientBatches>>,
|
||||
deadline: Mutex<Option<oneshot::Sender<()>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ClientBatches {
|
||||
client: RedisClient,
|
||||
pre: RedisBatch,
|
||||
post: RedisBatch,
|
||||
}
|
||||
|
||||
impl RequestRedis {
|
||||
pub fn new(post_call_deadline: Duration) -> Self {
|
||||
Self(Arc::new(Shared {
|
||||
post_call_deadline,
|
||||
batches: Mutex::new(Vec::new()),
|
||||
deadline: Mutex::new(None),
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn pre(&self, client: &RedisClient) -> RedisBatch {
|
||||
self.batches_for(client).pre
|
||||
}
|
||||
|
||||
/// The post-call batch, which goes out on `flush_post` or `post_call_deadline` after the first
|
||||
/// declaration, whichever comes first. Must be called inside a Tokio runtime.
|
||||
pub fn post(&self, client: &RedisClient) -> RedisBatch {
|
||||
self.arm_deadline();
|
||||
self.batches_for(client).post
|
||||
}
|
||||
|
||||
pub async fn flush_pre(&self) {
|
||||
join_all(self.batches().iter().map(|batches| batches.pre.flush())).await;
|
||||
}
|
||||
|
||||
/// Sends every post-call op declared so far; the shutdown drain calls it too.
|
||||
pub async fn flush_post(&self) {
|
||||
lock(&self.0.deadline).take();
|
||||
join_all(self.batches().iter().map(|batches| batches.post.flush())).await;
|
||||
}
|
||||
|
||||
fn batches(&self) -> Vec<ClientBatches> {
|
||||
lock(&self.0.batches).clone()
|
||||
}
|
||||
|
||||
fn batches_for(&self, client: &RedisClient) -> ClientBatches {
|
||||
let mut batches = lock(&self.0.batches);
|
||||
if let Some(existing) = batches
|
||||
.iter()
|
||||
.find(|batches| batches.client.same_client(client))
|
||||
{
|
||||
return existing.clone();
|
||||
}
|
||||
let created = ClientBatches {
|
||||
client: client.clone(),
|
||||
pre: client.batch(),
|
||||
post: client.batch(),
|
||||
};
|
||||
batches.push(created.clone());
|
||||
created
|
||||
}
|
||||
|
||||
fn arm_deadline(&self) {
|
||||
let mut deadline = lock(&self.0.deadline);
|
||||
if deadline.is_some() {
|
||||
return;
|
||||
}
|
||||
let (disarm, disarmed) = oneshot::channel();
|
||||
*deadline = Some(disarm);
|
||||
let request = self.clone();
|
||||
tokio::spawn(async move {
|
||||
tokio::select! {
|
||||
() = tokio::time::sleep(request.0.post_call_deadline) => request.flush_post().await,
|
||||
_ = disarmed => {}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
|
||||
mutex.lock().unwrap_or_else(PoisonError::into_inner)
|
||||
}
|
||||
282
litellm-rust/crates/redis-store/src/router_cooldowns.rs
Normal file
282
litellm-rust/crates/redis-store/src/router_cooldowns.rs
Normal file
|
|
@ -0,0 +1,282 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
Clock, DeclaredRead, 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, RedisBatch, RedisClient};
|
||||
|
||||
use crate::{DecodeError, Number, RequestRedis, 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 keys(model_ids: &[String]) -> Vec<String> {
|
||||
model_ids.iter().map(|model_id| key(model_id)).collect()
|
||||
}
|
||||
|
||||
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())
|
||||
}
|
||||
|
||||
/// An `active` read declared on a request batch, answered by `resolve_active`.
|
||||
pub struct DeclaredCooldowns {
|
||||
model_ids: Vec<String>,
|
||||
read: DeclaredRead,
|
||||
}
|
||||
|
||||
/// 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 declared = self.declare(RedisClient::batch, model_ids)?;
|
||||
self.resolve_active(declared).await
|
||||
}
|
||||
|
||||
/// `active`, with its Redis read declared on the request's pre-call batch.
|
||||
pub fn declare_active(
|
||||
&self,
|
||||
request: &RequestRedis,
|
||||
model_ids: &[String],
|
||||
) -> Result<DeclaredCooldowns, MemoryError> {
|
||||
self.declare(|client| request.pre(client), model_ids)
|
||||
}
|
||||
|
||||
pub async fn resolve_active(
|
||||
&self,
|
||||
declared: DeclaredCooldowns,
|
||||
) -> Result<Vec<(String, Cooldown)>, MemoryError> {
|
||||
let DeclaredCooldowns { model_ids, read } = declared;
|
||||
let read = self.cooldowns.resolve_many(read, Cooldown::decode).await?;
|
||||
let now = self.cooldowns.now();
|
||||
let corrected = model_ids
|
||||
.into_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_some((model_id, 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> {
|
||||
self.cooldowns
|
||||
.get_many(&keys(model_ids), Cooldown::decode)
|
||||
.await
|
||||
}
|
||||
|
||||
fn declare(
|
||||
&self,
|
||||
batch_for: impl FnOnce(&RedisClient) -> RedisBatch,
|
||||
model_ids: &[String],
|
||||
) -> Result<DeclaredCooldowns, MemoryError> {
|
||||
Ok(DeclaredCooldowns {
|
||||
model_ids: model_ids.to_vec(),
|
||||
read: self.cooldowns.declare_many(batch_for, &keys(model_ids))?,
|
||||
})
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
261
litellm-rust/crates/redis-store/src/router_usage.rs
Normal file
261
litellm-rust/crates/redis-store/src/router_usage.rs
Normal file
|
|
@ -0,0 +1,261 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
Clock, DualCache, Error as MemoryError, OnBreakerOpen, TierPolicy, system_clock,
|
||||
};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient};
|
||||
|
||||
use crate::{DecodeError, Number, RequestRedis, 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 increment = Increment::new(deployment_id, model, minute, tokens, requests);
|
||||
let Some(declared) = self.declare(RedisClient::batch, increment)? else {
|
||||
return Ok(());
|
||||
};
|
||||
let totals = totals(declared.totals, Declared::resolve).await;
|
||||
self.mirror(&declared.keys, declared.local, totals)
|
||||
}
|
||||
|
||||
/// `increment_deployment` on the request's post-call batch. Memory counts at once; the Redis
|
||||
/// totals reach memory whenever that batch goes out, from a task this spawns, so it must be
|
||||
/// called inside a Tokio runtime.
|
||||
pub fn increment_deployment_post_call(
|
||||
self: &Arc<Self>,
|
||||
request: &RequestRedis,
|
||||
deployment_id: &str,
|
||||
model: Option<&str>,
|
||||
minute: &str,
|
||||
tokens: Number,
|
||||
requests: i64,
|
||||
) -> Result<(), MemoryError> {
|
||||
let increment = Increment::new(deployment_id, model, minute, tokens, requests);
|
||||
let Some(declared) = self.declare(|client| request.post(client), increment)? else {
|
||||
return Ok(());
|
||||
};
|
||||
let store = Arc::clone(self);
|
||||
tokio::spawn(async move {
|
||||
let totals = totals(declared.totals, Declared::settled).await;
|
||||
if let Err(error) = store.mirror(&declared.keys, declared.local, totals) {
|
||||
tracing::warn!(%error, "mirroring Redis usage totals into memory failed");
|
||||
}
|
||||
});
|
||||
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 declare(
|
||||
&self,
|
||||
batch_for: impl FnOnce(&RedisClient) -> RedisBatch,
|
||||
increment: Increment,
|
||||
) -> Result<Option<DeclaredIncrement>, MemoryError> {
|
||||
let local = increment
|
||||
.counts
|
||||
.iter()
|
||||
.map(|(key, value)| self.add_local(key, *value, COUNTER_TTL))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let Some(remote) = self.0.remote() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let batch = batch_for(remote);
|
||||
let totals = increment
|
||||
.counts
|
||||
.iter()
|
||||
.map(|(key, value)| {
|
||||
batch.incr_by_float(key, value.as_f64(), Expiry::Refresh(COUNTER_TTL))
|
||||
})
|
||||
.collect();
|
||||
Ok(Some(DeclaredIncrement {
|
||||
keys: increment.counts.into_iter().map(|(key, _)| key).collect(),
|
||||
local,
|
||||
totals,
|
||||
}))
|
||||
}
|
||||
|
||||
fn mirror(
|
||||
&self,
|
||||
keys: &[String],
|
||||
local: Vec<Number>,
|
||||
totals: Result<Vec<f64>, litellm_redis_client::Error>,
|
||||
) -> Result<(), MemoryError> {
|
||||
let totals = match totals {
|
||||
Ok(totals) => totals.into_iter().map(Number::Float).collect(),
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "Redis usage pipeline failed, keeping the in-memory counts");
|
||||
local
|
||||
}
|
||||
};
|
||||
keys.iter().zip(totals).try_for_each(|(key, total)| {
|
||||
self.0
|
||||
.set_memory(key, Number::Int(total.truncate()), COUNTER_TTL)
|
||||
})
|
||||
}
|
||||
|
||||
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
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Increment {
|
||||
counts: [(String, Number); 2],
|
||||
}
|
||||
|
||||
impl Increment {
|
||||
fn new(
|
||||
deployment_id: &str,
|
||||
model: Option<&str>,
|
||||
minute: &str,
|
||||
tokens: Number,
|
||||
requests: i64,
|
||||
) -> Self {
|
||||
let counter = |kind| Counter {
|
||||
deployment_id: deployment_id.to_owned(),
|
||||
model: model.map(str::to_owned),
|
||||
minute: minute.to_owned(),
|
||||
kind,
|
||||
};
|
||||
Self {
|
||||
counts: [
|
||||
(counter(Kind::Tpm).key(), tokens),
|
||||
(counter(Kind::Rpm).key(), Number::Int(requests)),
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct DeclaredIncrement {
|
||||
keys: Vec<String>,
|
||||
local: Vec<Number>,
|
||||
totals: Vec<Declared<f64>>,
|
||||
}
|
||||
|
||||
async fn totals<F: Future<Output = Result<f64, litellm_redis_client::Error>>>(
|
||||
declared: Vec<Declared<f64>>,
|
||||
wait: impl Fn(Declared<f64>) -> F,
|
||||
) -> Result<Vec<f64>, litellm_redis_client::Error> {
|
||||
let mut totals = Vec::with_capacity(declared.len());
|
||||
for total in declared {
|
||||
totals.push(wait(total).await?);
|
||||
}
|
||||
Ok(totals)
|
||||
}
|
||||
114
litellm-rust/crates/redis-store/tests/request.rs
Normal file
114
litellm-rust/crates/redis-store/tests/request.rs
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_client::{Expiry, Operation};
|
||||
use litellm_redis_store::RequestRedis;
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
use support::{connect, connect_observed};
|
||||
|
||||
const KEEP: Expiry = Expiry::KeepExisting(Duration::from_secs(60));
|
||||
const LONG_DEADLINE: Duration = Duration::from_secs(60);
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn pre_declarations_on_one_client_share_one_pipeline_and_another_client_gets_its_own(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let (other, other_trips) = connect_observed(&deployment);
|
||||
let request = RequestRedis::new(LONG_DEADLINE);
|
||||
|
||||
let first = request.pre(&client).incr_by("a", 1, KEEP);
|
||||
let second = request.pre(&client.clone()).incr_by("b", 2, KEEP);
|
||||
let third = request.pre(&other).incr_by("c", 3, KEEP);
|
||||
request.flush_pre().await;
|
||||
|
||||
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
|
||||
assert_eq!(other_trips.seen(), [Operation::Pipeline { ops: 1 }]);
|
||||
assert_eq!(
|
||||
[
|
||||
first.resolve().await.unwrap(),
|
||||
second.resolve().await.unwrap(),
|
||||
third.resolve().await.unwrap(),
|
||||
],
|
||||
[1, 2, 3]
|
||||
);
|
||||
assert_eq!(trips.seen().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_post_declaration_waits_for_flush_post_and_flush_pre_leaves_it_alone(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let request = RequestRedis::new(LONG_DEADLINE);
|
||||
|
||||
let total = tokio::spawn(
|
||||
request
|
||||
.post(&client)
|
||||
.incr_by_float("spend", 2.5, KEEP)
|
||||
.settled(),
|
||||
);
|
||||
request.flush_pre().await;
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
|
||||
assert!(!total.is_finished());
|
||||
assert!(trips.seen().is_empty());
|
||||
assert_eq!(raw.get::<String>("spend").await.unwrap(), None);
|
||||
|
||||
request.flush_post().await;
|
||||
|
||||
assert_eq!(total.await.unwrap().unwrap(), 2.5);
|
||||
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 1 }]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_post_batch_nobody_flushes_goes_out_at_the_deadline(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let request = RequestRedis::new(Duration::from_millis(100));
|
||||
|
||||
let first = request.post(&client).incr_by("a", 1, KEEP);
|
||||
let second = request.post(&client).incr_by("b", 2, KEEP);
|
||||
let settled = tokio::time::timeout(Duration::from_secs(5), async {
|
||||
(first.settled().await, second.settled().await)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!((settled.0.unwrap(), settled.1.unwrap()), (1, 2));
|
||||
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_post_declaration_after_a_flush_arms_a_new_deadline(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let request = RequestRedis::new(Duration::from_millis(300));
|
||||
|
||||
let early = request.post(&client).incr_by("a", 1, KEEP);
|
||||
request.flush_post().await;
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
let late = request.post(&client).incr_by("a", 1, KEEP);
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
let flushed_before_its_deadline = trips.seen().len() > 1;
|
||||
let late = tokio::time::timeout(Duration::from_secs(5), late.settled())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(early.settled().await.unwrap(), 1);
|
||||
assert!(!flushed_before_its_deadline);
|
||||
assert_eq!(late.unwrap(), 2);
|
||||
}
|
||||
394
litellm-rust/crates/redis-store/tests/router_cooldowns.rs
Normal file
394
litellm-rust/crates/redis-store/tests/router_cooldowns.rs
Normal file
|
|
@ -0,0 +1,394 @@
|
|||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_client::{Expiry, Operation, Ttl};
|
||||
use litellm_redis_store::{
|
||||
Number, RequestRedis,
|
||||
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())]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_declared_active_read_rides_the_request_batch_and_answers_what_active_does(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = support::connect_observed(&deployment);
|
||||
let store = RouterCooldowns::with_clock(Some(client), POLICY, fixture_clock());
|
||||
let models = cooldowns().map(|case| case.model_id.to_owned());
|
||||
for case in cooldowns() {
|
||||
raw.set(case.key, case.str_dict, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let request = RequestRedis::new(Duration::from_secs(60));
|
||||
|
||||
let declared = store.declare_active(&request, &models).unwrap();
|
||||
assert!(trips.seen().is_empty());
|
||||
request.flush_pre().await;
|
||||
let active = store.resolve_active(declared).await.unwrap();
|
||||
|
||||
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
|
||||
assert_eq!(
|
||||
active,
|
||||
cooldowns()
|
||||
.map(|case| (case.model_id.to_owned(), case.cooldown))
|
||||
.to_vec()
|
||||
);
|
||||
assert_eq!(store.active(&models).await.unwrap(), active);
|
||||
assert_eq!(trips.seen().len(), 1, "the declared read backfilled memory");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_cooldown_recorded_after_the_declaration_wins_over_the_redis_miss(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock());
|
||||
let [case, _] = cooldowns();
|
||||
let request = RequestRedis::new(Duration::from_secs(60));
|
||||
|
||||
let declared = store
|
||||
.declare_active(&request, &[case.model_id.to_owned()])
|
||||
.unwrap();
|
||||
request.flush_pre().await;
|
||||
store
|
||||
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store.resolve_active(declared).await.unwrap(),
|
||||
vec![(case.model_id.to_owned(), case.cooldown)]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_breaker_refusal_of_a_declared_read_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 guarded = support::connect_guarded(&deployment, breaker);
|
||||
let store = RouterCooldowns::with_clock(Some(guarded.clone()), POLICY, fixture_clock());
|
||||
let [case, _] = cooldowns();
|
||||
let models = [case.model_id.to_owned()];
|
||||
deployment.stop().await.unwrap();
|
||||
let opened = RequestRedis::new(Duration::from_secs(60));
|
||||
let failed = store
|
||||
.declare_active(&opened, &["dep-other".to_owned()])
|
||||
.unwrap();
|
||||
opened.flush_pre().await;
|
||||
store.resolve_active(failed).await.unwrap();
|
||||
let refused_request = RequestRedis::new(Duration::from_secs(60));
|
||||
|
||||
let refused = store.declare_active(&refused_request, &models).unwrap();
|
||||
refused_request.flush_pre().await;
|
||||
let refused = store.resolve_active(refused).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();
|
||||
|
||||
assert_eq!(refused, vec![]);
|
||||
assert_eq!(
|
||||
store.active(&models).await.unwrap(),
|
||||
vec![(case.model_id.to_owned(), case.cooldown.clone())]
|
||||
);
|
||||
}
|
||||
240
litellm-rust/crates/redis-store/tests/router_usage.rs
Normal file
240
litellm-rust/crates/redis-store/tests/router_usage.rs
Normal file
|
|
@ -0,0 +1,240 @@
|
|||
mod support;
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_redis_client::{Expiry, Operation, Ttl};
|
||||
use litellm_redis_store::{
|
||||
Number, RequestRedis,
|
||||
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))]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn post_call_increments_wait_for_flush_post_then_mirror_the_redis_totals_into_memory(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = support::connect_observed(&deployment);
|
||||
let store = Arc::new(RouterUsage::new(Some(client), POLICY));
|
||||
let counters = tpm_and_rpm(Some("model"));
|
||||
let [tpm_key, rpm_key] =
|
||||
["tpm", "rpm"].map(|kind| format!("global_router:dep:model:{kind}:12-34"));
|
||||
raw.incr_by_float(&tpm_key, 10.0, Expiry::Refresh(COUNTER_TTL))
|
||||
.await
|
||||
.unwrap();
|
||||
let request = RequestRedis::new(Duration::from_secs(60));
|
||||
|
||||
store
|
||||
.increment_deployment_post_call(
|
||||
&request,
|
||||
"dep",
|
||||
Some("model"),
|
||||
"12-34",
|
||||
Number::Float(5.5),
|
||||
1,
|
||||
)
|
||||
.unwrap();
|
||||
let before_flush = store.get_many(&counters).await.unwrap();
|
||||
let unsent: Option<String> = raw.get(&rpm_key).await.unwrap();
|
||||
let trips_before_flush = trips.seen();
|
||||
request.flush_post().await;
|
||||
let mirrored = litellm_redis_testing::eventually(async || {
|
||||
let counted = store.get_many(&counters).await.unwrap();
|
||||
(counted == [Some(Number::Int(15)), Some(Number::Int(1))]).then_some(counted)
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
before_flush,
|
||||
vec![Some(Number::Float(5.5)), Some(Number::Int(1))]
|
||||
);
|
||||
assert_eq!(unsent, None);
|
||||
assert!(trips_before_flush.is_empty());
|
||||
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
|
||||
assert_eq!(mirrored, vec![Some(Number::Int(15)), Some(Number::Int(1))]);
|
||||
let written: [Option<String>; 2] = [
|
||||
raw.get(&tpm_key).await.unwrap(),
|
||||
raw.get(&rpm_key).await.unwrap(),
|
||||
];
|
||||
assert_eq!(written, [Some("15.5".to_owned()), Some("1".to_owned())]);
|
||||
let Ttl::Expires(remaining) = raw.ttl(&rpm_key).await.unwrap() else {
|
||||
panic!("{rpm_key} has no TTL");
|
||||
};
|
||||
assert!(within(remaining, 60), "{remaining:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_failed_post_call_pipeline_keeps_the_memory_counts_as_integers() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let store = Arc::new(RouterUsage::new(Some(connect(&deployment)), POLICY));
|
||||
let counters = tpm_and_rpm(Some("model"));
|
||||
let request = RequestRedis::new(Duration::from_secs(60));
|
||||
deployment.stop().await.unwrap();
|
||||
|
||||
store
|
||||
.increment_deployment_post_call(
|
||||
&request,
|
||||
"dep",
|
||||
Some("model"),
|
||||
"12-34",
|
||||
Number::Float(5.75),
|
||||
1,
|
||||
)
|
||||
.unwrap();
|
||||
request.flush_post().await;
|
||||
|
||||
litellm_redis_testing::eventually(async || {
|
||||
(store.get_many(&counters).await.unwrap() == [Some(Number::Int(5)), Some(Number::Int(1))])
|
||||
.then_some(())
|
||||
})
|
||||
.await;
|
||||
}
|
||||
188
litellm-rust/crates/redis-store/tests/support/mod.rs
Normal file
188
litellm-rust/crates/redis-store/tests/support/mod.rs
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
|
||||
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_circuit_breaker::{State, Transition};
|
||||
use litellm_redis_client::{
|
||||
CallEvent, ClusterEndpoint, ClusterNode, Credentials, Endpoint, Namespace, Operation,
|
||||
RedisClient, RedisConfig, RedisEvents, 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, None)
|
||||
}
|
||||
|
||||
/// Every round trip a client made, as the client reports it.
|
||||
#[derive(Default)]
|
||||
pub struct RoundTrips(Mutex<Vec<Operation>>);
|
||||
|
||||
impl RoundTrips {
|
||||
pub fn seen(&self) -> Vec<Operation> {
|
||||
self.0.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisEvents for RoundTrips {
|
||||
fn call(&self, event: &CallEvent) {
|
||||
self.0.lock().unwrap().push(event.operation);
|
||||
}
|
||||
|
||||
fn breaker(&self, _transition: &Transition<State>) {}
|
||||
}
|
||||
|
||||
pub fn connect_observed(deployment: &RedisDeployment) -> (RedisClient, Arc<RoundTrips>) {
|
||||
let trips = Arc::new(RoundTrips::default());
|
||||
(connect_with(deployment, None, Some(trips.clone())), trips)
|
||||
}
|
||||
|
||||
pub fn connect_guarded(
|
||||
deployment: &RedisDeployment,
|
||||
breaker: litellm_circuit_breaker::Breaker,
|
||||
) -> RedisClient {
|
||||
connect_with(deployment, Some(breaker), None)
|
||||
}
|
||||
|
||||
fn connect_with(
|
||||
deployment: &RedisDeployment,
|
||||
breaker: Option<litellm_circuit_breaker::Breaker>,
|
||||
events: Option<Arc<dyn RedisEvents>>,
|
||||
) -> 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,
|
||||
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))
|
||||
}
|
||||
|
|
@ -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?;
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ from contextvars import ContextVar, Token
|
|||
from dataclasses import dataclass, field
|
||||
from datetime import timedelta
|
||||
from types import MappingProxyType, TracebackType
|
||||
from typing import Final, Generic, Protocol, TypeVar
|
||||
from typing import TYPE_CHECKING, Final, Generic, Protocol, TypeAlias, TypeVar
|
||||
|
||||
from litellm._internal_context import (
|
||||
REDIS_FAMILIES_METADATA_KEY,
|
||||
|
|
@ -38,7 +38,11 @@ from litellm.caching.redis_cache import (
|
|||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.rust_bridge._native import NativeRequestRedis
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_T_co = TypeVar("_T_co", covariant=True)
|
||||
_ScriptArg = str | bytes | int | float
|
||||
SettledHook = Callable[[asyncio.Future[_T]], Awaitable[None] | None]
|
||||
POST_CALL_FLUSH_DEADLINE_SECONDS: Final = 1.0
|
||||
|
|
@ -443,6 +447,37 @@ def _backend_key(redis_cache: RedisCache) -> object:
|
|||
return (type(redis_cache), redis_cache.namespace, settings)
|
||||
|
||||
|
||||
class NativePending(Protocol[_T_co]):
|
||||
def resolve(self) -> Awaitable[_T_co]: ...
|
||||
|
||||
|
||||
class NativeBatchResult(Generic[_T]):
|
||||
"""Awaitable handle for a read declared on the request's native pre-call batch; awaiting it runs the
|
||||
native flush hooks, flushes that batch, then answers from what it read."""
|
||||
|
||||
__slots__ = ("_pending", "_request")
|
||||
|
||||
def __init__(self, request: RequestRedisBatches, pending: NativePending[_T]) -> None:
|
||||
self._request: Final = request
|
||||
self._pending: Final = pending
|
||||
|
||||
def __await__(self) -> Generator[object, None, _T]:
|
||||
return self._wait().__await__()
|
||||
|
||||
async def _wait(self) -> _T:
|
||||
await self._request.flush_native_pre()
|
||||
return await self._pending.resolve()
|
||||
|
||||
|
||||
NativeRequestFactory: TypeAlias = Callable[[float], "NativeRequestRedis"]
|
||||
|
||||
|
||||
def _native_request(post_call_deadline: float) -> NativeRequestRedis:
|
||||
from litellm.rust_bridge._native import NativeRequestRedis
|
||||
|
||||
return NativeRequestRedis(post_call_deadline)
|
||||
|
||||
|
||||
_open_post_call: Final[weakref.WeakSet[RequestRedisBatches]] = weakref.WeakSet()
|
||||
"""Requests whose post-call batch still holds declared ops, so a shutdown can send them before Redis goes away."""
|
||||
|
||||
|
|
@ -460,15 +495,25 @@ class RequestRedisBatches:
|
|||
"_batches",
|
||||
"_deadline",
|
||||
"_deadline_flush",
|
||||
"_native",
|
||||
"_native_factory",
|
||||
"_native_flush_hooks",
|
||||
"_post_call",
|
||||
"post_call_deadline",
|
||||
"prefetched",
|
||||
)
|
||||
|
||||
def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS,
|
||||
native_factory: NativeRequestFactory = _native_request,
|
||||
) -> None:
|
||||
self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
|
||||
self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
|
||||
self.post_call_deadline: Final = post_call_deadline
|
||||
self._native_factory: Final = native_factory
|
||||
self._native: NativeRequestRedis | None = None
|
||||
self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry
|
||||
self._deadline: asyncio.TimerHandle | None = None
|
||||
self._deadline_flush: asyncio.Task[None] | None = None
|
||||
# Reads declared early for a consumer that runs later in the request, keyed by consumer name.
|
||||
|
|
@ -495,20 +540,49 @@ class RequestRedisBatches:
|
|||
_open_post_call.add(self)
|
||||
return batch
|
||||
|
||||
def native(self) -> NativeRequestRedis:
|
||||
"""The request's native batches, built on first use so requests that touch no native store pay nothing."""
|
||||
if self._native is None:
|
||||
self._native = self._native_factory(self.post_call_deadline)
|
||||
return self._native
|
||||
|
||||
def native_post_call(self) -> NativeRequestRedis:
|
||||
"""``native`` for a post-call declaration, so a shutdown drain also flushes it."""
|
||||
_open_post_call.add(self)
|
||||
return self.native()
|
||||
|
||||
def native_result(self, pending: NativePending[_T]) -> NativeBatchResult[_T]:
|
||||
return NativeBatchResult(self, pending)
|
||||
|
||||
def add_native_flush_hook(self, hook: Callable[[], None]) -> None:
|
||||
"""Called before every native pre-call flush so late readers can declare into the same trip."""
|
||||
self._native_flush_hooks.append(hook)
|
||||
|
||||
async def flush_native_pre(self) -> None:
|
||||
for hook in self._native_flush_hooks:
|
||||
hook()
|
||||
if self._native is not None:
|
||||
await self._native.flush_pre()
|
||||
|
||||
def _flush_on_deadline(self) -> None:
|
||||
self._deadline = None
|
||||
self._deadline_flush = asyncio.ensure_future(self.flush_post_call())
|
||||
|
||||
async def flush_all(self) -> None:
|
||||
"""Send whatever is still declared (write-backs nobody awaits) before the request scope closes."""
|
||||
await asyncio.gather(*(batch.flush() for batch in self._batches.values() if batch.pending))
|
||||
await asyncio.gather(
|
||||
*(batch.flush() for batch in self._batches.values() if batch.pending), self.flush_native_pre()
|
||||
)
|
||||
|
||||
async def flush_post_call(self) -> None:
|
||||
"""One pipeline per backend for the post-call writes; the deadline is disarmed since this is that flush."""
|
||||
if self._deadline is not None:
|
||||
self._deadline.cancel()
|
||||
self._deadline = None
|
||||
await asyncio.gather(*(batch.flush() for batch in self._post_call.values() if batch.pending))
|
||||
await asyncio.gather(
|
||||
*(batch.flush() for batch in self._post_call.values() if batch.pending),
|
||||
*(() if self._native is None else (self._native.flush_post(),)),
|
||||
)
|
||||
if not any(batch.pending for batch in self._post_call.values()):
|
||||
_open_post_call.discard(self)
|
||||
|
||||
|
|
@ -557,17 +631,24 @@ async def drain_post_call_redis_batches() -> None:
|
|||
class request_redis_batch_scope:
|
||||
"""Redis reads declared inside share one pipeline per backend; nested scopes join the outer one."""
|
||||
|
||||
__slots__ = ("_post_call_deadline", "_token")
|
||||
__slots__ = ("_native_factory", "_post_call_deadline", "_token")
|
||||
|
||||
def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS,
|
||||
native_factory: NativeRequestFactory = _native_request,
|
||||
) -> None:
|
||||
self._token: Token[RequestRedisBatches | None] | None = None
|
||||
self._post_call_deadline: Final = post_call_deadline
|
||||
self._native_factory: Final = native_factory
|
||||
|
||||
def __enter__(self) -> RequestRedisBatches:
|
||||
outer: Final = _active_request_batches.get()
|
||||
if outer is not None:
|
||||
return outer
|
||||
batches: Final = RequestRedisBatches(post_call_deadline=self._post_call_deadline)
|
||||
batches: Final = RequestRedisBatches(
|
||||
post_call_deadline=self._post_call_deadline, native_factory=self._native_factory
|
||||
)
|
||||
self._token = _active_request_batches.set(batches)
|
||||
return batches
|
||||
|
||||
|
|
|
|||
|
|
@ -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
14
litellm/redis/stores/_connection.py
Normal 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
|
||||
)
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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,11 @@ 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_batch import NativeBatchResult, NativePending, RequestRedisBatches
|
||||
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 NativeRequestRedis, NativeRouterCooldowns
|
||||
|
||||
from . import _routing_prefetch
|
||||
from ._routing_prefetch import TierRead
|
||||
from ._connection import redis_kwargs
|
||||
|
||||
|
||||
class CooldownCacheValue(TypedDict):
|
||||
|
|
@ -28,156 +22,96 @@ 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 declare_active(
|
||||
self, request: RequestRedisBatches, model_ids: Sequence[str]
|
||||
) -> NativeBatchResult[ActiveCooldowns] | None: ...
|
||||
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 declare_active(
|
||||
self, request: NativeRequestRedis, model_ids: Sequence[str]
|
||||
) -> NativePending[ActiveCooldowns]: ...
|
||||
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)
|
||||
pending: Final = self._handle.aactive(model_ids)
|
||||
return tuple(await pending)
|
||||
|
||||
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float:
|
||||
def declare_active(
|
||||
self, request: RequestRedisBatches, model_ids: Sequence[str]
|
||||
) -> NativeBatchResult[ActiveCooldowns] | None:
|
||||
"""``active`` declared on the request's native pre-call batch; None when there is no Redis to read."""
|
||||
if not self._connected:
|
||||
return None
|
||||
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.declare_active(request.native(), model_ids)
|
||||
return request.native_result(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()
|
||||
|
|
|
|||
|
|
@ -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,12 @@ 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_batch import active_request_redis_batches
|
||||
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 NativeRequestRedis, 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 +26,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 +42,62 @@ 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,
|
||||
*,
|
||||
request: NativeRequestRedis | None = None,
|
||||
) -> 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 +109,18 @@ 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
|
||||
request: Final = active_request_redis_batches()
|
||||
await self._handle.increment_deployment(
|
||||
deployment_id,
|
||||
model,
|
||||
minute,
|
||||
tokens,
|
||||
requests,
|
||||
request=None if request is None else request.native_post_call(),
|
||||
)
|
||||
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)
|
||||
|
||||
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()
|
||||
|
|
|
|||
165
litellm/router_utils/routing_prefetch.py
Normal file
165
litellm/router_utils/routing_prefetch.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, 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
|
||||
from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, ActiveCooldowns, RouterCooldowns
|
||||
|
||||
_PREFETCH_SLOT: Final = "routing_read"
|
||||
_COOLDOWNS_SLOT: Final = "routing_cooldowns"
|
||||
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
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ArmedCooldowns:
|
||||
model_ids: frozenset[str]
|
||||
result: Awaitable[ActiveCooldowns]
|
||||
|
||||
|
||||
def arm_cooldowns(store: RouterCooldowns, model_ids: Sequence[str]) -> None:
|
||||
request: Final = active_request_redis_batches()
|
||||
if request is None or _COOLDOWNS_SLOT in request.prefetched:
|
||||
return
|
||||
result: Final = store.declare_active(request, model_ids)
|
||||
if result is not None:
|
||||
request.prefetched[_COOLDOWNS_SLOT] = ArmedCooldowns(frozenset(model_ids), result)
|
||||
|
||||
|
||||
def armed() -> bool:
|
||||
request: Final = active_request_redis_batches()
|
||||
return request is not None and (_PREFETCH_SLOT in request.prefetched or _COOLDOWNS_SLOT in request.prefetched)
|
||||
|
||||
|
||||
async def _armed_cooldowns(model_ids: Sequence[str]) -> ActiveCooldowns | None:
|
||||
request: Final = active_request_redis_batches()
|
||||
taken: Final = None if request is None else request.prefetched.pop(_COOLDOWNS_SLOT, None)
|
||||
if not isinstance(taken, ArmedCooldowns):
|
||||
return None
|
||||
try:
|
||||
active: Final = await taken.result
|
||||
except Exception as e: # noqa: BLE001 # the read falls back to the store's own Redis read
|
||||
verbose_router_logger.debug("cooldown prefetch failed, reading again: %s", e)
|
||||
return None
|
||||
if not taken.model_ids.issuperset(model_ids):
|
||||
return None
|
||||
wanted: Final = frozenset(model_ids)
|
||||
return tuple((model_id, cooldown) for model_id, cooldown in active if model_id in wanted)
|
||||
|
||||
|
||||
async def read_cooldowns(
|
||||
store: RouterCooldowns, model_ids: Sequence[str], parent_otel_span: Span | None
|
||||
) -> ActiveCooldowns:
|
||||
"""The armed cooldown read when it covers ``model_ids``, taken once; otherwise the store reads on its own."""
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
prefetched: Final = await _armed_cooldowns(model_ids)
|
||||
if prefetched is not None:
|
||||
return prefetched
|
||||
return await store.active(model_ids, parent_otel_span)
|
||||
|
||||
|
||||
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]
|
||||
|
|
@ -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,22 @@ 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)
|
||||
)
|
||||
routing_prefetch.arm_cooldowns(ctx.health().cooldowns.cooldown_store, ctx.registry().names(kind="ids"))
|
||||
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 +79,16 @@ 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(
|
||||
routing_prefetch.read_cooldowns(
|
||||
ctx.health().cooldowns.cooldown_store, 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]
|
||||
|
|
|
|||
|
|
@ -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,61 @@ 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 NativeRequestRedis:
|
||||
def __new__(cls, post_call_deadline_seconds: float) -> NativeRequestRedis: ...
|
||||
def flush_pre(self) -> Future[None]: ...
|
||||
def flush_post(self) -> Future[None]: ...
|
||||
|
||||
@final
|
||||
class NativeDeclaredCooldowns:
|
||||
def resolve(self) -> Future[list[tuple[str, _NativeCooldown]]]: ...
|
||||
|
||||
@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 declare_active(self, request: NativeRequestRedis, model_ids: Sequence[str]) -> NativeDeclaredCooldowns: ...
|
||||
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,
|
||||
*,
|
||||
request: NativeRequestRedis | None = None,
|
||||
) -> 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
|
||||
|
|
@ -215,7 +271,11 @@ __all__ = [
|
|||
"NativeClickHouseSpendConfig",
|
||||
"NativeClickHouseSpendStorage",
|
||||
"NativeDatabase",
|
||||
"NativeDeclaredCooldowns",
|
||||
"NativeDiagnosticProcessor",
|
||||
"NativeRequestRedis",
|
||||
"NativeRouterCooldowns",
|
||||
"NativeRouterUsage",
|
||||
"ProcessReservedForForking",
|
||||
"ResponsesWebSocketConnection",
|
||||
"RustBridgeDeclined",
|
||||
|
|
@ -233,6 +293,7 @@ __all__ = [
|
|||
"embedding",
|
||||
"gil_stats",
|
||||
"messages",
|
||||
"native_redis_breaker_states",
|
||||
"ocr",
|
||||
"ocr_health_check_document",
|
||||
"ocr_passthrough_response",
|
||||
|
|
|
|||
|
|
@ -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
46
litellm/rust_bridge/host/redis_events.py
Normal 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)
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -6,8 +6,9 @@ import asyncio
|
|||
import hashlib
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
import pytest
|
||||
from redis.exceptions import NoScriptError
|
||||
|
|
@ -28,6 +29,7 @@ from litellm.caching.redis_cache import (
|
|||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
|
||||
SCRIPT = "return redis.call('GET', KEYS[1])"
|
||||
_Answer = TypeVar("_Answer")
|
||||
SHA = hashlib.sha1(SCRIPT.encode()).hexdigest() # noqa: S324
|
||||
|
||||
|
||||
|
|
@ -93,6 +95,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
|
||||
|
|
@ -147,6 +150,43 @@ def replies(command: tuple[Any, ...]) -> Any:
|
|||
raise AssertionError(command)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeNativeRequest:
|
||||
"""Stands in for ``NativeRequestRedis``: records each flush in ``events`` next to what the test's pendings log."""
|
||||
|
||||
post_call_deadline: float
|
||||
events: list[str] = field(default_factory=list)
|
||||
|
||||
async def flush_pre(self) -> None:
|
||||
self.events.append("flush_pre")
|
||||
|
||||
async def flush_post(self) -> None:
|
||||
self.events.append("flush_post")
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeNatives:
|
||||
built: list[FakeNativeRequest] = field(default_factory=list)
|
||||
|
||||
def __call__(self, post_call_deadline: float) -> FakeNativeRequest:
|
||||
self.built.append(FakeNativeRequest(post_call_deadline))
|
||||
return self.built[-1]
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeNativePending(Generic[_Answer]):
|
||||
"""A read declared on a ``FakeNativeRequest``; it only answers once that request's pre-call batch was flushed."""
|
||||
|
||||
request: FakeNativeRequest
|
||||
answer: Callable[[], _Answer]
|
||||
|
||||
async def resolve(self) -> _Answer:
|
||||
if "flush_pre" not in self.request.events:
|
||||
raise ConnectionError("resolved before its batch was flushed")
|
||||
self.request.events.append("resolve")
|
||||
return self.answer()
|
||||
|
||||
|
||||
def make(fail: Exception | None = None, namespace: str | None = None) -> tuple[FakeRedisCache, FakeClient]:
|
||||
client = FakeClient(replies, fail)
|
||||
return FakeRedisCache(client, namespace), client
|
||||
|
|
|
|||
|
|
@ -38,7 +38,7 @@ from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
|
|||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
from .test_redis_batch import FakeClient, FakeRedisCache
|
||||
from .test_redis_batch import FakeClient, FakeNatives, FakeRedisCache
|
||||
|
||||
|
||||
async def _script_outside_the_pipeline(keys: Sequence[str], args: Sequence[object]) -> object:
|
||||
|
|
@ -514,3 +514,31 @@ async def test_a_cancelled_post_call_flush_keeps_the_shared_spend_counter_and_co
|
|||
assert redis_cache.alone == [], "a cancel says nothing about the shared counter, so Redis keeps it"
|
||||
assert spend_cache.in_memory_cache.get_cache("spend:key:k1") == 3.5, "the local copy counts the cancelled spend"
|
||||
assert spend_cache.in_memory_cache.get_cache("spend:team:t1") is None, "an absent local copy is not seeded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_request_boundary():
|
||||
natives = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
request.native_post_call()
|
||||
await request.flush_all()
|
||||
at_the_boundary = list(natives.built[0].events)
|
||||
await flush_post_call_redis_batches()
|
||||
await drain_post_call_redis_batches()
|
||||
|
||||
assert at_the_boundary == ["flush_pre"]
|
||||
assert natives.built[0].events == ["flush_pre", "flush_post"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_post_call_ops_still_waiting_on_their_callbacks_are_drained_at_shutdown():
|
||||
natives = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
request.native_post_call()
|
||||
await request.flush_all()
|
||||
await drain_post_call_redis_batches()
|
||||
await drain_post_call_redis_batches()
|
||||
|
||||
assert natives.built[0].events == ["flush_pre", "flush_post"]
|
||||
|
|
|
|||
|
|
@ -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,21 +24,19 @@ 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,
|
||||
)
|
||||
|
||||
from .test_redis_batch import FakeClient, FakeRedisCache, replies
|
||||
from .test_redis_batch import FakeClient, FakeNativePending, FakeNatives, FakeRedisCache, replies
|
||||
|
||||
|
||||
def _cooldown_key(model_id: str) -> str:
|
||||
|
|
@ -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 set(request.prefetched) == {"routing_cooldowns"}
|
||||
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)
|
||||
|
|
@ -1081,17 +885,31 @@ async def test_identity_prefetch_is_one_mget_after_which_hits_and_misses_alike_c
|
|||
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
|
||||
@pytest.mark.asyncio
|
||||
async def test_awaiting_a_native_read_runs_the_native_hooks_then_flushes_the_native_pre_batch_once_per_request():
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(post_call_deadline=7, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
native: Final = request.native()
|
||||
request.add_native_flush_hook(lambda: native.events.append("hook"))
|
||||
result: Final = request.native_result(FakeNativePending(native, lambda: "answer"))
|
||||
declared_only: Final = list(native.events)
|
||||
answer: Final = await result
|
||||
await request.flush_all()
|
||||
|
||||
assert request.native() is native
|
||||
assert [built.post_call_deadline for built in natives.built] == [7]
|
||||
assert declared_only == []
|
||||
assert answer == "answer"
|
||||
assert native.events == ["hook", "flush_pre", "resolve", "hook", "flush_pre"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_request_that_no_native_store_asks_for_builds_no_native_batches():
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
await request.flush_all()
|
||||
await request.flush_post_call()
|
||||
|
||||
assert natives.built == []
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -1,112 +1,267 @@
|
|||
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_batch import RequestRedisBatches
|
||||
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
|
||||
from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNativeRequest, FakeNatives
|
||||
|
||||
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
|
||||
declared: list[tuple[str, ...]] = field(default_factory=list)
|
||||
|
||||
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 declare_active(
|
||||
self, request: FakeNativeRequest, model_ids: Sequence[str]
|
||||
) -> FakeNativePending[list[tuple[str, CooldownCacheValue]]]:
|
||||
self.targets.append(current_service_target())
|
||||
self.declared.append(tuple(model_ids))
|
||||
return FakeNativePending(request, lambda: self.active(model_ids))
|
||||
|
||||
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.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
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
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, ...] = ()
|
||||
|
||||
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")
|
||||
|
||||
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
|
||||
|
||||
|
||||
@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
|
||||
async def test_a_declared_active_read_answers_after_the_native_pre_flush_under_the_router_cooldowns_target() -> None:
|
||||
handles: Final = Handles()
|
||||
store: Final = RouterCooldownStore(_redis(), handle=handles)
|
||||
natives: Final = FakeNatives()
|
||||
request: Final = RequestRedisBatches(native_factory=natives) # pyright: ignore[reportArgumentType] # test double
|
||||
store.put("dep", _entry(), 60.0)
|
||||
|
||||
declared: Final = store.declare_active(request, ["dep", "other"])
|
||||
assert declared is not None
|
||||
assert natives.built[0].events == []
|
||||
|
||||
assert tuple(await declared) == (("dep", _entry()),)
|
||||
assert natives.built[0].events == ["flush_pre", "resolve"]
|
||||
assert handles.built[0].declared == [("dep", "other")]
|
||||
assert handles.built[0].targets[1] == "router_cooldowns"
|
||||
|
||||
|
||||
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) == ()
|
||||
def test_a_store_without_redis_declares_nothing_on_the_request() -> None:
|
||||
natives: Final = FakeNatives()
|
||||
request: Final = RequestRedisBatches(native_factory=natives) # pyright: ignore[reportArgumentType] # test double
|
||||
|
||||
|
||||
@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)
|
||||
|
||||
with request_redis_batch_scope() as request:
|
||||
reader.arm_routing_read(("dep",), None)
|
||||
clock.failure = RuntimeError("backfill clock unavailable")
|
||||
await request.flush_all()
|
||||
|
||||
clock.failure = None
|
||||
assert await reader.active(("dep",), None) == (("dep", value),)
|
||||
assert RouterCooldownStore(None, handle=Handles()).declare_active(request, ["dep"]) is None
|
||||
assert natives.built == []
|
||||
|
|
|
|||
|
|
@ -1,12 +1,96 @@
|
|||
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_batch import drain_post_call_redis_batches, request_redis_batch_scope
|
||||
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
|
||||
from tests.unit.caching.test_redis_batch import FakeNatives
|
||||
|
||||
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)
|
||||
requests: list[object] = 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,
|
||||
*,
|
||||
request: object = None,
|
||||
) -> "asyncio.Future[None]":
|
||||
self.calls.append(("deployment", (deployment_id, model, minute, tokens, requests), current_service_target()))
|
||||
self.requests.append(request)
|
||||
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
|
||||
|
|
@ -25,63 +109,15 @@ async def test_router_reset_clears_usage_and_failure_counters(operation: str) ->
|
|||
|
||||
|
||||
@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
|
||||
async def test_inside_a_request_deployment_increments_ride_its_native_post_call_batch_and_the_shutdown_drain():
|
||||
handles: Final = Handles()
|
||||
store: Final = RouterUsageStore(None, handle=handles)
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
await store.increment_deployment("dep", "model", "12-34", 7.5, 1, None)
|
||||
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
await store.increment_deployment("dep", "model", "12-34", 2, 1, None)
|
||||
await drain_post_call_redis_batches()
|
||||
|
||||
@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]
|
||||
assert handles.built[0].requests == [None, request.native()]
|
||||
assert natives.built[0].events == ["flush_post"]
|
||||
|
|
|
|||
|
|
@ -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",)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,11 +1,10 @@
|
|||
"""
|
||||
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
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -13,8 +12,15 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.caching.redis_batch import request_redis_batch_scope
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore
|
||||
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
|
||||
from litellm.router_utils import routing_prefetch
|
||||
from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNativeRequest, FakeNatives
|
||||
from tests.unit.conftest import redis_cache_double
|
||||
from tests.unit.redis.stores.test_router_cooldowns import FakeCooldowns
|
||||
from tests.unit.redis.stores.test_router_cooldowns import Handles as CooldownHandles
|
||||
|
||||
_MODEL_GROUP = "claude"
|
||||
_MESSAGES = [{"role": "user", "content": "ping"}]
|
||||
|
|
@ -34,7 +40,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 +61,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 +81,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 +118,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 +155,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 +167,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 +178,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,19 +192,100 @@ 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")
|
||||
|
||||
deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
|
||||
|
||||
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
|
||||
|
||||
|
||||
def _armed_store() -> tuple[RouterCooldownStore, CooldownHandles]:
|
||||
handles: Final = CooldownHandles()
|
||||
store: Final = RouterCooldownStore(RedisCache(host="redis.internal"), handle=handles)
|
||||
store.put("dep-a", {"exception_received": "429", "status_code": "429", "timestamp": 1.0, "cooldown_time": 60}, 60)
|
||||
return store, handles
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_takes_the_armed_cooldown_read_from_the_native_batch_once_then_reads_on_its_own():
|
||||
store, handles = _armed_store()
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
|
||||
routing_prefetch.arm_cooldowns(store, ["dep-a", "dep-b", "dep-c"])
|
||||
armed: Final = routing_prefetch.armed()
|
||||
first: Final = await routing_prefetch.read_cooldowns(store, ["dep-a", "dep-b"], None)
|
||||
reads_after_first: Final = len(handles.built[0].targets)
|
||||
second: Final = await routing_prefetch.read_cooldowns(store, ["dep-a", "dep-b"], None)
|
||||
|
||||
assert armed
|
||||
assert handles.built[0].declared == [("dep-a", "dep-b", "dep-c")]
|
||||
assert first == second == (("dep-a", store.snapshot()["dep-a"]),)
|
||||
assert natives.built[0].events == ["flush_pre", "resolve"]
|
||||
assert len(handles.built[0].targets) == reads_after_first + 1, "a second read goes to the store"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_armed_cooldown_read_that_does_not_cover_the_routing_ids_is_read_again():
|
||||
store, _ = _armed_store()
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
|
||||
routing_prefetch.arm_cooldowns(store, ["dep-b"])
|
||||
active: Final = await routing_prefetch.read_cooldowns(store, ["dep-a", "dep-b"], None)
|
||||
|
||||
assert active == (("dep-a", store.snapshot()["dep-a"]),)
|
||||
assert natives.built[0].events == ["flush_pre", "resolve"]
|
||||
|
||||
|
||||
def _redis_down() -> list[tuple[str, CooldownCacheValue]]:
|
||||
raise ConnectionError("redis down")
|
||||
|
||||
|
||||
class _UnreachableCooldowns(FakeCooldowns):
|
||||
def declare_active(
|
||||
self, request: FakeNativeRequest, model_ids: Sequence[str]
|
||||
) -> FakeNativePending[list[tuple[str, CooldownCacheValue]]]:
|
||||
super().declare_active(request, model_ids)
|
||||
return FakeNativePending(request, _redis_down)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_armed_cooldown_read_falls_back_to_the_stores_own_read():
|
||||
built: Final[list[FakeCooldowns]] = []
|
||||
|
||||
def unreachable(connection: RedisCache | None) -> FakeCooldowns:
|
||||
built.append(_UnreachableCooldowns(None, None))
|
||||
return built[-1]
|
||||
|
||||
store: Final = RouterCooldownStore(RedisCache(host="redis.internal"), handle=unreachable)
|
||||
store.put("dep-a", {"exception_received": "429", "status_code": "429", "timestamp": 1.0, "cooldown_time": 60}, 60)
|
||||
|
||||
with request_redis_batch_scope(native_factory=FakeNatives()): # pyright: ignore[reportArgumentType] # test double
|
||||
routing_prefetch.arm_cooldowns(store, ["dep-a"])
|
||||
active: Final = await routing_prefetch.read_cooldowns(store, ["dep-a"], None)
|
||||
|
||||
assert built[0].declared == [("dep-a",)]
|
||||
assert active == (("dep-a", store.snapshot()["dep-a"]),)
|
||||
|
||||
|
||||
def test_arming_the_cooldown_read_outside_a_request_or_without_redis_declares_nothing():
|
||||
store, handles = _armed_store()
|
||||
offline: Final = CooldownHandles()
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
routing_prefetch.arm_cooldowns(store, ["dep-a"])
|
||||
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
|
||||
routing_prefetch.arm_cooldowns(RouterCooldownStore(None, handle=offline), ["dep-a"])
|
||||
armed: Final = routing_prefetch.armed()
|
||||
|
||||
assert handles.built[0].declared == []
|
||||
assert offline.built[0].declared == []
|
||||
assert not armed
|
||||
assert natives.built == []
|
||||
|
|
|
|||
81
tests/unit/rust_bridge/host/test_redis_events.py
Normal file
81
tests/unit/rust_bridge/host/test_redis_events.py
Normal 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
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1972,6 +1972,8 @@ def _tls_view(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,
|
||||
|
|
@ -2006,6 +2008,7 @@ def _connection_view(connection: AbstractConnection, kwargs: dict[str, object])
|
|||
"auth": _auth_view(kwargs),
|
||||
"socket_timeout": connection.socket_timeout,
|
||||
"socket_connect_timeout": connection.socket_connect_timeout,
|
||||
"health_check_interval": float(connection.health_check_interval) or None,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -2047,6 +2050,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"tls": None,
|
||||
"socket_timeout": 5.0,
|
||||
"socket_connect_timeout": 5.0,
|
||||
"health_check_interval": None,
|
||||
**fields,
|
||||
}
|
||||
|
||||
|
|
@ -2064,31 +2068,10 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
_expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}),
|
||||
id="host_defaults",
|
||||
),
|
||||
pytest.param(
|
||||
{"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"},
|
||||
{},
|
||||
_expected_view(
|
||||
{
|
||||
"db": 2,
|
||||
"host": "redis",
|
||||
"mode": "host",
|
||||
"port": 6380,
|
||||
"socket_timeout": 10.5,
|
||||
"socket_connect_timeout": 10.5,
|
||||
}
|
||||
),
|
||||
id="padded_and_underscored_numbers_parse_like_python_int_and_float",
|
||||
),
|
||||
pytest.param(
|
||||
{},
|
||||
{"REDIS_HOST": "redis", "REDIS_PORT": "6_380", "REDIS_DB": "3\n"},
|
||||
_expected_view({"db": 3, "host": "redis", "mode": "host", "port": 6380}),
|
||||
id="padded_environment_numbers_parse_like_python_int",
|
||||
),
|
||||
pytest.param(
|
||||
{"startup_nodes": []},
|
||||
{"REDIS_CLUSTER_NODES": '[{"host": "node-a", "port": 7000}]'},
|
||||
_expected_view({"mode": "cluster", "nodes": ["node-a:7000"]}),
|
||||
_expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}),
|
||||
id="empty_cluster_nodes_use_the_environment",
|
||||
),
|
||||
pytest.param(
|
||||
|
|
@ -2158,6 +2141,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"socket_timeout": 1.5,
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2184,6 +2169,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"port": 6390,
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2208,6 +2195,27 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
),
|
||||
id="string_numbers_are_coerced",
|
||||
),
|
||||
pytest.param(
|
||||
{"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"},
|
||||
{},
|
||||
_expected_view(
|
||||
{
|
||||
"db": 2,
|
||||
"host": "redis",
|
||||
"mode": "host",
|
||||
"port": 6380,
|
||||
"socket_timeout": 10.5,
|
||||
"socket_connect_timeout": 10.5,
|
||||
}
|
||||
),
|
||||
id="padded_and_underscored_numbers_parse_like_python_int_and_float",
|
||||
),
|
||||
pytest.param(
|
||||
{},
|
||||
{"REDIS_HOST": "redis", "REDIS_PORT": "6_380", "REDIS_DB": "3\n"},
|
||||
_expected_view({"db": 3, "host": "redis", "mode": "host", "port": 6380}),
|
||||
id="padded_environment_numbers_parse_like_python_int",
|
||||
),
|
||||
pytest.param(
|
||||
{"host": "redis"},
|
||||
{"REDIS_SOCKET_TIMEOUT": "0.75"},
|
||||
|
|
@ -2250,6 +2258,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"mode": "url",
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2273,6 +2283,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"mode": "url",
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2293,6 +2305,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
_expected_view(
|
||||
{
|
||||
"auth": {"kind": "static", "password": "cluster-secret", "username": None},
|
||||
"health_check_interval": 25.0,
|
||||
"mode": "cluster",
|
||||
"nodes": ["node-a:7000"],
|
||||
}
|
||||
|
|
@ -2302,7 +2315,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
pytest.param(
|
||||
{},
|
||||
{"REDIS_CLUSTER_NODES": '[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]'},
|
||||
_expected_view({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"]}),
|
||||
_expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"]}),
|
||||
id="environment_cluster_nodes",
|
||||
),
|
||||
pytest.param(
|
||||
|
|
@ -2315,10 +2328,13 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
{},
|
||||
_expected_view(
|
||||
{
|
||||
"health_check_interval": 25.0,
|
||||
"mode": "cluster",
|
||||
"nodes": ["node-a:7000"],
|
||||
"tls": {
|
||||
"ca_certs": "/ca.pem",
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": True,
|
||||
"keyfile": None,
|
||||
|
|
@ -2430,6 +2446,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"port": 6379,
|
||||
"tls": {
|
||||
"ca_certs": "/gcp-ca.pem",
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2512,6 +2530,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"port": 6379,
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2545,6 +2565,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"port": 6379,
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2576,6 +2598,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"mode": "url",
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2632,6 +2656,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"port": 6379,
|
||||
"tls": {
|
||||
"ca_certs": "/ca.pem",
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": True,
|
||||
"keyfile": None,
|
||||
|
|
@ -2652,6 +2678,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"port": 6379,
|
||||
"tls": {
|
||||
"ca_certs": "/env-ca.pem",
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2707,6 +2735,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"port": 6379,
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": "/client.pem",
|
||||
"check_hostname": False,
|
||||
"keyfile": "/client.key",
|
||||
|
|
@ -2724,6 +2754,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"mode": "url",
|
||||
"tls": {
|
||||
"ca_certs": "/ca.pem",
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": True,
|
||||
"keyfile": None,
|
||||
|
|
@ -2755,6 +2787,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"mode": "url",
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2790,7 +2824,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
pytest.param(
|
||||
{"db": False, "startup_nodes": [{"host": "node-a", "port": 7000}]},
|
||||
{},
|
||||
_expected_view({"mode": "cluster", "nodes": ["node-a:7000"]}),
|
||||
_expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}),
|
||||
id="cluster_db_false",
|
||||
),
|
||||
pytest.param(
|
||||
|
|
@ -2807,7 +2841,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
pytest.param(
|
||||
{"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]},
|
||||
{},
|
||||
_expected_view({"mode": "cluster", "nodes": ["node-a:7000"]}),
|
||||
_expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}),
|
||||
id="cluster_db_float_zero",
|
||||
),
|
||||
pytest.param(
|
||||
|
|
@ -2824,7 +2858,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
pytest.param(
|
||||
{"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]},
|
||||
{},
|
||||
_expected_view({"mode": "cluster", "nodes": ["node-a:7000"]}),
|
||||
_expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}),
|
||||
id="cluster_db_empty_list",
|
||||
),
|
||||
pytest.param(
|
||||
|
|
@ -2838,6 +2872,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
"port": 6379,
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": None,
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
|
|
@ -2895,6 +2931,78 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]:
|
|||
_expected_view({"mode": "unix", "path": "/tmp/redis.sock"}),
|
||||
id="unix_socket_path",
|
||||
),
|
||||
pytest.param(
|
||||
{"host": "redis", "ssl": True, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"},
|
||||
{},
|
||||
_expected_view({"mode": "unix", "path": "/tmp/redis.sock"}),
|
||||
id="unix_socket_path_keeps_tls_settings_out_of_the_target",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"host": "redis",
|
||||
"ssl": True,
|
||||
"ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----",
|
||||
},
|
||||
{},
|
||||
_expected_view(
|
||||
{
|
||||
"db": 0,
|
||||
"host": "redis",
|
||||
"mode": "host",
|
||||
"port": 6379,
|
||||
"tls": {
|
||||
"ca_certs": None,
|
||||
"ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----",
|
||||
"ca_path": None,
|
||||
"certfile": None,
|
||||
"check_hostname": False,
|
||||
"keyfile": None,
|
||||
"verify_certificate": True,
|
||||
},
|
||||
}
|
||||
),
|
||||
id="ssl_ca_data",
|
||||
),
|
||||
pytest.param(
|
||||
{"host": "redis"},
|
||||
{"REDIS_SSL": "true", "REDIS_SSL_CA_PATH": "/etc/ssl/private-cas"},
|
||||
{
|
||||
"rejected": {
|
||||
"message": "Redis.__init__() got an unexpected keyword argument 'ssl_ca_path'",
|
||||
"type": "TypeError",
|
||||
}
|
||||
},
|
||||
id="ssl_ca_path_from_the_environment",
|
||||
),
|
||||
pytest.param(
|
||||
{"health_check_interval": "7", "host": "redis"},
|
||||
{},
|
||||
_expected_view({"db": 0, "health_check_interval": 7.0, "host": "redis", "mode": "host", "port": 6379}),
|
||||
id="health_check_interval_is_whole_seconds",
|
||||
),
|
||||
pytest.param(
|
||||
{"health_check_interval": 0, "host": "redis"},
|
||||
{},
|
||||
_expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}),
|
||||
id="health_check_interval_zero_is_off",
|
||||
),
|
||||
pytest.param(
|
||||
{"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]},
|
||||
{},
|
||||
{
|
||||
"rejected": {
|
||||
"message": "Argument 'db' must be 0 or None in cluster mode",
|
||||
"type": "RedisClusterException",
|
||||
}
|
||||
},
|
||||
id="a_cluster_rejects_a_nonzero_db",
|
||||
),
|
||||
pytest.param(
|
||||
{"unix_socket_path": "/tmp/redis.sock"},
|
||||
{},
|
||||
{"rejected": {"message": "Either 'host' or 'url' must be specified for redis.", "type": "ValueError"}},
|
||||
id="unix_socket_path_without_a_host_fails_like_python",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_redis_settings_resolve_like_the_rust_resolver_expects(monkeypatch, kwargs, environment, expected):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -15606,11 +15606,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)
|
||||
|
|
@ -15625,7 +15624,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")
|
||||
|
|
@ -15635,7 +15633,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)
|
||||
|
||||
|
|
@ -15650,10 +15648,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)
|
||||
|
|
@ -15677,10 +15674,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)
|
||||
|
|
@ -15708,11 +15704,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()
|
||||
|
|
@ -24299,11 +24294,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=[
|
||||
{
|
||||
|
|
@ -24313,8 +24304,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",
|
||||
|
|
@ -24328,11 +24317,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
|
||||
|
|
@ -24352,9 +24348,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."""
|
||||
|
|
@ -24374,9 +24368,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={})
|
||||
|
||||
|
|
@ -24385,8 +24376,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]]]:
|
||||
|
|
@ -25141,7 +25130,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=[
|
||||
{
|
||||
|
|
@ -25152,26 +25141,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=[
|
||||
{
|
||||
|
|
@ -25182,9 +25163,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",
|
||||
|
|
@ -25192,12 +25170,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
|
||||
|
|
@ -25352,7 +25325,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue