diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs index 6ccac4df659..a8f4eb9245e 100644 --- a/litellm-rust/crates/python-bridge/src/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -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 }, + /// Answered with `None`. + PutCooldown { + id: String, + cooldown: Value, + ttl: f64, + }, + /// Answered with the count after this failure. + IncrementAllowedFails { + id: String, + suffix: Option, + 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, + }, +} + +/// 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 { + 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), - /// `batch_get`: one entry per key. Values(Vec>), - /// `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, - keys: Option>, - value: Option, - ttl: Option, - /// 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) -> Option { } /// 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, - ttl: Option, -) -> CacheOutcome { - call_cache(op, key, keys, value, ttl, false).await -} - -async fn call_cache( - op: CacheOp, - key: Option<&str>, - keys: Option<&[String]>, - value: Option, - ttl: Option, - 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 { + let CacheOutcome::Number(next) = cache_step(call).await else { + unreachable!("{CHECKED}"); + }; + Ok(next) +} + +async fn value(call: CacheCall) -> Result, litellm_cache::Error> { + let CacheOutcome::Value(value) = cache_step(call).await else { + unreachable!("{CHECKED}"); + }; + Ok(value) +} + +async fn values(call: CacheCall) -> Result>, 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, 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>, 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, + 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 { + 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 { + 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, litellm_cache::Error> { + value(CacheCall::HealthStates).await + } + + async fn get_local(&self, key: &str) -> Result, litellm_cache::Error> { + value(CacheCall::GetLocal { + key: key.to_owned(), + }) + .await + } + + async fn increment_local( &self, key: &str, amount: f64, ttl: Option, ) -> Result { - 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>, 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, - ) -> Result<(), litellm_cache::Error> { - let keys: Vec = 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, 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, - ) -> Result { - 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, + cache: InMemoryRouterCache, calls: Vec, } 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 { .. })] )); } diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs index 8157d0f5a32..674b85013a2 100644 --- a/litellm-rust/crates/python-bridge/src/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -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 { - 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> { - 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> { 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>) -> PyResult>) -> PyResult { let outcome = outcome.filter(|outcome| !outcome.is_none()); let json = |value: Option| value.filter(|value| !value.is_null()); match op { - CacheOp::Get => outcome + CacheOp::HealthStates | CacheOp::GetLocal => outcome .map(|outcome| from_py_argument::(&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::() { - 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> = 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::()) .and_then(|outcome| outcome.extract::().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>) -> PyResult 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> { + 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" "# ), diff --git a/litellm-rust/crates/redis-store/src/router_usage/keys.rs b/litellm-rust/crates/redis-store/src/router_usage/keys.rs index 261fbeb1d1d..1482668ab2b 100644 --- a/litellm-rust/crates/redis-store/src/router_usage/keys.rs +++ b/litellm-rust/crates/redis-store/src/router_usage/keys.rs @@ -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!( diff --git a/litellm-rust/crates/router-core/src/cache.rs b/litellm-rust/crates/router-core/src/cache.rs index 9811e34a3e3..f53ff875feb 100644 --- a/litellm-rust/crates/router-core/src/cache.rs +++ b/litellm-rust/crates/router-core/src/cache.rs @@ -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, + pub requests: Option, +} + #[async_trait] pub trait RouterCache: Send + Sync { - async fn async_get_cache(&self, key: &str) -> Result, Error>; + /// The stored cooldown of each of `ids`, in order, `None` for one without. + async fn cooldowns(&self, ids: &[String]) -> Result>, 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, + id: &str, + suffix: Option<&str>, + ttl: Duration, + ) -> Result; + + /// `_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; + + /// The background health checks' per-deployment states, as `DeploymentHealthCache` keeps them. + async fn health_states(&self) -> Result, Error>; + + /// Read this worker's in-memory tier only (`local_only=True`). + async fn get_local(&self, key: &str) -> Result, 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, ) -> Result; - - async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, 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, 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, - ) -> Result { - 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, - ) -> 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), + Tokens(UsageWindow), + Requests(UsageWindow), + HealthStates, + Local(String), +} + struct Entry { value: Value, expires: Option, } -/// 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>>, + entries: Arc>>, } impl InMemoryRouterCache { - fn live(entries: &mut HashMap, key: &str) -> Option { + fn live(entries: &mut HashMap, slot: &Slot) -> Option { 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 { - Self::live(&mut self.entries.lock(), key) + fn get(&self, slot: &Slot) -> Option { + Self::live(&mut self.entries.lock(), slot) } - pub fn increment(&self, key: &str, amount: f64, ttl: Option) -> f64 { + fn set(&self, slot: Slot, value: Value, ttl: Option) { + 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) -> 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 { + 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, Error> { - Ok(self.get(key)) + async fn cooldowns(&self, ids: &[String]) -> Result>, 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, - ) -> 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 { + 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 { + Ok(self.counted(window)) + } + + async fn health_states(&self) -> Result, Error> { + Ok(self.get(&Slot::HealthStates)) + } + + async fn get_local(&self, key: &str) -> Result, Error> { + Ok(self.local(key)) + } + + async fn increment_local( &self, key: &str, amount: f64, ttl: Option, ) -> Result { - Ok(self.increment(key, amount, ttl)) - } - - async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, 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)) } } diff --git a/litellm-rust/crates/router-core/src/health.rs b/litellm-rust/crates/router-core/src/health.rs index eaa61dac7ba..3b6ce607ff1 100644 --- a/litellm-rust/crates/router-core/src/health.rs +++ b/litellm-rust/crates/router-core/src/health.rs @@ -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, Error> { - let keys: Vec = 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 { 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 { 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 { - 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, Error> { + ) -> Result, 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| 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), + }, ); } } diff --git a/litellm-rust/crates/router-core/src/health_routing.rs b/litellm-rust/crates/router-core/src/health_routing.rs index bc3e43a6e92..9fe93f9c70e 100644 --- a/litellm-rust/crates/router-core/src/health_routing.rs +++ b/litellm-rust/crates/router-core/src/health_routing.rs @@ -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 { 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) diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 80c8f597157..69d4ee90f95 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -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, Error> { + ) -> Result, Error> { report .callbacks .record( @@ -459,7 +459,7 @@ impl Router { total_tokens: f64, requests: f64, report: &SuccessReport, - ) -> Result, Error> { + ) -> Result, Error> { let snapshot = self.core.snapshot.load_full(); let Some(deployment) = snapshot.registry.deployment(id) else { return Ok(None); diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs index a2a741d6135..facd7e1a633 100644 --- a/litellm-rust/crates/router-core/tests/call_loop.rs +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -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)); } diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 40ae7e0e554..471f5c8ad91 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -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")) diff --git a/litellm/router_rust/routing_stores.py b/litellm/router_rust/routing_stores.py index af3ef7cd375..1dd678d377d 100644 --- a/litellm/router_rust/routing_stores.py +++ b/litellm/router_rust/routing_stores.py @@ -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) diff --git a/litellm/router_rust/surface/health_routing.py b/litellm/router_rust/surface/health_routing.py index 46542bf9af2..44a045a0d5d 100644 --- a/litellm/router_rust/surface/health_routing.py +++ b/litellm/router_rust/surface/health_routing.py @@ -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", diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index 1c9cb10a48b..ea9833ccda5 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -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.""" diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 773985ad427..d0deb8000af 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -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]]: ... diff --git a/tests/unit/router_rust/test_routing_stores.py b/tests/unit/router_rust/test_routing_stores.py index 6fc19cdcfa5..f27c3499ff3 100644 --- a/tests/unit/router_rust/test_routing_stores.py +++ b/tests/unit/router_rust/test_routing_stores.py @@ -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"]