refactor(router): move Redis key construction out of router-core into the stores

router-core's RouterCache now names the routing state it touches (cooldowns,
allowed fails, per-minute usage, health states) instead of spelling Redis keys.
The stores own every key, and the Python host answers each typed cache step
from its store without parsing keys back into ids.
This commit is contained in:
Yujong Lee 2026-11-06 15:37:43 +09:00
parent 08f2a98250
commit 6d019c9dc8
14 changed files with 938 additions and 887 deletions

View file

@ -6,8 +6,9 @@
//! current thread, so a host future finds its exchange without the router passing it through.
//!
//! The routing cache is a host operation too: [`StepCache`] parks every [`RouterCache`] op as a
//! `cache` step, and Python runs it against the router's own `DualCache` (memory plus Redis), so
//! Python and Rust read and write the same cooldowns, fail counters and usage.
//! `cache` step, and Python answers it from the router's stores, so Python and Rust read and write
//! the same cooldowns, fail counters and usage. A step names the state it touches, never a Redis
//! key: the keys are the stores'.
use std::{
cell::RefCell,
@ -22,8 +23,8 @@ use async_trait::async_trait;
use litellm_router_core::{
Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, Deployment, Error, Host, Output,
Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, RouterEvent, RouterRequest,
RouterResponse, RoutingContext, SecretRef, SelectStage, ThreadRandomness,
build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo,
RouterResponse, RoutingContext, SecretRef, SelectStage, ThreadRandomness, UsageCounts,
UsageWindow, build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo,
precall::PreCallFacts,
};
use pyo3::{PyResult, Python, types::PyAnyMethods};
@ -102,55 +103,115 @@ pub(crate) enum Stage {
ProviderRequest,
}
/// The `op` of a `cache` step: one per [`RouterCache`] method.
/// One [`RouterCache`] op, as a `cache` step carries it to Python.
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum CacheCall {
/// Answered with one stored cooldown or `None` per id.
Cooldowns { ids: Vec<String> },
/// Answered with `None`.
PutCooldown {
id: String,
cooldown: Value,
ttl: f64,
},
/// Answered with the count after this failure.
IncrementAllowedFails {
id: String,
suffix: Option<String>,
ttl: f64,
},
/// Answered with `None`.
AddUsage {
window: UsageWindow,
tokens: f64,
requests: f64,
},
/// Answered with `None`.
AddRequest { window: UsageWindow, ttl: f64 },
/// Answered with `[tokens, requests]`, each a number or `None`.
Usage { window: UsageWindow },
/// Answered with the health checks' states, or `None`.
HealthStates,
/// Answered with the local value, or `None`.
GetLocal { key: String },
/// Answered with the local counter after the increment.
IncrementLocal {
key: String,
amount: f64,
ttl: Option<f64>,
},
}
/// The `op` of a `cache` step.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)]
pub(crate) enum CacheOp {
/// `async_get_cache(key)`: answered with the stored JSON value, or `None`.
#[strum(serialize = "get")]
Get,
/// `async_set_cache(key, value, ttl)`: answered with `None`.
#[strum(serialize = "set")]
Set,
/// `async_increment(key, value, ttl)`: answered with the new number. The TTL applies when
/// the key is created.
#[strum(serialize = "increment")]
Increment,
/// `async_increment_usage(keys, value, ttl)`: `value` is one amount per key, added as
/// `_increment_deployment_usage` adds them. Answered with `None`.
#[strum(serialize = "increment_usage")]
IncrementUsage,
/// `async_batch_get_cache(keys)`: answered with one JSON value or `None` per key, in order.
#[strum(serialize = "batch_get")]
BatchGet,
/// `async_delete_cache(key)`: answered with `None`.
#[strum(serialize = "delete")]
Delete,
#[strum(serialize = "cooldowns")]
Cooldowns,
#[strum(serialize = "put_cooldown")]
PutCooldown,
#[strum(serialize = "increment_allowed_fails")]
IncrementAllowedFails,
#[strum(serialize = "add_usage")]
AddUsage,
#[strum(serialize = "add_request")]
AddRequest,
#[strum(serialize = "usage")]
Usage,
#[strum(serialize = "health_states")]
HealthStates,
#[strum(serialize = "get_local")]
GetLocal,
#[strum(serialize = "increment_local")]
IncrementLocal,
}
impl CacheCall {
pub(crate) fn op(&self) -> CacheOp {
match self {
Self::Cooldowns { .. } => CacheOp::Cooldowns,
Self::PutCooldown { .. } => CacheOp::PutCooldown,
Self::IncrementAllowedFails { .. } => CacheOp::IncrementAllowedFails,
Self::AddUsage { .. } => CacheOp::AddUsage,
Self::AddRequest { .. } => CacheOp::AddRequest,
Self::Usage { .. } => CacheOp::Usage,
Self::HealthStates => CacheOp::HealthStates,
Self::GetLocal { .. } => CacheOp::GetLocal,
Self::IncrementLocal { .. } => CacheOp::IncrementLocal,
}
}
/// How many entries a `Values` answer holds.
fn values_len(&self) -> Option<usize> {
match self {
Self::Cooldowns { ids } => Some(ids.len()),
Self::Usage { .. } => Some(2),
_ => None,
}
}
}
/// Python's outcome of a `cache` step.
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum CacheOutcome {
/// `get`: the stored value (`None` when missing).
Value(Option<Value>),
/// `batch_get`: one entry per key.
Values(Vec<Option<Value>>),
/// `increment`: the counter after the increment.
Number(f64),
/// `set`, `delete` and `increment_usage`.
Stored,
}
impl CacheOutcome {
/// The op this outcome answers (`Stored` answers both `set` and `delete`).
fn answers(&self, op: CacheOp) -> bool {
matches!(
(self, op),
(Self::Value(_), CacheOp::Get)
| (Self::Values(_), CacheOp::BatchGet)
| (Self::Number(_), CacheOp::Increment)
(Self::Value(_), CacheOp::HealthStates | CacheOp::GetLocal)
| (Self::Values(_), CacheOp::Cooldowns | CacheOp::Usage)
| (
Self::Number(_),
CacheOp::IncrementAllowedFails | CacheOp::IncrementLocal
)
| (
Self::Stored,
CacheOp::Set | CacheOp::Delete | CacheOp::IncrementUsage
CacheOp::PutCooldown | CacheOp::AddUsage | CacheOp::AddRequest
)
)
}
@ -219,18 +280,7 @@ pub(crate) enum HostCall {
fallback_depth: u32,
retry: u32,
},
/// One [`RouterCache`] op against the host's cache. `key` is set for every op but
/// `batch_get`, which sets `keys`; `value` is set for `set` (the JSON to store) and
/// `increment` (the amount); `ttl` is in seconds.
Cache {
op: CacheOp,
key: Option<String>,
keys: Option<Vec<String>>,
value: Option<Value>,
ttl: Option<f64>,
/// The op reads or writes this worker's in-memory tier only (`local_only=True`).
local: bool,
},
Cache(CacheCall),
/// The deployment build's `os.environ/` references: answered with one value per ref.
Secrets {
index: usize,
@ -263,7 +313,7 @@ impl HostCall {
Self::RefuseFallbacks { .. } => StepKind::RefuseFallbacks,
Self::Select { .. } => StepKind::Select,
Self::PreRoute { .. } => StepKind::PreRoute,
Self::Cache { .. } => StepKind::Cache,
Self::Cache(_) => StepKind::Cache,
Self::Secrets { .. } => StepKind::Secrets,
Self::Build { kind, .. } => *kind,
Self::ModelInfo { .. } => StepKind::ModelInfo,
@ -373,15 +423,17 @@ impl HostCallResponse {
got: values.len(),
});
}
if let (Self::Cache(outcome), HostCall::Cache { op, keys, .. }) = (self, call) {
if !outcome.answers(*op) {
return Err(DriveError::WrongCacheOutcome { op: *op });
if let (Self::Cache(outcome), HostCall::Cache(cache)) = (self, call) {
let op = cache.op();
if !outcome.answers(op) {
return Err(DriveError::WrongCacheOutcome { op });
}
if let (CacheOutcome::Values(values), Some(keys)) = (outcome, keys)
&& values.len() != keys.len()
if let (CacheOutcome::Values(values), Some(expected)) = (outcome, cache.values_len())
&& values.len() != expected
{
return Err(DriveError::BatchLength {
expected: keys.len(),
op,
expected,
got: values.len(),
});
}
@ -447,8 +499,12 @@ pub(crate) enum DriveError {
UnknownCandidate { id: String },
#[error("the outcome does not answer a `{op}` cache step")]
WrongCacheOutcome { op: CacheOp },
#[error("a `batch_get` cache step for {expected} keys was answered with {got} values")]
BatchLength { expected: usize, got: usize },
#[error("a `{op}` cache step for {expected} entries was answered with {got} values")]
BatchLength {
op: CacheOp,
expected: usize,
got: usize,
},
#[error("a `secrets` step for {expected} references was answered with {got} values")]
SecretsLength { expected: usize, got: usize },
#[error("the router call is pending on something other than a step")]
@ -539,7 +595,7 @@ fn ttl_seconds(ttl: Option<Duration>) -> Option<f64> {
}
/// An increment amount as Python's counters store it: an int when it is whole.
fn amount_value(amount: f64) -> Value {
pub(crate) fn amount_value(amount: f64) -> Value {
if amount.fract() == 0.0 && amount.abs() < i64::MAX as f64 {
Value::from(amount as i64)
} else {
@ -547,131 +603,139 @@ fn amount_value(amount: f64) -> Value {
}
}
/// The host's cache: every op is a `cache` step Python runs against the router's `DualCache`.
/// The host's cache: every op is a `cache` step Python answers from the router's stores.
pub(crate) struct StepCache;
async fn cache_step(
op: CacheOp,
key: Option<&str>,
keys: Option<&[String]>,
value: Option<Value>,
ttl: Option<Duration>,
) -> CacheOutcome {
call_cache(op, key, keys, value, ttl, false).await
}
async fn call_cache(
op: CacheOp,
key: Option<&str>,
keys: Option<&[String]>,
value: Option<Value>,
ttl: Option<Duration>,
local: bool,
) -> CacheOutcome {
let call = HostCall::Cache {
op,
key: key.map(str::to_owned),
keys: keys.map(<[String]>::to_vec),
value,
ttl: ttl_seconds(ttl),
local,
};
let HostCallResponse::Cache(outcome) = call_host(call).await else {
async fn cache_step(call: CacheCall) -> CacheOutcome {
let HostCallResponse::Cache(outcome) = call_host(HostCall::Cache(call)).await else {
unreachable!("{CHECKED}");
};
outcome
}
async fn stored(call: CacheCall) -> Result<(), litellm_cache::Error> {
let CacheOutcome::Stored = cache_step(call).await else {
unreachable!("{CHECKED}");
};
Ok(())
}
async fn counted(call: CacheCall) -> Result<f64, litellm_cache::Error> {
let CacheOutcome::Number(next) = cache_step(call).await else {
unreachable!("{CHECKED}");
};
Ok(next)
}
async fn value(call: CacheCall) -> Result<Option<Value>, litellm_cache::Error> {
let CacheOutcome::Value(value) = cache_step(call).await else {
unreachable!("{CHECKED}");
};
Ok(value)
}
async fn values(call: CacheCall) -> Result<Vec<Option<Value>>, litellm_cache::Error> {
let CacheOutcome::Values(values) = cache_step(call).await else {
unreachable!("{CHECKED}");
};
Ok(values)
}
#[async_trait]
impl RouterCache for StepCache {
async fn async_get_cache(&self, key: &str) -> Result<Option<Value>, litellm_cache::Error> {
let CacheOutcome::Value(value) =
cache_step(CacheOp::Get, Some(key), None, None, None).await
else {
unreachable!("{CHECKED}");
};
Ok(value)
async fn cooldowns(&self, ids: &[String]) -> Result<Vec<Option<Value>>, litellm_cache::Error> {
values(CacheCall::Cooldowns { ids: ids.to_vec() }).await
}
async fn async_set_cache(
async fn put_cooldown(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
id: &str,
cooldown: Value,
ttl: Duration,
) -> Result<(), litellm_cache::Error> {
cache_step(CacheOp::Set, Some(key), None, Some(value), ttl).await;
Ok(())
stored(CacheCall::PutCooldown {
id: id.to_owned(),
cooldown,
ttl: ttl.as_secs_f64(),
})
.await
}
async fn async_increment(
async fn increment_allowed_fails(
&self,
id: &str,
suffix: Option<&str>,
ttl: Duration,
) -> Result<f64, litellm_cache::Error> {
counted(CacheCall::IncrementAllowedFails {
id: id.to_owned(),
suffix: suffix.map(str::to_owned),
ttl: ttl.as_secs_f64(),
})
.await
}
async fn add_usage(
&self,
window: &UsageWindow,
tokens: f64,
requests: f64,
) -> Result<(), litellm_cache::Error> {
stored(CacheCall::AddUsage {
window: window.clone(),
tokens,
requests,
})
.await
}
async fn add_request(
&self,
window: &UsageWindow,
ttl: Duration,
) -> Result<(), litellm_cache::Error> {
stored(CacheCall::AddRequest {
window: window.clone(),
ttl: ttl.as_secs_f64(),
})
.await
}
async fn usage(&self, window: &UsageWindow) -> Result<UsageCounts, litellm_cache::Error> {
let read = values(CacheCall::Usage {
window: window.clone(),
})
.await?;
let count = |index: usize| read.get(index).cloned().flatten().and_then(|v| v.as_f64());
Ok(UsageCounts {
tokens: count(0),
requests: count(1),
})
}
async fn health_states(&self) -> Result<Option<Value>, litellm_cache::Error> {
value(CacheCall::HealthStates).await
}
async fn get_local(&self, key: &str) -> Result<Option<Value>, litellm_cache::Error> {
value(CacheCall::GetLocal {
key: key.to_owned(),
})
.await
}
async fn increment_local(
&self,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, litellm_cache::Error> {
let value = Some(amount_value(amount));
let CacheOutcome::Number(next) =
cache_step(CacheOp::Increment, Some(key), None, value, ttl).await
else {
unreachable!("{CHECKED}");
};
Ok(next)
}
async fn async_batch_get_cache(
&self,
keys: &[String],
) -> Result<Vec<Option<Value>>, litellm_cache::Error> {
let CacheOutcome::Values(values) =
cache_step(CacheOp::BatchGet, None, Some(keys), None, None).await
else {
unreachable!("{CHECKED}");
};
Ok(values)
}
async fn async_delete_cache(&self, key: &str) -> Result<(), litellm_cache::Error> {
cache_step(CacheOp::Delete, Some(key), None, None, None).await;
Ok(())
}
async fn async_increment_usage(
&self,
increments: &[(String, f64)],
ttl: Option<Duration>,
) -> Result<(), litellm_cache::Error> {
let keys: Vec<String> = increments.iter().map(|(key, _)| key.clone()).collect();
let amounts = increments
.iter()
.map(|(_, amount)| amount_value(*amount))
.collect();
let value = Some(Value::Array(amounts));
cache_step(CacheOp::IncrementUsage, None, Some(&keys), value, ttl).await;
Ok(())
}
async fn async_get_local(&self, key: &str) -> Result<Option<Value>, litellm_cache::Error> {
let CacheOutcome::Value(value) =
call_cache(CacheOp::Get, Some(key), None, None, None, true).await
else {
unreachable!("{CHECKED}");
};
Ok(value)
}
async fn async_increment_local(
&self,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, litellm_cache::Error> {
let value = Some(amount_value(amount));
let CacheOutcome::Number(next) =
call_cache(CacheOp::Increment, Some(key), None, value, ttl, true).await
else {
unreachable!("{CHECKED}");
};
Ok(next)
counted(CacheCall::IncrementLocal {
key: key.to_owned(),
amount,
ttl: ttl_seconds(ttl),
})
.await
}
}
@ -1129,13 +1193,12 @@ impl Driver {
#[cfg(test)]
mod tests {
use std::{
collections::HashMap,
time::{SystemTime, UNIX_EPOCH},
};
use std::time::{SystemTime, UNIX_EPOCH};
use futures_util::FutureExt;
use litellm_router_core::{
FailureKind, Op, Router, RouterConfig, RouterDynamicConfig, SeededRandomness,
FailureKind, InMemoryRouterCache, Op, Router, RouterConfig, RouterDynamicConfig,
SeededRandomness,
};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
@ -1173,61 +1236,68 @@ mod tests {
HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::RateLimit, Some(429))))
}
/// A dict-backed stand-in for Python's `DualCache` that answers `cache` steps and keeps
/// every one it answered.
/// Router-core's in-memory cache standing in for Python's stores: it answers `cache` steps
/// and keeps every one it answered.
#[derive(Default)]
struct Store {
entries: HashMap<String, Value>,
cache: InMemoryRouterCache,
calls: Vec<HostCall>,
}
impl Store {
fn respond(&mut self, call: &HostCall) -> HostCallResponse {
let HostCall::Cache {
op,
key,
keys,
value,
..
} = call
else {
let HostCall::Cache(cache_call) = call else {
panic!("not a cache step: {call:?}");
};
self.calls.push(call.clone());
let key = key.clone().unwrap_or_default();
HostCallResponse::Cache(match op {
CacheOp::Get => CacheOutcome::Value(self.entries.get(&key).cloned()),
CacheOp::Set => {
self.entries.insert(key, value.clone().unwrap());
CacheOutcome::Stored
}
CacheOp::Increment => {
let current = self.entries.get(&key).and_then(Value::as_f64);
let next =
current.unwrap_or(0.0) + value.as_ref().and_then(Value::as_f64).unwrap();
self.entries.insert(key, amount_value(next));
CacheOutcome::Number(next)
}
CacheOp::IncrementUsage => {
let amounts = value.as_ref().and_then(Value::as_array).unwrap();
for (key, amount) in keys.iter().flatten().zip(amounts) {
let current = self.entries.get(key).and_then(Value::as_f64);
let next = current.unwrap_or(0.0) + amount.as_f64().unwrap();
self.entries.insert(key.clone(), amount_value(next));
let cache = &self.cache;
let seconds = Duration::from_secs_f64;
let outcome = async {
Ok::<_, litellm_cache::Error>(match cache_call {
CacheCall::Cooldowns { ids } => {
CacheOutcome::Values(cache.cooldowns(ids).await?)
}
CacheOutcome::Stored
}
CacheOp::BatchGet => CacheOutcome::Values(
keys.iter()
.flatten()
.map(|key| self.entries.get(key).cloned())
.collect(),
),
CacheOp::Delete => {
self.entries.remove(&key);
CacheOutcome::Stored
}
})
CacheCall::PutCooldown { id, cooldown, ttl } => {
cache
.put_cooldown(id, cooldown.clone(), seconds(*ttl))
.await?;
CacheOutcome::Stored
}
CacheCall::IncrementAllowedFails { id, suffix, ttl } => CacheOutcome::Number(
cache
.increment_allowed_fails(id, suffix.as_deref(), seconds(*ttl))
.await?,
),
CacheCall::AddUsage {
window,
tokens,
requests,
} => {
cache.add_usage(window, *tokens, *requests).await?;
CacheOutcome::Stored
}
CacheCall::AddRequest { window, ttl } => {
cache.add_request(window, seconds(*ttl)).await?;
CacheOutcome::Stored
}
CacheCall::Usage { window } => {
let counts = cache.usage(window).await?;
CacheOutcome::Values(
[counts.tokens, counts.requests]
.map(|count| count.map(Value::from))
.to_vec(),
)
}
CacheCall::HealthStates => CacheOutcome::Value(cache.health_states().await?),
CacheCall::GetLocal { key } => CacheOutcome::Value(cache.get_local(key).await?),
CacheCall::IncrementLocal { key, amount, ttl } => CacheOutcome::Number(
cache
.increment_local(key, *amount, ttl.map(seconds))
.await?,
),
})
};
HostCallResponse::Cache(outcome.now_or_never().unwrap().unwrap())
}
/// Resume with `response`, answering every cache step on the way, up to the next other
@ -1283,18 +1353,10 @@ mod tests {
fn cooldown_reads_are_cache_steps(pair: Router) {
let mut driver = start(&pair, "primary");
let step = driver.resume(None).unwrap();
let Step::HostCall(HostCall::Cache {
op: CacheOp::BatchGet,
key: None,
keys: Some(keys),
value: None,
ttl: None,
local: false,
}) = step
else {
let Step::HostCall(HostCall::Cache(CacheCall::Cooldowns { ids })) = step else {
panic!("the call should read cooldowns first, got {step:?}");
};
assert_eq!(keys, ["deployment:a1:cooldown", "deployment:a2:cooldown"]);
assert_eq!(ids, ["a1", "a2"]);
}
#[rstest]
@ -1343,10 +1405,16 @@ mod tests {
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs_f64();
store.entries.insert(
"deployment:a1:cooldown".into(),
json!({"exception_received": "slow down", "status_code": "429", "timestamp": now, "cooldown_time": 60}),
);
store
.cache
.put_cooldown(
"a1",
json!({"exception_received": "slow down", "status_code": "429", "timestamp": now, "cooldown_time": 60}),
Duration::from_secs(60),
)
.now_or_never()
.unwrap()
.unwrap();
for _ in 0..8 {
let mut call = start(&pair_without_retries, "primary");
let Step::HostCall(HostCall::Attempt { deployment_id, .. }) =
@ -1387,7 +1455,7 @@ mod tests {
.calls
.iter()
.filter(|call| {
matches!(call, HostCall::Cache { op, .. } if !matches!(op, CacheOp::Get | CacheOp::BatchGet))
matches!(call, HostCall::Cache(cache) if !matches!(cache, CacheCall::Cooldowns { .. } | CacheCall::Usage { .. } | CacheCall::HealthStates | CacheCall::GetLocal { .. }))
})
.collect();
assert!(!writes.is_empty(), "the core wrote no outcome");
@ -1482,13 +1550,7 @@ mod tests {
let mut store = Store::default();
let mut driver = start(&pair, "primary");
let step = driver.resume(None).unwrap();
let Step::HostCall(
call @ HostCall::Cache {
op: CacheOp::BatchGet,
..
},
) = &step
else {
let Step::HostCall(call @ HostCall::Cache(CacheCall::Cooldowns { .. })) = &step else {
panic!("expected the cooldown read, got {step:?}");
};
let error = driver.resume(Some(response)).unwrap_err();
@ -1596,10 +1658,7 @@ mod tests {
assert_eq!(model_group, "primary");
assert!(matches!(
store.calls.as_slice(),
[HostCall::Cache {
op: CacheOp::BatchGet,
..
}]
[HostCall::Cache(CacheCall::Cooldowns { .. })]
));
}

View file

@ -10,7 +10,7 @@ mod stream;
use std::{
sync::{Arc, Mutex, PoisonError},
time::{Duration, SystemTime},
time::Duration,
};
use litellm_error::Report;
@ -19,9 +19,8 @@ use litellm_router_core::lookup::{ListScope, Lookup, Row};
use litellm_router_core::{
AttemptError, Blocked, BuildPath, Built, Deployment, DeploymentSpec, Error, FailureKind,
FailureReport, Op, OutcomeCallbacks, PreRouted, Provided, RetryCause, Router, RouterEvent,
SecretAt, SuccessReport,
SecretAt, SuccessReport, UsageWindow,
build::Install,
health::{UsageMetric, current_minute, usage_key},
registry::{DeploymentFilter, ResolveKind, ResolveScope},
strategy::routing_context,
};
@ -32,8 +31,10 @@ use pyo3::{
};
use serde_json::Value;
use litellm_redis_store::router_usage::{Counter, Kind};
use self::host::{
CacheOp, CacheOutcome, DriveError, Driver, Finish, HostCall, HostCallResponse,
CacheCall, CacheOp, CacheOutcome, DriveError, Driver, Finish, HostCall, HostCallResponse,
PythonRandomness, Stage, Step, StepHost, StepKind,
};
use crate::errors::RustBridgeDeclined;
@ -384,13 +385,19 @@ impl RouterCoreHandle {
.collect()
}
/// The routing-cache key deployment `id` serving litellm `model` counts `metric` (`tpm`,
/// `rpm`, `itpm` or `otpm`) under in `minute` (`HH-MM`).
/// The key deployment `id` serving litellm `model` counts `metric` (`tpm`, `rpm`, `itpm` or
/// `otpm`) under in `minute` (`HH-MM`), as the usage store spells it.
fn usage_key(&self, id: &str, model: &str, metric: &str, minute: &str) -> PyResult<String> {
let metric: UsageMetric = metric
let kind: Kind = metric
.parse()
.map_err(|_| PyValueError::new_err(format!("`{metric}` is not a usage metric")))?;
Ok(usage_key(id, model, metric, minute))
Ok(Counter {
deployment_id: id.to_owned(),
model: Some(model.to_owned()),
minute: minute.to_owned(),
kind,
}
.key())
}
/// Keep the providers `other`'s build resolved for the deployments this router serves too.
@ -457,7 +464,8 @@ impl RouterCoreHandle {
}
/// Record a success a logging callback saw, driven like `report_failure`; it ends with
/// `done` (`key`: the tokens counter's key when usage was added, else None).
/// `done` (`key`: the tokens counter's key when usage was added, else None, as
/// `deployment_callback_on_success` returns it).
/// `usage_model` is the model its usage is keyed by, and `tracks_usage` the host's answer to
/// whether it counts any (see `Output::tracks_usage`). `callbacks`: see `report_failure`.
#[pyo3(signature = (id, total_tokens, requests, usage_model=None, tracks_usage=None, callbacks=None))]
@ -481,7 +489,9 @@ impl RouterCoreHandle {
router
.report_success(&id, total_tokens, requests, report)
.await
.map(|key| Finish::Counted { key })
.map(|window| Finish::Counted {
key: window.map(tokens_key),
})
}))
}
@ -504,26 +514,6 @@ impl RouterCoreHandle {
}))
}
/// The usage keys of this minute for every deployment of `model_group`:
/// `[(id, tpm_key, rpm_key)]`, or None when the group has no deployment.
fn usage_keys(&self, model_group: &str) -> Option<Vec<(String, String, String)>> {
let minute = current_minute(SystemTime::now());
let group = self.router.registry().group(model_group)?;
Some(
group
.iter()
.map(|deployment| {
let model = &deployment.params.model;
(
deployment.id.clone(),
usage_key(&deployment.id, model, UsageMetric::Tpm, &minute),
usage_key(&deployment.id, model, UsageMetric::Rpm, &minute),
)
})
.collect(),
)
}
/// Every configured deployment, as `model_list` entries with their ids, in config order.
fn deployments<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyList>> {
let registry = self.router.registry();
@ -1237,7 +1227,7 @@ fn host_call_response(
result => HostCallResponse::Attempt(result),
})
}
HostCall::Cache { op, .. } => cache_outcome(*op, outcome).map(HostCallResponse::Cache),
HostCall::Cache(call) => cache_outcome(call.op(), outcome).map(HostCallResponse::Cache),
HostCall::Secrets { .. } => Ok(HostCallResponse::Secrets(from_py_argument(
&outcome.ok_or_else(missing)?,
)?)),
@ -1307,36 +1297,39 @@ fn pre_route_outcome(outcome: Option<Bound<'_, PyAny>>) -> PyResult<Option<PreRo
Ok(Some(PreRouted::Deployment(Arc::new(deployment))))
}
/// `get`: the JSON value or None; `batch_get`: a list of JSON values or None, one per key;
/// `increment`: the new int or float; `set` and `delete`: None.
/// `cooldowns` and `usage`: a list of JSON values or None, one per entry; `health_states` and
/// `get_local`: the JSON value or None; `increment_allowed_fails` and `increment_local`: the new
/// int or float; the writes: None.
fn cache_outcome(op: CacheOp, outcome: Option<Bound<'_, PyAny>>) -> PyResult<CacheOutcome> {
let outcome = outcome.filter(|outcome| !outcome.is_none());
let json = |value: Option<Value>| value.filter(|value| !value.is_null());
match op {
CacheOp::Get => outcome
CacheOp::HealthStates | CacheOp::GetLocal => outcome
.map(|outcome| from_py_argument::<Value>(&outcome))
.transpose()
.map(|value| CacheOutcome::Value(json(value))),
CacheOp::BatchGet => {
CacheOp::Cooldowns | CacheOp::Usage => {
let outcome = outcome.ok_or_else(|| {
PyValueError::new_err("a `batch_get` cache outcome must be a list, not None")
PyValueError::new_err(format!("a `{op}` cache outcome must be a list, not None"))
})?;
if !outcome.is_instance_of::<PyList>() {
return Err(PyValueError::new_err(
"a `batch_get` cache outcome must be a list with one entry per key",
));
return Err(PyValueError::new_err(format!(
"a `{op}` cache outcome must be a list with one entry per value asked for"
)));
}
let values: Vec<Option<Value>> = from_py_argument(&outcome)?;
Ok(CacheOutcome::Values(values.into_iter().map(json).collect()))
}
CacheOp::Increment => outcome
CacheOp::IncrementAllowedFails | CacheOp::IncrementLocal => outcome
.filter(|outcome| !outcome.is_instance_of::<PyBool>())
.and_then(|outcome| outcome.extract::<f64>().ok())
.map(CacheOutcome::Number)
.ok_or_else(|| {
PyValueError::new_err("an `increment` cache outcome must be the new int or float")
PyValueError::new_err(format!(
"an `{op}` cache outcome must be the new int or float"
))
}),
CacheOp::Set | CacheOp::Delete | CacheOp::IncrementUsage => match outcome {
CacheOp::PutCooldown | CacheOp::AddUsage | CacheOp::AddRequest => match outcome {
None => Ok(CacheOutcome::Stored),
Some(_) => Err(PyValueError::new_err(format!(
"a `{op}` cache step has no outcome, so pass None"
@ -1345,6 +1338,69 @@ fn cache_outcome(op: CacheOp, outcome: Option<Bound<'_, PyAny>>) -> PyResult<Cac
}
}
/// The tokens counter of `window`, as the usage store keys it.
fn tokens_key(window: UsageWindow) -> String {
Counter {
deployment_id: window.deployment_id,
model: Some(window.model),
minute: window.minute,
kind: Kind::Tpm,
}
.key()
}
/// A [`UsageWindow`] as the `window` of a `cache` step.
fn window_dict<'py>(py: Python<'py>, window: UsageWindow) -> PyResult<Bound<'py, PyDict>> {
let dict = PyDict::new(py);
dict.set_item("deployment_id", window.deployment_id)?;
dict.set_item("model", window.model)?;
dict.set_item("minute", window.minute)?;
Ok(dict)
}
fn cache_step_dict(py: Python<'_>, dict: &Bound<'_, PyDict>, call: CacheCall) -> PyResult<()> {
dict.set_item("op", <&'static str>::from(call.op()))?;
let amount = |amount: f64| to_py(py, &host::amount_value(amount));
match call {
CacheCall::Cooldowns { ids } => dict.set_item("ids", PyList::new(py, ids)?),
CacheCall::PutCooldown { id, cooldown, ttl } => {
dict.set_item("id", id)?;
dict.set_item("cooldown", to_py(py, &cooldown)?)?;
dict.set_item("ttl", ttl)
}
CacheCall::IncrementAllowedFails { id, suffix, ttl } => {
dict.set_item("id", id)?;
dict.set_item("suffix", suffix)?;
dict.set_item("ttl", ttl)
}
CacheCall::AddUsage {
window,
tokens,
requests,
} => {
dict.set_item("window", window_dict(py, window)?)?;
dict.set_item("tokens", amount(tokens)?)?;
dict.set_item("requests", amount(requests)?)
}
CacheCall::AddRequest { window, ttl } => {
dict.set_item("window", window_dict(py, window)?)?;
dict.set_item("ttl", ttl)
}
CacheCall::Usage { window } => dict.set_item("window", window_dict(py, window)?),
CacheCall::HealthStates => Ok(()),
CacheCall::GetLocal { key } => dict.set_item("key", key),
CacheCall::IncrementLocal {
key,
amount: by,
ttl,
} => {
dict.set_item("key", key)?;
dict.set_item("amount", amount(by)?)?;
dict.set_item("ttl", ttl)
}
}
}
/// `{"ok": True, "total_tokens": int}` (an attempt's may add `"blocked": str`, `"deferred": bool`
/// and `"tracks_usage": bool`) or `{"ok": False, "kind": str, "status": int | None,
/// "retry_after": float | None, "num_retries": int | None, "detail": str | None,
@ -1804,21 +1860,7 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) ->
dict.set_item("fallback_depth", fallback_depth)?;
dict.set_item("retry", retry)
}
HostCall::Cache {
op,
key,
keys,
value,
ttl,
local,
} => {
dict.set_item("op", <&'static str>::from(op))?;
dict.set_item("key", key)?;
dict.set_item("keys", keys.map(|keys| PyList::new(py, keys)).transpose()?)?;
dict.set_item("value", value.map(|value| to_py(py, &value)).transpose()?)?;
dict.set_item("ttl", ttl)?;
dict.set_item("local", local)
}
HostCall::Cache(call) => cache_step_dict(py, dict, call),
HostCall::Secrets { index, refs } => {
dict.set_item("index", index)?;
let refs = refs
@ -1973,22 +2015,39 @@ import time
store = {}
asked = []
def add(key, amount):
store[key] = store.get(key, 0) + amount
return store[key]
def window(step):
w = step["window"]
return (w["deployment_id"], w["model"], w["minute"])
def answer_cache(step):
asked.append(step)
op, key = step["op"], step["key"]
if op == "get":
return store.get(key)
if op == "batch_get":
return [store.get(k) for k in step["keys"]]
if op == "set":
store[key] = step["value"]
op = step["op"]
if op == "cooldowns":
return [store.get(("cooldown", id)) for id in step["ids"]]
if op == "put_cooldown":
store[("cooldown", step["id"])] = step["cooldown"]
return None
if op == "increment":
store[key] = store.get(key, 0) + step["value"]
return store[key]
if op == "delete":
store.pop(key, None)
if op == "increment_allowed_fails":
return add(("allowed_fails", step["id"], step["suffix"]), 1)
if op == "add_usage":
add(("tokens", *window(step)), step["tokens"])
add(("requests", *window(step)), step["requests"])
return None
if op == "add_request":
add(("requests", *window(step)), 1)
return None
if op == "usage":
return [store.get(("tokens", *window(step))), store.get(("requests", *window(step)))]
if op == "health_states":
return None
if op == "get_local":
return store.get(("local", step["key"]))
if op == "increment_local":
return add(("local", step["key"]), step["amount"])
raise AssertionError(op)
def drive(call, outcome):
@ -2000,8 +2059,9 @@ def drive(call, outcome):
handle = native.RouterCoreHandle(config, {"num_retries": 1})
picked = drive(handle.select("primary", {}), None)
assert (picked["kind"], picked["deployment_id"], picked["model_group"]) == ("done", "a1", "primary"), picked
assert [step["op"] for step in asked] == ["batch_get"]
assert list(asked[0]) == ["kind", "op", "key", "keys", "value", "ttl", "local", "events"], list(asked[0])
assert [step["op"] for step in asked] == ["cooldowns"]
assert list(asked[0]) == ["kind", "op", "ids", "events"], list(asked[0])
assert asked[0]["ids"] == ["a1"], asked[0]
call = handle.start("acompletion", '{"model": "primary"}')
version = handle.registry_version()
assert handle.apply({"delete": ["gone"]}) == version + 1
@ -2023,20 +2083,20 @@ assert cooled.resolve("primary")["ids"] == ["a1", "a2"]
assert cooled.resolve("nope") is None
assert drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})["kind"] == "failed"
# The core records the failure itself: the 429 cools the deployment down in the host's store.
assert f"deployment:{tried}:cooldown" in store, store
assert ("cooldown", tried) in store, store
other = drive(cooled.select("primary", {}), None)
assert other["kind"] == "done" and other["deployment_id"] != tried, other
call = handle.select("primary", {})
step = call.next(None)
assert step["op"] == "batch_get"
assert step["op"] == "cooldowns"
for bad in ([], [None, None], "nope", None):
try:
call.next(bad)
except ValueError:
pass
else:
raise AssertionError(f"{bad!r} should not answer a batch_get of one key")
raise AssertionError(f"{bad!r} should not answer a cooldowns read of one id")
assert drive(call, [None])["kind"] == "done"
"#
),

View file

@ -5,6 +5,12 @@ pub enum Kind {
Tpm,
#[strum(serialize = "rpm")]
Rpm,
/// Input tokens per minute, which the io-token rate limit check reads.
#[strum(serialize = "itpm")]
Itpm,
/// Output tokens per minute.
#[strum(serialize = "otpm")]
Otpm,
}
#[derive(Clone, Debug, PartialEq, Eq)]
@ -16,8 +22,9 @@ pub struct Counter {
}
impl Counter {
/// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`.
pub(super) fn key(&self) -> String {
/// `global_router:{id}:{model}:{kind}:{HH-MM}`, where Python formats a missing model as `None`.
/// Public only for the hosts that hand a counter's key back to their callers.
pub fn key(&self) -> String {
let kind: &str = self.kind.into();
let model = self.model.as_deref().unwrap_or("None");
format!(

View file

@ -1,6 +1,7 @@
//! [`RouterCache`]: the DualCache-style store all routing state lives in, under the keys,
//! encodings and TTLs the Python router uses, so Python and Rust workers share it. Method names
//! follow `litellm-cache`'s `BaseCache`, `BatchCache`, `CounterCache` and `DeleteCache`.
//! [`RouterCache`]: the routing state router-core reads and writes, named by what it is (a
//! deployment's cooldown, its allowed-fails count, its per-minute usage, the health checks'
//! states). Where that state lives in Redis, and under which keys, is the host's store's business:
//! router-core never spells a Redis key.
use std::{
collections::HashMap,
@ -13,90 +14,110 @@ use litellm_cache::Error;
use parking_lot::Mutex;
use serde_json::Value;
/// One deployment's per-minute usage counters: `model` is the model its usage is counted under,
/// and `minute` is `HH-MM` in UTC.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct UsageWindow {
pub deployment_id: String,
pub model: String,
pub minute: String,
}
/// A [`UsageWindow`]'s counts, `None` for a counter nothing has written yet.
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct UsageCounts {
pub tokens: Option<f64>,
pub requests: Option<f64>,
}
#[async_trait]
pub trait RouterCache: Send + Sync {
async fn async_get_cache(&self, key: &str) -> Result<Option<Value>, Error>;
/// The stored cooldown of each of `ids`, in order, `None` for one without.
async fn cooldowns(&self, ids: &[String]) -> Result<Vec<Option<Value>>, Error>;
async fn async_set_cache(
async fn put_cooldown(&self, id: &str, cooldown: Value, ttl: Duration) -> Result<(), Error>;
/// Count one more failure against `id`'s allowed fails (per `suffix` when one is given) and
/// return the count. `ttl` applies when the counter is created.
async fn increment_allowed_fails(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
id: &str,
suffix: Option<&str>,
ttl: Duration,
) -> Result<f64, Error>;
/// `_increment_deployment_usage`'s write: tokens and requests in one round trip.
async fn add_usage(
&self,
window: &UsageWindow,
tokens: f64,
requests: f64,
) -> Result<(), Error>;
/// Add `amount` to the number at `key` (0 when missing) and return the new value. The TTL
/// applies when the key is created, as Redis `INCRBYFLOAT` + `EXPIRE NX` does.
async fn async_increment(
/// One more request on `window`, whose counter expires `ttl` after it is created.
async fn add_request(&self, window: &UsageWindow, ttl: Duration) -> Result<(), Error>;
async fn usage(&self, window: &UsageWindow) -> Result<UsageCounts, Error>;
/// The background health checks' per-deployment states, as `DeploymentHealthCache` keeps them.
async fn health_states(&self) -> Result<Option<Value>, Error>;
/// Read this worker's in-memory tier only (`local_only=True`).
async fn get_local(&self, key: &str) -> Result<Option<Value>, Error>;
/// Add to a counter in this worker's in-memory tier only (`local_only=True`).
async fn increment_local(
&self,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, Error>;
async fn async_batch_get_cache(&self, keys: &[String]) -> Result<Vec<Option<Value>>, Error>;
async fn async_delete_cache(&self, key: &str) -> Result<(), Error>;
/// Read this worker's in-memory tier only (`local_only=True`).
async fn async_get_local(&self, key: &str) -> Result<Option<Value>, Error> {
self.async_get_cache(key).await
}
/// Add to a counter in this worker's in-memory tier only (`local_only=True`).
async fn async_increment_local(
&self,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
self.async_increment(key, amount, ttl).await
}
/// `_increment_deployment_usage`'s write: add each `(key, amount)` in one pipeline. With a
/// shared tier, this worker's in-memory tier then holds the shared totals, so its reads see
/// what other workers counted up to now.
async fn async_increment_usage(
&self,
increments: &[(String, f64)],
ttl: Option<Duration>,
) -> Result<(), Error> {
for (key, amount) in increments {
self.async_increment(key, *amount, ttl).await?;
}
Ok(())
}
}
#[derive(Default)]
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
enum Slot {
Cooldown(String),
AllowedFails(String, Option<String>),
Tokens(UsageWindow),
Requests(UsageWindow),
HealthStates,
Local(String),
}
struct Entry {
value: Value,
expires: Option<Instant>,
}
/// An in-process [`RouterCache`], for tests and for state Python keeps `local_only`.
/// An in-process [`RouterCache`], for tests.
#[derive(Clone, Default)]
pub struct InMemoryRouterCache {
entries: Arc<Mutex<HashMap<String, Entry>>>,
entries: Arc<Mutex<HashMap<Slot, Entry>>>,
}
impl InMemoryRouterCache {
fn live(entries: &mut HashMap<String, Entry>, key: &str) -> Option<Value> {
fn live(entries: &mut HashMap<Slot, Entry>, slot: &Slot) -> Option<Value> {
let expired = entries
.get(key)
.get(slot)
.is_some_and(|entry| entry.expires.is_some_and(|at| at <= Instant::now()));
if expired {
entries.remove(key);
entries.remove(slot);
}
entries.get(key).map(|entry| entry.value.clone())
entries.get(slot).map(|entry| entry.value.clone())
}
pub fn get(&self, key: &str) -> Option<Value> {
Self::live(&mut self.entries.lock(), key)
fn get(&self, slot: &Slot) -> Option<Value> {
Self::live(&mut self.entries.lock(), slot)
}
pub fn increment(&self, key: &str, amount: f64, ttl: Option<Duration>) -> f64 {
fn set(&self, slot: Slot, value: Value, ttl: Option<Duration>) {
let expires = ttl.map(|ttl| Instant::now() + ttl);
self.entries.lock().insert(slot, Entry { value, expires });
}
fn increment(&self, slot: Slot, amount: f64, ttl: Option<Duration>) -> f64 {
let mut entries = self.entries.lock();
let current = Self::live(&mut entries, key).and_then(|value| value.as_f64());
let current = Self::live(&mut entries, &slot).and_then(|value| value.as_f64());
let next = current.unwrap_or(0.0) + amount;
let value = serde_json::Number::from_f64(next)
.map(|number| {
@ -107,53 +128,92 @@ impl InMemoryRouterCache {
}
})
.unwrap_or(Value::Null);
let expires = match (current, entries.get(key)) {
let expires = match (current, entries.get(&slot)) {
(Some(_), Some(entry)) => entry.expires,
_ => ttl.map(|ttl| Instant::now() + ttl),
};
entries.insert(key.to_owned(), Entry { value, expires });
entries.insert(slot, Entry { value, expires });
next
}
/// `window`'s counts as stored now.
pub fn counted(&self, window: &UsageWindow) -> UsageCounts {
let count = |slot: Slot| self.get(&slot).and_then(|value| value.as_f64());
UsageCounts {
tokens: count(Slot::Tokens(window.clone())),
requests: count(Slot::Requests(window.clone())),
}
}
/// A local counter as stored now.
pub fn local(&self, key: &str) -> Option<Value> {
self.get(&Slot::Local(key.to_owned()))
}
pub fn set_health_states(&self, states: Value) {
self.set(Slot::HealthStates, states, None);
}
}
#[async_trait]
impl RouterCache for InMemoryRouterCache {
async fn async_get_cache(&self, key: &str) -> Result<Option<Value>, Error> {
Ok(self.get(key))
async fn cooldowns(&self, ids: &[String]) -> Result<Vec<Option<Value>>, Error> {
let mut entries = self.entries.lock();
Ok(ids
.iter()
.map(|id| Self::live(&mut entries, &Slot::Cooldown(id.clone())))
.collect())
}
async fn async_set_cache(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<(), Error> {
let expires = ttl.map(|ttl| Instant::now() + ttl);
self.entries
.lock()
.insert(key.to_owned(), Entry { value, expires });
async fn put_cooldown(&self, id: &str, cooldown: Value, ttl: Duration) -> Result<(), Error> {
self.set(Slot::Cooldown(id.to_owned()), cooldown, Some(ttl));
Ok(())
}
async fn async_increment(
async fn increment_allowed_fails(
&self,
id: &str,
suffix: Option<&str>,
ttl: Duration,
) -> Result<f64, Error> {
let slot = Slot::AllowedFails(id.to_owned(), suffix.map(str::to_owned));
Ok(self.increment(slot, 1.0, Some(ttl)))
}
async fn add_usage(
&self,
window: &UsageWindow,
tokens: f64,
requests: f64,
) -> Result<(), Error> {
self.increment(Slot::Tokens(window.clone()), tokens, None);
self.increment(Slot::Requests(window.clone()), requests, None);
Ok(())
}
async fn add_request(&self, window: &UsageWindow, ttl: Duration) -> Result<(), Error> {
self.increment(Slot::Requests(window.clone()), 1.0, Some(ttl));
Ok(())
}
async fn usage(&self, window: &UsageWindow) -> Result<UsageCounts, Error> {
Ok(self.counted(window))
}
async fn health_states(&self) -> Result<Option<Value>, Error> {
Ok(self.get(&Slot::HealthStates))
}
async fn get_local(&self, key: &str) -> Result<Option<Value>, Error> {
Ok(self.local(key))
}
async fn increment_local(
&self,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
Ok(self.increment(key, amount, ttl))
}
async fn async_batch_get_cache(&self, keys: &[String]) -> Result<Vec<Option<Value>>, Error> {
let mut entries = self.entries.lock();
Ok(keys
.iter()
.map(|key| Self::live(&mut entries, key))
.collect())
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
self.entries.lock().remove(key);
Ok(())
Ok(self.increment(Slot::Local(key.to_owned()), amount, ttl))
}
}

View file

@ -1,17 +1,10 @@
//! [`Health`]: cooldowns, allowed-fails counters and per-minute usage, kept in the shared
//! [`RouterCache`] under the Python router's keys, encodings and TTLs:
//!
//! | state | key | value | TTL |
//! |---|---|---|---|
//! | cooldown | `deployment:{id}:cooldown` | `{exception_received, status_code, timestamp, cooldown_time}` | cooldown time |
//! | allowed fails | `deployment:{id}:allowed_fails[:{suffix}]` | int counter | cooldown time |
//! | usage | `global_router:{id}:{model}:tpm:{HH-MM}` and `:rpm:` | int counter | 60s |
//! | outcomes (local only) | `{id}:successes`, `{id}:fails` | int counter | 60s |
//! [`Health`]: cooldowns, allowed-fails counters and per-minute usage, kept in the host's
//! [`RouterCache`] with the Python router's encodings and TTLs, and this worker's per-minute
//! outcomes (`{id}:successes`, `{id}:fails`, kept `local_only`).
//!
//! From `router_utils/cooldown_cache.py`, `router_utils/cooldown_handlers.py`,
//! `router_utils/router_callbacks/track_deployment_metrics.py`, `RouterCacheEnum` in
//! `types/router.py` and `Router._increment_usage` / `deployment_callback_on_failure` in
//! `router.py`.
//! `router_utils/router_callbacks/track_deployment_metrics.py` and `Router._increment_usage` /
//! `deployment_callback_on_failure` in `router.py`.
use std::{
sync::Arc,
@ -22,7 +15,7 @@ use serde::{Deserialize, Serialize};
use tracing::debug;
use crate::{
cache::RouterCache,
cache::{RouterCache, UsageWindow},
error::{AttemptError, Error, FailureKind},
outcome::{OutcomeCallbacks, OutcomeWrite},
registry::{Deployment, Registry},
@ -68,37 +61,6 @@ pub fn current_minute(time: SystemTime) -> String {
format!("{:02}-{:02}", seconds / 3600, (seconds % 3600) / 60)
}
pub fn cooldown_key(id: &str) -> String {
format!("deployment:{id}:cooldown")
}
pub fn allowed_fails_key(id: &str, suffix: Option<&str>) -> String {
match suffix {
Some(suffix) => format!("deployment:{id}:allowed_fails:{suffix}"),
None => format!("deployment:{id}:allowed_fails"),
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::EnumString)]
pub enum UsageMetric {
#[strum(serialize = "tpm")]
Tpm,
#[strum(serialize = "rpm")]
Rpm,
/// Input tokens per minute, which the io-token rate limit check counts.
#[strum(serialize = "itpm")]
Itpm,
/// Output tokens per minute.
#[strum(serialize = "otpm")]
Otpm,
}
/// `RouterCacheEnum.TPM` / `RPM` / `ITPM` / `OTPM`: `model` is the deployment's
/// `litellm_params.model`.
pub fn usage_key(id: &str, model: &str, metric: UsageMetric, minute: &str) -> String {
format!("global_router:{id}:{model}:{metric}:{minute}")
}
/// `CooldownCacheValue`.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CooldownValue {
@ -260,8 +222,7 @@ impl Health {
&self,
ids: &[String],
) -> Result<Vec<(String, CooldownValue)>, Error> {
let keys: Vec<String> = ids.iter().map(|id| cooldown_key(id)).collect();
let values = self.cache.async_batch_get_cache(&keys).await?;
let values = self.cache.cooldowns(ids).await?;
let now = self.now();
Ok(ids
.iter()
@ -287,16 +248,14 @@ impl Health {
cooldown_time: cooldown.as_secs_f64(),
};
let value = serde_json::to_value(value).unwrap_or_default();
self.cache
.async_set_cache(&cooldown_key(id), value, Some(cooldown))
.await?;
self.cache.put_cooldown(id, value, cooldown).await?;
Ok(())
}
async fn local_count(&self, key: &str) -> Result<u64, Error> {
Ok(self
.cache
.async_get_local(key)
.get_local(key)
.await?
.and_then(|value| value.as_f64())
.map_or(0, |count| count as u64))
@ -315,7 +274,7 @@ impl Health {
pub async fn count_success(&self, id: &str) -> Result<String, Error> {
let key = format!("{id}:successes");
self.cache
.async_increment_local(&key, 1.0, Some(USAGE_TTL))
.increment_local(&key, 1.0, Some(USAGE_TTL))
.await?;
Ok(key)
}
@ -323,7 +282,7 @@ impl Health {
/// `increment_deployment_failures_for_current_minute`.
pub async fn count_failure(&self, id: &str) -> Result<(), Error> {
self.cache
.async_increment_local(&format!("{id}:fails"), 1.0, Some(USAGE_TTL))
.increment_local(&format!("{id}:fails"), 1.0, Some(USAGE_TTL))
.await?;
Ok(())
}
@ -534,10 +493,7 @@ impl Health {
allowed: f64,
ttl: Duration,
) -> Result<bool, Error> {
let fails = self
.cache
.async_increment(&allowed_fails_key(id, suffix), 1.0, Some(ttl))
.await?;
let fails = self.cache.increment_allowed_fails(id, suffix, ttl).await?;
Ok(fails > allowed)
}
@ -568,34 +524,22 @@ impl Health {
/// `_increment_deployment_usage`, for a deployment that tracks usage: tokens and requests
/// under `model` (metadata's `deployment`, which a wildcard route fills; else the
/// deployment's own), unless there is nothing to add. Returns the tokens counter's key when
/// it added to it.
/// deployment's own), unless there is nothing to add. Returns the window it added to.
pub async fn add_usage(
&self,
deployment: &Deployment,
model: Option<&str>,
total_tokens: f64,
requests: f64,
) -> Result<Option<String>, Error> {
) -> Result<Option<UsageWindow>, Error> {
if total_tokens <= 0.0 && requests <= 0.0 {
return Ok(None);
}
let minute = current_minute(self.clock.now());
let model = model.unwrap_or(&deployment.params.model);
let tpm_key = usage_key(&deployment.id, model, UsageMetric::Tpm, &minute);
let window = self.usage_window(deployment, model);
self.cache
.async_increment_usage(
&[
(tpm_key.clone(), total_tokens),
(
usage_key(&deployment.id, model, UsageMetric::Rpm, &minute),
requests,
),
],
Some(USAGE_TTL),
)
.add_usage(&window, total_tokens, requests)
.await?;
Ok(Some(tpm_key))
Ok(Some(window))
}
/// `async_deployment_callback_on_failure`: a failed request counts against the deployment's
@ -605,38 +549,28 @@ impl Health {
deployment: &Deployment,
model: Option<&str>,
) -> Result<(), Error> {
let minute = current_minute(self.clock.now());
self.cache
.async_increment(
&usage_key(
&deployment.id,
model.unwrap_or(&deployment.params.model),
UsageMetric::Rpm,
&minute,
),
1.0,
Some(USAGE_TTL),
)
.await?;
let window = self.usage_window(deployment, model);
self.cache.add_request(&window, USAGE_TTL).await?;
Ok(())
}
/// This minute's (tpm, rpm) for a deployment.
pub async fn usage(&self, deployment: &Deployment) -> Result<(u64, u64), Error> {
let minute = current_minute(self.clock.now());
let model = &deployment.params.model;
let keys = [UsageMetric::Tpm, UsageMetric::Rpm]
.map(|metric| usage_key(&deployment.id, model, metric, &minute));
let values = self.cache.async_batch_get_cache(&keys).await?;
let count = |index: usize| {
values
.get(index)
.cloned()
.flatten()
.and_then(|value| value.as_f64())
.map_or(0, |count| count as u64)
};
Ok((count(0), count(1)))
let counts = self
.cache
.usage(&self.usage_window(deployment, None))
.await?;
let count = |value: Option<f64>| value.map_or(0, |count| count as u64);
Ok((count(counts.tokens), count(counts.requests)))
}
/// This minute's window for `deployment`, counted under `model` or the deployment's own.
fn usage_window(&self, deployment: &Deployment, model: Option<&str>) -> UsageWindow {
UsageWindow {
deployment_id: deployment.id.clone(),
model: model.unwrap_or(&deployment.params.model).to_owned(),
minute: current_minute(self.clock.now()),
}
}
}
@ -648,7 +582,7 @@ mod tests {
use super::*;
use crate::{
cache::InMemoryRouterCache,
cache::{InMemoryRouterCache, UsageCounts},
config::{RouterConfig, RouterDynamicConfig},
registry::{DeploymentSpec, RoutingGroupSpec},
};
@ -936,29 +870,17 @@ mod tests {
.await
.unwrap();
let minute = current_minute(clock.now());
let window = UsageWindow {
deployment_id: "a".to_owned(),
model: keyed.to_owned(),
minute: current_minute(clock.now()),
};
assert_eq!(
(
cache.get(&usage_key("a", keyed, UsageMetric::Tpm, &minute)),
cache.get(&usage_key("a", keyed, UsageMetric::Rpm, &minute)),
),
(Some(serde_json::json!(15)), Some(serde_json::json!(1))),
);
}
#[rstest]
fn keys_match_python() {
assert_eq!(
(
cooldown_key("abc"),
allowed_fails_key("abc", Some("generic")),
usage_key("abc", "openai/gpt-4o", UsageMetric::Tpm, "09-05"),
),
(
"deployment:abc:cooldown".to_owned(),
"deployment:abc:allowed_fails:generic".to_owned(),
"global_router:abc:openai/gpt-4o:tpm:09-05".to_owned(),
)
cache.counted(&window),
UsageCounts {
tokens: Some(15.0),
requests: Some(1.0),
},
);
}
}

View file

@ -1,6 +1,6 @@
//! Health-check-driven routing: the background health checks' results, kept in the shared
//! [`RouterCache`](crate::cache::RouterCache) under one key as Python's `DeploymentHealthCache`
//! writes them, and the filter `_async_filter_health_check_unhealthy_deployments` /
//! Health-check-driven routing: the background health checks' results, read from the host's
//! [`RouterCache`](crate::cache::RouterCache) as Python's `DeploymentHealthCache` writes them, and
//! the filter `_async_filter_health_check_unhealthy_deployments` /
//! `_filter_health_check_unhealthy_deployments` run before the cooldown filter.
use std::{collections::BTreeSet, sync::Arc, time::UNIX_EPOCH};
@ -9,9 +9,6 @@ use serde_json::Value;
use crate::{health::Health, registry::Deployment, settings::Settings};
/// `DeploymentHealthCache.CACHE_KEY`.
pub const HEALTH_STATE_KEY: &str = "litellm:health_check:deployment_health_state";
/// `_extract_unhealthy_ids`: the ids marked unhealthy whose result is younger than `staleness`.
pub fn unhealthy_ids(raw: Option<&Value>, now: f64, staleness: f64) -> BTreeSet<String> {
let Some(Value::Object(states)) = raw else {
@ -71,12 +68,7 @@ impl Health {
return deployments;
}
// A read that fails counts as no health state.
let raw = self
.cache()
.async_get_cache(HEALTH_STATE_KEY)
.await
.ok()
.flatten();
let raw = self.cache().health_states().await.ok().flatten();
let now = self
.now_time()
.duration_since(UNIX_EPOCH)

View file

@ -118,7 +118,7 @@ use arc_swap::ArcSwap;
use parking_lot::Mutex;
pub use crate::{
cache::{InMemoryRouterCache, RouterCache},
cache::{InMemoryRouterCache, RouterCache, UsageCounts, UsageWindow},
call::RouterResponse,
config::{Op, RouterConfig, RouterDynamicConfig, RoutingStrategy},
error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported},
@ -436,14 +436,14 @@ impl Router {
/// A success a logging callback saw (`deployment_callback_on_success`): counted for a
/// configured deployment, with its tokens and requests when it tracks usage. Returns the
/// tokens counter's key when it added usage (`_increment_deployment_usage`), else `None`.
/// window it added usage to (`_increment_deployment_usage`), else `None`.
pub async fn report_success(
&self,
id: &str,
total_tokens: f64,
requests: f64,
report: SuccessReport,
) -> Result<Option<String>, Error> {
) -> Result<Option<UsageWindow>, Error> {
report
.callbacks
.record(
@ -459,7 +459,7 @@ impl Router {
total_tokens: f64,
requests: f64,
report: &SuccessReport,
) -> Result<Option<String>, Error> {
) -> Result<Option<UsageWindow>, Error> {
let snapshot = self.core.snapshot.load_full();
let Some(deployment) = snapshot.registry.deployment(id) else {
return Ok(None);

View file

@ -8,12 +8,8 @@ use litellm_router_core::{
Attempt, AttemptError, Deployment, Error, FailureKind, FailureReport, Host,
InMemoryRouterCache, Op, Output, Payload, Router, RouterCache, RouterConfig,
RouterDynamicConfig, RouterEvent, RouterRequest, RoutingContext, SeededRandomness, SelectStage,
SuccessReport,
call::StreamOrigin,
fallback::Fallbacks,
health::{UsageMetric, current_minute, usage_key},
host::UncountedRetries,
select::PassThrough,
SuccessReport, UsageWindow, call::StreamOrigin, fallback::Fallbacks, health::current_minute,
host::UncountedRetries, select::PassThrough,
};
use parking_lot::Mutex;
use rstest::{fixture, rstest};
@ -303,7 +299,7 @@ async fn a_deferred_attempt_counts_its_request_and_tokens_but_no_success(
let deployment = router.registry().deployment("a1").unwrap();
let successes = host
.cache
.get("a1:successes")
.local("a1:successes")
.and_then(|count| count.as_f64());
let usage = router.health().usage(&deployment).await.unwrap();
assert_eq!((successes, usage), expected);
@ -375,13 +371,18 @@ async fn a_reported_success_keeps_the_filled_wildcard_model() {
};
router.report_success("w1", 7.0, 0.0, report).await.unwrap();
let minute = current_minute(std::time::SystemTime::now());
let key = |model: &str| usage_key("w1", model, UsageMetric::Tpm, &minute);
let tokens = |model: &str| {
host.cache
.counted(&UsageWindow {
deployment_id: "w1".to_owned(),
model: model.to_owned(),
minute: minute.clone(),
})
.tokens
};
assert_eq!(
(
host.cache.get(&key("openai/gpt-4o")),
host.cache.get(&key("openai/*")),
),
(Some(json!(7)), None),
(tokens("openai/gpt-4o"), tokens("openai/*")),
(Some(7.0), None),
);
}
@ -412,7 +413,10 @@ async fn a_failed_attempt_counts_its_request_only_when_its_callback_ran(
let _ = router.call(Op::from("acompletion"), request()).await;
let deployment = router.registry().deployment("a1").unwrap();
let usage = router.health().usage(&deployment).await.unwrap();
let counted = host.cache.get("a1:fails").and_then(|count| count.as_f64());
let counted = host
.cache
.local("a1:fails")
.and_then(|count| count.as_f64());
assert_eq!((usage.1, counted), (requests, fails));
}
@ -452,7 +456,10 @@ async fn a_reported_failure_counts_its_request_when_asked(
.await
.unwrap();
let deployment = router.registry().deployment("a1").unwrap();
let fails = host.cache.get("a1:fails").and_then(|count| count.as_f64());
let fails = host
.cache
.local("a1:fails")
.and_then(|count| count.as_f64());
let usage = router.health().usage(&deployment).await.unwrap();
assert_eq!((fails, usage.1), (Some(1.0), requests));
}

View file

@ -34,7 +34,7 @@ from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, TypeVar,
from pydantic import TypeAdapter
from litellm._internal_context import service_target
from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, HEALTH_CHECKS_TARGET
from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS
from litellm.integrations.custom_guardrail import is_guardrail_intervention
from litellm.litellm_core_utils.add_retry_fallback_headers import (
add_fallback_headers_to_response,
@ -64,8 +64,8 @@ from litellm.litellm_core_utils.routing_common_utils import (
from litellm.litellm_core_utils.routing_groups import validate_routing_strategy
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.redis.stores.connection import connect
from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, RouterCooldownStore
from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage
from litellm.redis.stores.router_cooldowns import RouterCooldownStore
from litellm.redis.stores.router_usage import RouterUsage
from litellm.router_api import (
ROUTER_INIT_ARGS,
SETTING_FIELDS,
@ -113,7 +113,14 @@ from litellm.router_rust.invoke.streaming import (
shape_attempt_kwargs,
)
from litellm.router_rust.invoke.walk import FallbackOutcomes, WalkMetadata, await_all, log_retry_into
from litellm.router_rust.routing_stores import RoutingStores
from litellm.router_rust.routing_stores import (
CACHE_STEP,
CooldownsStep,
GetLocalStep,
IncrementLocalStep,
RoutingStores,
service_target_of,
)
from litellm.router_rust.steps.access import caller_fields, selection_refusal
from litellm.router_rust.steps.callbacks import OptionalChecks, RustOptionalChecks
from litellm.router_rust.steps.plugins import SYNC_PATH_MESSAGE, candidate_models, run_plugins
@ -156,7 +163,6 @@ from litellm.router_rust.surface.health import (
logged_call,
)
from litellm.router_rust.surface.health_routing import (
HEALTH_STATE_KEY,
HealthRouting,
failover_document,
fully_unhealthy_model_names,
@ -477,13 +483,6 @@ class RustCoreUnsupported(AttributeError):
"""A member or request the Rust core does not serve. It is never forwarded to a Python core."""
def _cache_target(key: str) -> str:
"""The service target a routing cache read or write is attributed to, as the Python core's."""
if key == HEALTH_STATE_KEY:
return HEALTH_CHECKS_TARGET
return ROUTER_COOLDOWNS_TARGET if key.endswith(":cooldown") else ROUTER_USAGE_TARGET
# --- failures -----------------------------------------------------------------------------------
@ -3459,80 +3458,32 @@ class RustCore(
# -- cache steps -----------------------------------------------------------------------------
async def _batch_read(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list
"""One batch read: cooldowns and tpm/rpm usage from their stores, every other key from the routing cache."""
stored: Final = await self._stored_with_usage(keys)
if stored is not None:
return stored
read: Final[object] = await self._cache.async_batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped
return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above
async def _stored_with_usage(self, keys: Sequence[str]) -> tuple[object, ...] | None:
async def _cooldowns_with_usage(self, step: CooldownsStep) -> object:
"""A usage-based-routing-v2 pick reads its tpm/rpm counters alongside its cooldowns (``_usage_read_scope``)."""
usage_read: Final = _USAGE_READ.get()
if usage_read is None or not keys or not keys[0].endswith(":cooldown"):
return await self._stores.batch_get(keys)
if usage_read is None:
return await self._stores.answer(step)
batch, deployments = usage_read
stored, _ = await asyncio.gather(self._stores.batch_get(keys), batch.read_usage(deployments, None))
stored, _ = await asyncio.gather(self._stores.answer(step), batch.read_usage(deployments, None))
return stored
def _batch_read_sync(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list
"""``_batch_read`` for a sync pick."""
stored: Final = self._stores.batch_get_sync(keys)
if stored is not None:
return stored
read: Final[object] = self._cache.batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped
return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above
async def _increment_usage(self, keys: Sequence[object], amounts: Sequence[object]) -> None:
try:
counted: Final = await self._stores.increment_usage(keys, amounts)
except Exception as e: # noqa: BLE001 # usage that cannot be written never fails the call, as on the Python router
from litellm._logging import verbose_router_logger
verbose_router_logger.debug("_increment_deployment_usage: usage was not recorded. Got - %s", e)
return
if not counted:
raise _not_usage(keys)
async def _answer_cache(self, step: Mapping[str, object]) -> object:
with service_target(_cache_target(str(step.get("key")))):
return await self._answer_cache_in_target(step)
async def _answer_cache_in_target(self, step: Mapping[str, object]) -> object:
key: Final = str(step.get("key"))
ttl: Final = _ttl(step)
local: Final = step.get("local") is True
stored: Final = None if local else await self._stores.answer(step.get("op"), key, step.get("value"), ttl)
if stored is not None:
return _plain(stored.value)
match step.get("op"):
case "get":
return _plain(await self._cache.async_get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped
case "batch_get":
return _plain(await self._batch_read([str(k) for k in _as_list(step.get("keys"))]))
case "set":
if ttl is None:
await self._cache.async_set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
else:
await self._cache.async_set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
return None
case "increment":
by: Final = _amount(step.get("value"))
total: Final = (
await self._cache.async_increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType] # see above
if ttl is None
else await self._cache.async_increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
)
return by if total is None else total
case "increment_usage":
await self._increment_usage(_as_list(step.get("keys")), _as_list(step.get("value")))
return None
case "delete":
await self._cache.async_delete_cache(key)
return None
case other:
raise RuntimeError(f"router-core asked for an unknown cache op {other!r}")
async def _answer_cache(self, raw: Mapping[str, object]) -> object:
step: Final = CACHE_STEP.validate_python(raw)
with service_target(service_target_of(step)):
match step:
case CooldownsStep():
return _plain(await self._cooldowns_with_usage(step))
case GetLocalStep(key=key):
return _plain(await self._cache.async_get_cache(key, local_only=True)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped
case IncrementLocalStep(key=key, amount=amount, ttl=ttl):
total: Final = (
await self._cache.async_increment_cache(key, amount, local_only=True) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
if ttl is None
else await self._cache.async_increment_cache(key, amount, local_only=True, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
)
return amount if total is None else total
case _:
return _plain(await self._stores.answer(step))
def run_routing_cache_sync(self, read_or_write: "Callable[[], Coroutine[object, object, object]]") -> object:
"""A routing-cache coroutine on a sync call, whose `make_call` runs on the one event loop
@ -3553,54 +3504,28 @@ class RustCore(
@property
def _stores(self) -> RoutingStores:
"""The stores router-core's cache steps reach, over the cooldowns as last assigned."""
return RoutingStores(self._cooldowns, self._usage)
return RoutingStores(self._cooldowns, self._usage, self._health_checks)
def io_token_usage_sync(self, model_group: str) -> tuple[int | None, int | None]:
"""``get_model_group_io_token_usage``, whose counters stay in the routing cache, on a sync call."""
read: Final = self.run_routing_cache_sync(lambda: self._hosted.get_model_group_io_token_usage(model_group))
return _IO_TOKEN_USAGE.validate_python(read)
def _answer_cache_sync(self, step: Mapping[str, object]) -> object:
with service_target(_cache_target(str(step.get("key")))):
return self._answer_cache_sync_in_target(step)
def _answer_cache_sync_in_target(self, step: Mapping[str, object]) -> object:
key: Final = str(step.get("key"))
ttl: Final = _ttl(step)
local: Final = step.get("local") is True
op: Final = step.get("op")
stored: Final = None if local else self._stores.answer_sync(op, key, step.get("value"), ttl)
if stored is not None:
return _plain(stored.value)
match op:
case "get":
return _plain(self._cache.get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped
case "batch_get":
return _plain(self._batch_read_sync([str(k) for k in _as_list(step.get("keys"))]))
case "set":
if ttl is None:
self._cache.set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
else:
self._cache.set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
return None
case "increment":
by: Final = _amount(step.get("value"))
# DualCache.increment_cache is annotated int but adds floats as given.
return (
self._cache.increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above
if ttl is None
else self._cache.increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above
)
case "increment_usage":
keys: Final = _as_list(step.get("keys"))
if not self._stores.increment_usage_sync(keys, _as_list(step.get("value"))):
raise _not_usage(keys)
return None
case "delete":
self._cache.delete_cache(key) # pyright: ignore[reportUnknownMemberType] # DualCache's delete is untyped
return None
case other:
raise RuntimeError(f"router-core asked for an unknown cache op {other!r}")
def _answer_cache_sync(self, raw: Mapping[str, object]) -> object:
step: Final = CACHE_STEP.validate_python(raw)
with service_target(service_target_of(step)):
match step:
case GetLocalStep(key=key):
return _plain(self._cache.get_cache(key, local_only=True)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped
case IncrementLocalStep(key=key, amount=amount, ttl=ttl):
# DualCache.increment_cache is annotated int but adds floats as given.
return (
self._cache.increment_cache(key, amount, local_only=True) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above
if ttl is None
else self._cache.increment_cache(key, amount, local_only=True, ttl=ttl) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above
)
case _:
return _plain(self._stores.answer_sync(step))
# -- errors the core raises when no attempt failed ------------------------------------------
@ -3960,19 +3885,6 @@ def _seconds(value: object, default: float) -> float | int:
_IO_TOKEN_USAGE: Final = TypeAdapter(tuple[int | None, int | None])
def _not_usage(keys: Sequence[object]) -> RuntimeError:
return RuntimeError(f"router-core asked to add usage to keys that are not one deployment's tpm/rpm: {keys!r}")
def _ttl(step: Mapping[str, object]) -> float | None:
ttl: Final = step.get("ttl")
return ttl if isinstance(ttl, int | float) else None
def _amount(value: object) -> float | int:
return value if isinstance(value, int | float) and not isinstance(value, bool) else 1
# The plans whose kwargs carry `original_generic_function`, which the fallback cooldown reads.
_GENERIC_API_KINDS: Final = frozenset(("generic", "search"))

View file

@ -2,24 +2,36 @@ from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Final, TypeAlias
from typing import Annotated, Final, Literal, Protocol, TypeAlias
from pydantic import TypeAdapter
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
from typing_extensions import assert_never
from litellm._internal_context import service_target, with_service_target
from litellm.caching.redis_batch import Batch
from litellm._logging import verbose_router_logger
from litellm.constants import HEALTH_CHECKS_TARGET
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET
from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter
from litellm.redis.stores.router_usage.types import CounterKind
from litellm.router_utils import routing_prefetch
from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue
__all__ = ("AllowedFailsKey", "CooldownKey", "Handled", "RoutingStores", "UsageKey", "parse_key")
__all__ = (
"CACHE_STEP",
"CacheStep",
"CooldownsStep",
"GetLocalStep",
"IncrementLocalStep",
"LocalStep",
"RoutingStores",
"StoreStep",
"service_target_of",
)
_COOLDOWN_VALUE: Final = TypeAdapter(CooldownCacheValue)
def _cooldown_value(value: object) -> CooldownCacheValue:
def _cooldown_value(value: JsonValue) -> CooldownCacheValue:
"""The cooldown router-core wrote, its whole-second ``cooldown_time`` kept an ``int`` as it serialized it."""
validated: Final = _COOLDOWN_VALUE.validate_python(value)
written: Final = as_str_mapping(value)
@ -29,82 +41,102 @@ def _cooldown_value(value: object) -> CooldownCacheValue:
return validated
@dataclass(frozen=True, slots=True)
class CooldownKey:
model_id: str
class _Step(BaseModel):
model_config = ConfigDict(frozen=True)
@dataclass(frozen=True, slots=True)
class AllowedFailsKey:
model_id: str
class UsageWindow(_Step):
deployment_id: str
model: str
minute: str
def counter(self, kind: CounterKind) -> UsageCounter:
return UsageCounter(self.deployment_id, self.model, self.minute, kind)
class CooldownsStep(_Step):
op: Literal["cooldowns"]
ids: tuple[str, ...]
class PutCooldownStep(_Step):
op: Literal["put_cooldown"]
id: str
cooldown: JsonValue
ttl: float
class IncrementAllowedFailsStep(_Step):
op: Literal["increment_allowed_fails"]
id: str
suffix: str | None
ttl: float
@dataclass(frozen=True, slots=True)
class UsageKey:
counter: UsageCounter
class AddUsageStep(_Step):
op: Literal["add_usage"]
window: UsageWindow
tokens: int | float
requests: int | float
RoutingKey: TypeAlias = CooldownKey | AllowedFailsKey | UsageKey
class AddRequestStep(_Step):
op: Literal["add_request"]
window: UsageWindow
ttl: float
@dataclass(frozen=True, slots=True)
class Handled:
value: object
class UsageStep(_Step):
op: Literal["usage"]
window: UsageWindow
def _deployment_key(rest: str) -> RoutingKey | None:
if rest.endswith(":cooldown"):
return CooldownKey(rest.removesuffix(":cooldown"))
model_id, marker, suffix = rest.partition(":allowed_fails")
if not marker or (suffix and not suffix.startswith(":")):
return None
return AllowedFailsKey(model_id, suffix.removeprefix(":") or None)
class HealthStatesStep(_Step):
op: Literal["health_states"]
def _usage_key(rest: str) -> UsageKey | None:
parts: Final = rest.rsplit(":", 2)
if len(parts) != 3:
return None
head, kind, minute = parts
model_id, separator, model = head.partition(":")
if not separator:
return None
match kind:
case "tpm" | "rpm":
return UsageKey(UsageCounter(model_id, model, minute, kind))
class GetLocalStep(_Step):
op: Literal["get_local"]
key: str
class IncrementLocalStep(_Step):
op: Literal["increment_local"]
key: str
amount: int | float
ttl: float | None
StoreStep: TypeAlias = (
CooldownsStep
| PutCooldownStep
| IncrementAllowedFailsStep
| AddUsageStep
| AddRequestStep
| UsageStep
| HealthStatesStep
)
LocalStep: TypeAlias = GetLocalStep | IncrementLocalStep
CacheStep: TypeAlias = Annotated[StoreStep | LocalStep, Field(discriminator="op")]
CACHE_STEP: Final[TypeAdapter[CacheStep]] = TypeAdapter(CacheStep)
def service_target_of(step: StoreStep | LocalStep) -> str:
"""The service target a cache step is attributed to, as the Python core's."""
match step:
case CooldownsStep() | PutCooldownStep() | IncrementAllowedFailsStep():
return ROUTER_COOLDOWNS_TARGET
case HealthStatesStep():
return HEALTH_CHECKS_TARGET
case AddUsageStep() | AddRequestStep() | UsageStep() | GetLocalStep() | IncrementLocalStep():
return ROUTER_USAGE_TARGET
case _:
return None
assert_never(step)
def parse_key(key: str) -> RoutingKey | None:
"""The store family a router-core key belongs to. Each store rebuilds the same key from what is parsed
here, so an id or model holding a ``:`` still names the key router-core asked for."""
if key.startswith("deployment:"):
return _deployment_key(key.removeprefix("deployment:"))
if key.startswith("global_router:"):
return _usage_key(key.removeprefix("global_router:"))
return None
def _service_target(key: RoutingKey | None) -> str | None:
if isinstance(key, CooldownKey | AllowedFailsKey):
return ROUTER_COOLDOWNS_TARGET
if isinstance(key, UsageKey):
return ROUTER_USAGE_TARGET
return None
def _cooldown_ids(keys: Sequence[str]) -> tuple[str, ...] | None:
ids: Final = tuple(parsed.model_id for parsed in map(parse_key, keys) if isinstance(parsed, CooldownKey))
return ids if keys and len(ids) == len(keys) else None
def _usage_counters(keys: Sequence[object]) -> tuple[UsageCounter, ...] | None:
counters: Final = tuple(
parsed.counter for parsed in (parse_key(str(key)) for key in keys) if isinstance(parsed, UsageKey)
)
return counters if keys and len(counters) == len(keys) else None
class HealthStates(Protocol):
async def async_get_states(self) -> object: ...
def get_states(self) -> object: ...
def _aligned(model_ids: Sequence[str], active: Sequence[tuple[str, CooldownCacheValue]]) -> tuple[object, ...]:
@ -112,119 +144,67 @@ def _aligned(model_ids: Sequence[str], active: Sequence[tuple[str, CooldownCache
return tuple(found.get(model_id) for model_id in model_ids)
def _one_deployment(tpm: UsageCounter, rpm: UsageCounter) -> bool:
return (tpm.kind, rpm.kind) == ("tpm", "rpm") and (tpm.deployment_id, tpm.model, tpm.minute) == (
rpm.deployment_id,
rpm.model,
rpm.minute,
)
def _deployment_increment(keys: Sequence[object], amounts: Sequence[object]) -> tuple[UsageCounter, float, int] | None:
"""``_increment_deployment_usage``'s tpm then rpm increment of one deployment."""
counters: Final = _usage_counters(keys)
if counters is None or len(counters) != 2 or len(amounts) != 2:
return None
tpm, rpm = counters
tokens, requests = amounts
if not _one_deployment(tpm, rpm) or not isinstance(tokens, int | float) or not isinstance(requests, int | float):
return None
return tpm, tokens, int(requests)
def _by(value: object) -> float:
return value if isinstance(value, int | float) and not isinstance(value, bool) else 1
def _usage_counters(window: UsageWindow) -> tuple[UsageCounter, UsageCounter]:
return window.counter("tpm"), window.counter("rpm")
@dataclass(frozen=True, slots=True)
class RoutingStores:
"""Router-core's cache steps for cooldowns, allowed fails and tpm/rpm usage, answered by the native stores
the Python core uses. ``None`` leaves the step to the routing cache."""
"""Router-core's cache steps answered by the stores the Python core uses. The stores own every key."""
def __init__(self, cooldowns: CooldownCache, usage: RouterUsage) -> None:
self._cooldowns: Final = cooldowns
self._usage: Final = usage
cooldowns: CooldownCache
usage: RouterUsage
health: HealthStates
def answer_sync(self, op: object, key: str, value: object, ttl: float | None) -> Handled | None:
store: Final = self._cooldowns.cooldown_store
parsed: Final = parse_key(key)
target: Final = _service_target(parsed)
if target is None:
return None
with service_target(target):
match op, parsed:
case "get", CooldownKey(model_id):
return Handled(_aligned((model_id,), store.active_sync((model_id,)))[0])
case "set", CooldownKey(model_id):
store.put(model_id, _cooldown_value(value), ttl or 0.0)
return Handled(None)
case "delete", CooldownKey(model_id):
store.evict_local(model_id)
return Handled(None)
case "increment", AllowedFailsKey(model_id, suffix):
return Handled(store.increment_allowed_fails(model_id, suffix, ttl or 0.0))
case "increment", UsageKey(counter):
return Handled(self._usage.increment_sync(counter, _by(value), int(ttl or 0)))
case _:
return None
async def answer(self, step: StoreStep) -> object:
store: Final = self.cooldowns.cooldown_store
match step:
case CooldownsStep(ids=ids):
return _aligned(ids, await routing_prefetch.read_cooldowns(store, ids, None))
case PutCooldownStep(id=model_id, cooldown=cooldown, ttl=ttl):
await store.async_put(model_id, _cooldown_value(cooldown), ttl)
return None
case IncrementAllowedFailsStep(id=model_id, suffix=suffix, ttl=ttl):
return await store.async_increment_allowed_fails(model_id, suffix, ttl)
case AddUsageStep(window=window, tokens=tokens, requests=requests):
try:
await self.usage.increment_deployment(
window.deployment_id, window.model, window.minute, tokens, int(requests)
)
except Exception as e: # noqa: BLE001 # usage that cannot be written never fails the call, as on the Python router
verbose_router_logger.debug("_increment_deployment_usage: usage was not recorded. Got - %s", e)
return None
case AddRequestStep(window=window, ttl=ttl):
await self.usage.increment(window.counter("rpm"), 1, int(ttl))
return None
case UsageStep(window=window):
return tuple(await self.usage.get_many(_usage_counters(window)))
case HealthStatesStep():
return await self.health.async_get_states()
case _:
assert_never(step)
async def answer(self, op: object, key: str, value: object, ttl: float | None) -> Handled | None:
parsed: Final = parse_key(key)
target: Final = _service_target(parsed)
if target is None:
return None
store: Final = self._cooldowns.cooldown_store
with service_target(target):
match op, parsed:
case "get", CooldownKey(model_id):
return Handled(_aligned((model_id,), await store.active((model_id,), batch=Batch.NONE))[0])
case "set", CooldownKey(model_id):
await store.async_put(model_id, _cooldown_value(value), ttl or 0.0)
return Handled(None)
case "increment", AllowedFailsKey(model_id, suffix):
return Handled(await store.async_increment_allowed_fails(model_id, suffix, ttl or 0.0))
case "increment", UsageKey(counter):
by: Final = _by(value)
return Handled(await self._usage.increment(counter, by, int(ttl or 0)))
case _:
return self.answer_sync(op, key, value, ttl)
def batch_get_sync(self, keys: Sequence[str]) -> tuple[object, ...] | None:
model_ids: Final = _cooldown_ids(keys)
if model_ids is not None:
with service_target(ROUTER_COOLDOWNS_TARGET):
return _aligned(model_ids, self._cooldowns.cooldown_store.active_sync(model_ids))
counters: Final = _usage_counters(keys)
if counters is None:
return None
with service_target(ROUTER_USAGE_TARGET):
return tuple(self._usage.get_many_sync(counters))
async def batch_get(self, keys: Sequence[str]) -> tuple[object, ...] | None:
model_ids: Final = _cooldown_ids(keys)
if model_ids is not None:
store: Final = self._cooldowns.cooldown_store
with service_target(ROUTER_COOLDOWNS_TARGET):
return _aligned(model_ids, await routing_prefetch.read_cooldowns(store, model_ids, None))
counters: Final = _usage_counters(keys)
if counters is None:
return None
with service_target(ROUTER_USAGE_TARGET):
return tuple(await self._usage.get_many(counters))
@with_service_target(ROUTER_USAGE_TARGET)
def increment_usage_sync(self, keys: Sequence[object], amounts: Sequence[object]) -> bool:
found: Final = _deployment_increment(keys, amounts)
if found is None:
return False
tpm, tokens, requests = found
self._usage.increment_deployment_sync(tpm.deployment_id, tpm.model, tpm.minute, tokens, requests)
return True
@with_service_target(ROUTER_USAGE_TARGET)
async def increment_usage(self, keys: Sequence[object], amounts: Sequence[object]) -> bool:
found: Final = _deployment_increment(keys, amounts)
if found is None:
return False
tpm, tokens, requests = found
await self._usage.increment_deployment(tpm.deployment_id, tpm.model, tpm.minute, tokens, requests)
return True
def answer_sync(self, step: StoreStep) -> object:
store: Final = self.cooldowns.cooldown_store
match step:
case CooldownsStep(ids=ids):
return _aligned(ids, store.active_sync(ids))
case PutCooldownStep(id=model_id, cooldown=cooldown, ttl=ttl):
store.put(model_id, _cooldown_value(cooldown), ttl)
return None
case IncrementAllowedFailsStep(id=model_id, suffix=suffix, ttl=ttl):
return store.increment_allowed_fails(model_id, suffix, ttl)
case AddUsageStep(window=window, tokens=tokens, requests=requests):
self.usage.increment_deployment_sync(
window.deployment_id, window.model, window.minute, tokens, int(requests)
)
return None
case AddRequestStep(window=window, ttl=ttl):
self.usage.increment_sync(window.counter("rpm"), 1, int(ttl))
return None
case UsageStep(window=window):
return tuple(self.usage.get_many_sync(_usage_counters(window)))
case HealthStatesStep():
return self.health.get_states()
case _:
assert_never(step)

View file

@ -21,7 +21,6 @@ if TYPE_CHECKING:
from litellm.router_config import RouterConfig
from litellm.types.router import Deployment, DeploymentTypedDict
HEALTH_STATE_KEY: Final = DeploymentHealthCache.CACHE_KEY
# The failover flags the Python core keeps as plain router attributes, with their defaults.
FAILOVER_FLAGS: Final = ("enable_weighted_failover", "fallback_on_unsupported_params")
@ -149,7 +148,6 @@ async def fully_unhealthy_model_names(
__all__ = [
"FAILOVER_FLAGS",
"HEALTH_STATE_KEY",
"HealthRouting",
"all_blocked",
"failover_document",

View file

@ -105,6 +105,24 @@ class DeploymentHealthCache:
and (now - _timestamp_of(state)) < self.staleness_threshold # pyright: ignore[reportUnknownArgumentType] # a cached record
}
@with_service_target(HEALTH_CHECKS_TARGET)
async def async_get_states(self) -> object:
"""Every deployment's raw state, as the router core's health filter reads it; a failed read is None."""
try:
return await self.cache.async_get_cache(key=self.CACHE_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped
except Exception as e:
verbose_logger.debug("DeploymentHealthCache::async_get_states - Exception: %s", str(e))
return None
@with_service_target(HEALTH_CHECKS_TARGET)
def get_states(self) -> object:
"""``async_get_states`` for a sync call."""
try:
return self.cache.get_cache(key=self.CACHE_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped
except Exception as e:
verbose_logger.debug("DeploymentHealthCache::get_states - Exception: %s", str(e))
return None
@with_service_target(HEALTH_CHECKS_TARGET)
async def async_get_unhealthy_deployment_ids(self, parent_otel_span: Span | None = None) -> set[str]:
"""Return set of deployment IDs currently marked unhealthy and not stale."""

View file

@ -764,20 +764,43 @@ class _PreRouteStep(_Step):
fallback_depth: int
retry: int
# One RouterCache op against the router's DualCache. `key` is set for every op but `batch_get`,
# and `increment_usage`, which set `keys`; `value` is the JSON to store (`set`), the amount
# (`increment`) or one amount per key (`increment_usage`); `ttl` is in seconds; `local` reads or
# writes this worker's in-memory tier only. Outcomes: get -> the JSON value or None; batch_get ->
# one JSON value or None per key, in order; increment -> the new int or float; set, delete and
# increment_usage -> None.
# One deployment's per-minute usage counters: `model` is the model its usage is counted under,
# `minute` is `HH-MM` in UTC.
class _UsageWindow(TypedDict):
deployment_id: str
model: str
minute: str
# One RouterCache op, answered from the router's stores. It names the state it touches, never a
# Redis key. Fields by op: cooldowns (`ids`) -> one stored cooldown or None per id; put_cooldown
# (`id`, `cooldown`, `ttl`) -> None; increment_allowed_fails (`id`, `suffix`, `ttl`) -> the new
# count; add_usage (`window`, `tokens`, `requests`) -> None; add_request (`window`, `ttl`) -> None;
# usage (`window`) -> [tokens, requests], each a number or None; health_states -> the health
# checks' states or None; get_local (`key`) -> this worker's in-memory value or None;
# increment_local (`key`, `amount`, `ttl`) -> the new in-memory count. `ttl` is in seconds.
class _CacheStep(_Step):
kind: Literal["cache"]
op: Literal["get", "set", "increment", "increment_usage", "batch_get", "delete"]
key: str | None
keys: list[str] | None
value: JsonValue
ttl: float | None
local: bool
op: Literal[
"cooldowns",
"put_cooldown",
"increment_allowed_fails",
"add_usage",
"add_request",
"usage",
"health_states",
"get_local",
"increment_local",
]
ids: NotRequired[list[str]]
id: NotRequired[str]
cooldown: NotRequired[JsonValue]
suffix: NotRequired[str | None]
ttl: NotRequired[float | None]
window: NotRequired[_UsageWindow]
tokens: NotRequired[int | float]
requests: NotRequired[int | float]
key: NotRequired[str]
amount: NotRequired[int | float]
# One lookup of a model-info query (router-core `ModelInfoLookup`): its `lookup` name and that
# lookup's arguments, answered as the lookup documents.
@ -1057,8 +1080,6 @@ class RouterCoreHandle:
) -> RouterCall[_SuccessReportSteps]: ...
# Count a success against any id, driven like `report_failure`; `done` carries `key`.
def count_success(self, id: str, callbacks: Mapping[str, bool] | None = None) -> RouterCall[_SuccessReportSteps]: ...
# [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment.
def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ...
# Every deployment as a `model_list` entry, in config order.
def deployments(self) -> list[dict[str, JsonValue]]: ...
def deployments_from(self, start: int) -> list[dict[str, JsonValue]]: ...

View file

@ -8,7 +8,12 @@ import pytest
from litellm.redis.stores.connection import connect
from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore
from litellm.redis.stores.router_usage import RouterUsageStore
from litellm.router_rust.routing_stores import Handled, RoutingStores
from litellm.router_rust.routing_stores import (
CACHE_STEP,
IncrementAllowedFailsStep,
PutCooldownStep,
RoutingStores,
)
from litellm.router_utils.cooldown_cache import CooldownCache
from tests.unit.redis.stores.test_router_cooldowns import FakeCooldowns
@ -49,33 +54,43 @@ class SlowRedis(FakeCooldowns):
return self.increment_allowed_fails(model_id, suffix, ttl_seconds)
class NoHealthStates:
async def async_get_states(self) -> object:
return None
def get_states(self) -> object:
return None
async def _another_request(ran: asyncio.Event) -> None:
ran.set()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("op", "key", "value", "answer"),
("raw", "answer"),
[
("set", "deployment:dep:cooldown", _ENTRY, None),
("increment", "deployment:dep:allowed_fails:timeout", 1, 1),
({"op": "put_cooldown", "id": "dep", "cooldown": _ENTRY, "ttl": 60.0}, None),
({"op": "increment_allowed_fails", "id": "dep", "suffix": "timeout", "ttl": 60.0}, 1),
],
)
async def test_an_async_cooldown_step_waiting_on_redis_lets_other_requests_run(
op: str, key: str, value: object, answer: object
raw: Mapping[str, object], answer: object
) -> None:
"""router-core's failure path writes cooldowns and allowed-fails counts; on an async call that Redis round trip
must not hold the event loop, or every request on the worker stalls behind it."""
other_request: Final = asyncio.Event()
handle: Final = SlowRedis(has_redis=True, other_request=other_request)
cooldowns: Final = CooldownCache(store=RouterCooldownStore(handle), default_cooldown_time=60)
stores: Final = RoutingStores(cooldowns, RouterUsageStore(connect(None).router_usage()))
stores: Final = RoutingStores(cooldowns, RouterUsageStore(connect(None).router_usage()), NoHealthStates())
step: Final = CACHE_STEP.validate_python(raw)
assert isinstance(step, PutCooldownStep | IncrementAllowedFailsStep)
step: Final = asyncio.create_task(stores.answer(op, key, value, 60.0))
answering: Final = asyncio.create_task(stores.answer(step))
other: Final = asyncio.create_task(_another_request(other_request))
handled: Final = await step
handled: Final = await answering
await other
assert handled == Handled(answer)
assert handled == answer
assert handle.overlapped is True
assert handle.targets == ["router_cooldowns"]