mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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:
parent
08f2a98250
commit
6d019c9dc8
14 changed files with 938 additions and 887 deletions
|
|
@ -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 { .. })]
|
||||
));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
"#
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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!(
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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]]: ...
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue