From 1b8f4d2fe471e6bcde13338f9c4698eac570add3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 5 Nov 2026 09:19:28 +0900 Subject: [PATCH] refactor(redis-store): own the router cooldown and usage encodings Cooldown converts to and from its CooldownCacheValue dict and router_usage::Kind derives its tpm and rpm key segment with strum, so the bridge carries those values through PyValue and FromStr instead of spelling the fields itself. PyValue moves down from the auth cache store to serve both. --- litellm-rust/Cargo.lock | 2 + litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../python-bridge/src/redis/store/mod.rs | 1 + .../src/redis/store/router_cooldowns.rs | 27 ++-- .../src/redis/store/router_usage.rs | 16 +-- .../python-bridge/src/redis/store/value.rs | 125 ++++++++++++++++++ litellm-rust/crates/redis-store/Cargo.toml | 1 + .../redis-store/src/router_cooldowns.rs | 23 +++- .../crates/redis-store/src/router_usage.rs | 10 +- .../crates/redis-store/tests/router_usage.rs | 15 +++ 10 files changed, 181 insertions(+), 40 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/value.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index d4cc4db147d..3b67588b506 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4377,6 +4377,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-python-compat", "litellm-redis-client", "litellm-redis-connect", "litellm-redis-store", @@ -4493,6 +4494,7 @@ dependencies = [ "litellm-secret-value", "rstest", "serde_json", + "strum", "thiserror 2.0.19", "tokio", "tracing", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index bd62f16fed9..b574e9c23c6 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -34,6 +34,7 @@ litellm-redis-client.workspace = true litellm-redis-connect.workspace = true litellm-redis-store.workspace = true litellm-redis-resolve.workspace = true +litellm-python-compat.workspace = true litellm-redis-types.workspace = true tracing.workspace = true litellm-cache-response.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs index e609316efe9..f4cd89b342d 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -2,6 +2,7 @@ mod policy; pub(crate) mod request; pub(crate) mod router_cooldowns; pub(crate) mod router_usage; +mod value; use std::{ pin::Pin, diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs index 69f005dcb3e..044c64dcf98 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs @@ -2,13 +2,14 @@ use std::sync::Arc; use litellm_redis_store::router_cooldowns::{Cooldown, RouterCooldowns}; use pyo3::{ + exceptions::PyValueError, prelude::*, types::{PyDict, PyTuple}, }; use super::{ NativePending, PyNumber, policy::router_tier_policies, redis_async, redis_sync, - request::NativeRequestRedis, store_error, ttl, + request::NativeRequestRedis, store_error, ttl, value::PyValue, }; use crate::redis::client::NativeRedisClient; @@ -17,35 +18,27 @@ use crate::redis::client::NativeRedisClient; #[pyclass(frozen, name = "NativeRouterCooldowns")] pub(crate) struct NativeRouterCooldowns(Arc); +/// A `CooldownCacheValue` dict, whose shape `Cooldown` owns. pub(crate) struct PyCooldown(Cooldown); impl<'a, 'py> FromPyObject<'a, 'py> for PyCooldown { type Error = PyErr; fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { - let fields = value.cast::()?; - let item = |name: &str| fields.as_any().get_item(name); - Ok(Self(Cooldown { - exception_received: item("exception_received")?.extract()?, - status_code: item("status_code")?.extract()?, - timestamp: item("timestamp")?.extract::()?.0, - cooldown_time: item("cooldown_time")?.extract::()?.0, - })) + let PyValue(value) = value.extract()?; + Cooldown::from_value(&value) + .map(Self) + .map_err(|error| PyValueError::new_err(error.to_string())) } } impl<'py> IntoPyObject<'py> for PyCooldown { - type Target = PyDict; - type Output = Bound<'py, PyDict>; + type Target = PyAny; + type Output = Bound<'py, PyAny>; type Error = PyErr; fn into_pyobject(self, py: Python<'py>) -> PyResult { - let fields = PyDict::new(py); - fields.set_item("exception_received", self.0.exception_received)?; - fields.set_item("status_code", self.0.status_code)?; - fields.set_item("timestamp", PyNumber(self.0.timestamp))?; - fields.set_item("cooldown_time", PyNumber(self.0.cooldown_time))?; - Ok(fields) + PyValue(self.0.to_value()).into_pyobject(py) } } diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs index fe2945e7497..26597a16731 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs @@ -1,6 +1,6 @@ use std::{sync::Arc, time::Duration}; -use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage}; +use litellm_redis_store::router_usage::{Counter, RouterUsage}; use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; use super::{PyNumber, policy::router_tier_policies, redis_async, store_error}; @@ -10,7 +10,7 @@ use crate::redis::client::NativeRedisClient; #[pyclass(frozen, name = "NativeRouterUsage")] pub(crate) struct NativeRouterUsage(Arc); -/// `(deployment_id, model, minute, "tpm" | "rpm")`. +/// `(deployment_id, model, minute, kind)`, the kind spelled as its key segment. pub(crate) struct PyCounter(Counter); impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter { @@ -19,15 +19,9 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter { fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { let (deployment_id, model, minute, kind): (String, Option, String, String) = value.extract()?; - let kind = match kind.as_str() { - "tpm" => Kind::Tpm, - "rpm" => Kind::Rpm, - _ => { - return Err(PyValueError::new_err(format!( - "unknown usage counter kind {kind:?}" - ))); - } - }; + let kind = kind + .parse() + .map_err(|_| PyValueError::new_err(format!("unknown usage counter kind {kind:?}")))?; Ok(Self(Counter { deployment_id, model, diff --git a/litellm-rust/crates/python-bridge/src/redis/store/value.rs b/litellm-rust/crates/python-bridge/src/redis/store/value.rs new file mode 100644 index 00000000000..c86c9e9bb20 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/value.rs @@ -0,0 +1,125 @@ +use litellm_python_compat::{BigInt, MAX_DEPTH, Value}; +use pyo3::{ + exceptions::{PyTypeError, PyValueError}, + prelude::*, + types::{ + PyBool, PyBytes, PyComplex, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PySet, PyString, + PyTuple, + }, +}; + +/// A plain Python value (what `json.loads` or `ast.literal_eval` returns) across the bridge. +/// Anything else, a live object included, is refused with `TypeError`. +pub(crate) struct PyValue(pub(crate) Value); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyValue { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + to_value(&value, 0).map(Self) + } +} + +impl<'py> IntoPyObject<'py> for PyValue { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + to_python(py, self.0) + } +} + +fn to_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult { + if depth > MAX_DEPTH { + return Err(PyValueError::new_err("value is nested too deeply")); + } + let nested = |items: Bound<'_, PyAny>| -> PyResult> { + items + .try_iter()? + .map(|item| to_value(&item?, depth + 1)) + .collect() + }; + if value.is_none() { + return Ok(Value::None); + } + if let Ok(flag) = value.cast::() { + return Ok(Value::Bool(flag.is_true())); + } + if value.is_instance_of::() { + return Ok(Value::Int(match value.extract::() { + Ok(small) => small.into(), + Err(_) => value + .str()? + .to_str()? + .parse::() + .map_err(|error| PyValueError::new_err(error.to_string()))?, + })); + } + if let Ok(number) = value.cast::() { + return Ok(Value::Float(number.value())); + } + if let Ok(text) = value.cast::() { + return Ok(Value::Str(text.to_str()?.to_owned())); + } + if let Ok(bytes) = value.cast::() { + return Ok(Value::Bytes(bytes.as_bytes().to_vec())); + } + if let Ok(complex) = value.cast::() { + return Ok(Value::Complex { + re: complex.real(), + im: complex.imag(), + }); + } + if let Ok(entries) = value.cast::() { + return entries + .iter() + .map(|(key, item)| Ok((to_value(&key, depth + 1)?, to_value(&item, depth + 1)?))) + .collect::>() + .map(Value::Dict); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::List); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::Tuple); + } + if value.is_instance_of::() || value.is_instance_of::() { + return nested(value.clone()).map(Value::Set); + } + Err(PyTypeError::new_err(format!( + "{} is not a plain data value", + value.get_type().name()? + ))) +} + +fn to_python(py: Python<'_>, value: Value) -> PyResult> { + let all = |items: Vec| -> PyResult>> { + items.into_iter().map(|item| to_python(py, item)).collect() + }; + Ok(match value { + Value::None => py.None().into_bound(py), + Value::Bool(flag) => PyBool::new(py, flag).to_owned().into_any(), + Value::Int(number) => match i64::try_from(&number) { + Ok(small) => small.into_pyobject(py)?.into_any(), + Err(_) => py + .get_type::() + .call1((number.to_string(),))? + .into_any(), + }, + Value::Float(number) => number.into_pyobject(py)?.into_any(), + Value::Complex { re, im } => PyComplex::from_doubles(py, re, im).into_any(), + Value::Str(text) => text.into_pyobject(py)?.into_any(), + Value::Bytes(bytes) => PyBytes::new(py, &bytes).into_any(), + Value::Tuple(items) => PyTuple::new(py, all(items)?)?.into_any(), + Value::List(items) => PyList::new(py, all(items)?)?.into_any(), + Value::Set(items) => PySet::new(py, all(items)?)?.into_any(), + Value::Dict(entries) => { + let dict = PyDict::new(py); + for (key, item) in entries { + dict.set_item(to_python(py, key)?, to_python(py, item)?)?; + } + dict.into_any() + } + }) +} diff --git a/litellm-rust/crates/redis-store/Cargo.toml b/litellm-rust/crates/redis-store/Cargo.toml index b5d3c0cf88c..7b4466ffa4a 100644 --- a/litellm-rust/crates/redis-store/Cargo.toml +++ b/litellm-rust/crates/redis-store/Cargo.toml @@ -12,6 +12,7 @@ litellm-cache-memory.workspace = true litellm-python-compat.workspace = true litellm-redis-client.workspace = true serde_json.workspace = true +strum.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } tracing.workspace = true diff --git a/litellm-rust/crates/redis-store/src/router_cooldowns.rs b/litellm-rust/crates/redis-store/src/router_cooldowns.rs index 52643eabd0b..fe194ce2c33 100644 --- a/litellm-rust/crates/redis-store/src/router_cooldowns.rs +++ b/litellm-rust/crates/redis-store/src/router_cooldowns.rs @@ -36,10 +36,10 @@ impl Cooldown { self.timestamp.as_f64() + self.cooldown_time.as_f64() - now.as_secs_f64() } - /// `str(dict)`, the bytes `RedisCache.set_cache` writes on `main`. - pub fn encode(&self) -> String { + /// The `CooldownCacheValue` dict the router reads and writes. + pub fn to_value(&self) -> Value { let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value); - repr(&Value::Dict(vec![ + Value::Dict(vec![ field( "exception_received", Value::Str(self.exception_received.clone()), @@ -47,12 +47,11 @@ impl Cooldown { field("status_code", Value::Str(self.status_code.clone())), field("timestamp", self.timestamp.to_value()), field("cooldown_time", self.cooldown_time.to_value()), - ])) + ]) } - /// Reads both `main`'s `str(dict)` and a JSON object, as `RedisCache._get_cache_logic` does. - pub fn decode(raw: &[u8]) -> Result { - let Value::Dict(fields) = python_value(raw)? else { + pub fn from_value(value: &Value) -> Result { + let Value::Dict(fields) = value else { return Err(DecodeError::Shape("a cooldown is a dict")); }; let field = |name: &str| { @@ -77,6 +76,16 @@ impl Cooldown { cooldown_time: number("cooldown_time")?, }) } + + /// `str(dict)`, the bytes `RedisCache.set_cache` writes on `main`. + pub fn encode(&self) -> String { + repr(&self.to_value()) + } + + /// Reads both `main`'s `str(dict)` and a JSON object, as `RedisCache._get_cache_logic` does. + pub fn decode(raw: &[u8]) -> Result { + Self::from_value(&python_value(raw)?) + } } fn key(model_id: &str) -> String { diff --git a/litellm-rust/crates/redis-store/src/router_usage.rs b/litellm-rust/crates/redis-store/src/router_usage.rs index 68353aa5050..d3ce0fdfea9 100644 --- a/litellm-rust/crates/redis-store/src/router_usage.rs +++ b/litellm-rust/crates/redis-store/src/router_usage.rs @@ -22,9 +22,12 @@ const MEMORY_ENTRIES: usize = 200; /// `RoutingArgs.ttl`, the window of a per-minute counter. pub const COUNTER_TTL: Duration = Duration::from_secs(60); -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +/// A counter's key segment, which is also how the Python router names it. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)] pub enum Kind { + #[strum(serialize = "tpm")] Tpm, + #[strum(serialize = "rpm")] Rpm, } @@ -39,10 +42,7 @@ pub struct Counter { impl Counter { /// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`. fn key(&self) -> String { - let kind = match self.kind { - Kind::Tpm => "tpm", - Kind::Rpm => "rpm", - }; + let kind: &str = self.kind.into(); let model = self.model.as_deref().unwrap_or("None"); format!( "global_router:{}:{model}:{kind}:{}", diff --git a/litellm-rust/crates/redis-store/tests/router_usage.rs b/litellm-rust/crates/redis-store/tests/router_usage.rs index 75cbfdf2ebb..ed27351c807 100644 --- a/litellm-rust/crates/redis-store/tests/router_usage.rs +++ b/litellm-rust/crates/redis-store/tests/router_usage.rs @@ -20,6 +20,21 @@ fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] { }) } +#[rstest] +#[case::tpm(Kind::Tpm, "tpm")] +#[case::rpm(Kind::Rpm, "rpm")] +fn a_counter_kind_round_trips_through_its_key_segment(#[case] kind: Kind, #[case] segment: &str) { + assert_eq!(<&str>::from(kind), segment); + assert_eq!(segment.parse::().unwrap(), kind); +} + +#[rstest] +#[case::unknown("tokens")] +#[case::wrong_case("TPM")] +fn other_text_is_not_a_counter_kind(#[case] segment: &str) { + assert!(segment.parse::().is_err()); +} + #[rstest] #[tokio::test] async fn reads_counters_python_incremented_under_mains_keys(