From ddaf8c6c2b78418b2ad374f811225a8188e04576 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 5 Nov 2026 06:35:09 +0900 Subject: [PATCH] feat(redis): bridge Python to a shared native Redis client and request batches NativeRedisClient resolves a RedisCache's redis_kwargs through the RedisCache kwargs surface and hands every native store in the process one shared client, breaker and ServiceLogging event sink. A setting Rust ignores is logged with its consequence, and a config it cannot use leaves the stores on memory NativeRequestRedis carries one request's pre-call and post-call batches from RequestRedisBatches into the stores. The keyspaces that use these land in the PRs above RedisSpec carries RedisCache.redis_kwargs into the resolver as plain values plus the names of keys whose value is a Python object, instead of replacing each object with its type name. ConfiguredRedis::with_host_objects takes those names, so the resolver reports them in Resolution::host_objects, warns for the ones Rust cannot honor and lets each hide its REDIS_* value. A credential_provider object without a ProviderHandle fails with ProviderWithoutHandle, so the stores stay on memory instead of connecting without auth Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test-unit.yml | 1 + litellm-rust/Cargo.lock | 25 ++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/host-python/src/lib.rs | 2 +- .../crates/host-python/src/runtime.rs | 5 + litellm-rust/crates/python-bridge/Cargo.toml | 7 + litellm-rust/crates/python-bridge/src/lib.rs | 10 + .../python-bridge/src/redis/client/mod.rs | 199 +++++++++++++ .../crates/python-bridge/src/redis/events.rs | 274 ++++++++++++++++++ .../crates/python-bridge/src/redis/mod.rs | 3 + .../python-bridge/src/redis/store/mod.rs | 16 + .../python-bridge/src/redis/store/request.rs | 41 +++ litellm-rust/crates/redis-connect/Cargo.toml | 1 + .../crates/redis-connect/src/connect.rs | 3 + litellm-rust/crates/redis-store/AGENTS.md | 7 + litellm-rust/crates/redis-store/Cargo.toml | 24 ++ litellm-rust/crates/redis-store/src/lib.rs | 3 + .../crates/redis-store/src/request.rs | 121 ++++++++ .../crates/redis-store/tests/request.rs | 181 ++++++++++++ .../crates/redis-store/tests/support/mod.rs | 83 ++++++ litellm/caching/redis_batch.py | 95 +++++- litellm/caching/redis_cache.py | 10 +- litellm/redis/__init__.py | 0 litellm/redis/stores/__init__.py | 0 litellm/redis/stores/_connection.py | 40 +++ litellm/rust_bridge/_native.pyi | 21 ++ litellm/rust_bridge/host/AGENTS.md | 2 +- litellm/rust_bridge/host/redis_events.py | 46 +++ tests/unit/caching/test_redis_batch.py | 42 ++- tests/unit/redis/__init__.py | 0 tests/unit/redis/stores/__init__.py | 0 tests/unit/redis/stores/test_connection.py | 28 ++ .../rust_bridge/host/test_redis_events.py | 81 ++++++ 33 files changed, 1360 insertions(+), 12 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/client/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/events.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/request.rs create mode 100644 litellm-rust/crates/redis-store/AGENTS.md create mode 100644 litellm-rust/crates/redis-store/Cargo.toml create mode 100644 litellm-rust/crates/redis-store/src/lib.rs create mode 100644 litellm-rust/crates/redis-store/src/request.rs create mode 100644 litellm-rust/crates/redis-store/tests/request.rs create mode 100644 litellm-rust/crates/redis-store/tests/support/mod.rs create mode 100644 litellm/redis/__init__.py create mode 100644 litellm/redis/stores/__init__.py create mode 100644 litellm/redis/stores/_connection.py create mode 100644 litellm/rust_bridge/host/redis_events.py create mode 100644 tests/unit/redis/__init__.py create mode 100644 tests/unit/redis/stores/__init__.py create mode 100644 tests/unit/redis/stores/test_connection.py create mode 100644 tests/unit/rust_bridge/host/test_redis_events.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index a090ec7bdb1..40618b2bacc 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -170,6 +170,7 @@ jobs: tests/unit/google_genai tests/unit/router_strategy tests/unit/router_utils + tests/unit/redis tests/unit/proxy/common_utils/test_cache_aware_routing.py tests/unit/enterprise/enterprise_callbacks/send_emails tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 0b49deedd0f..d4cc4db147d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4361,6 +4361,8 @@ dependencies = [ "litellm-cache", "litellm-cache-response", "litellm-callbacks-legacy-python", + "litellm-circuit-breaker", + "litellm-config", "litellm-core-utils", "litellm-error", "litellm-host", @@ -4375,6 +4377,10 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-redis-client", + "litellm-redis-connect", + "litellm-redis-store", + "litellm-redis-types", "litellm-router-types", "litellm-secret-value", "litellm-secrets", @@ -4473,6 +4479,25 @@ dependencies = [ "url", ] +[[package]] +name = "litellm-redis-store" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-cache", + "litellm-cache-memory", + "litellm-circuit-breaker", + "litellm-python-compat", + "litellm-redis-client", + "litellm-redis-testing", + "litellm-secret-value", + "rstest", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tracing", +] + [[package]] name = "litellm-redis-testing" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a24cfd02681..ed9881fb281 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -72,6 +72,7 @@ litellm-python-compat = { path = "crates/python-compat" } litellm-circuit-breaker = { path = "crates/circuit-breaker" } litellm-redis-client = { path = "crates/redis-client" } litellm-redis-connect = { path = "crates/redis-connect" } +litellm-redis-store = { path = "crates/redis-store" } litellm-redis-testing = { path = "crates/redis-testing" } litellm-redis-resolve = { path = "crates/redis-resolve" } litellm-redis-types = { path = "crates/redis-types" } diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 8e37ab02669..3b22f74d95e 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -37,7 +37,7 @@ pub use owned::PythonOwned; pub use runtime::{ ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value, ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync, - run_sync_value, runtime_started, + run_sync_value, runtime_handle, runtime_started, }; pub use services::PythonHostCalls; diff --git a/litellm-rust/crates/host-python/src/runtime.rs b/litellm-rust/crates/host-python/src/runtime.rs index e8872c9340a..4f91e2e085b 100644 --- a/litellm-rust/crates/host-python/src/runtime.rs +++ b/litellm-rust/crates/host-python/src/runtime.rs @@ -68,6 +68,11 @@ fn runtime() -> PyResult<&'static Runtime> { Ok(pyo3_async_runtimes::tokio::get_runtime()) } +/// The native runtime, for work that must outlive the call that schedules it. +pub fn runtime_handle() -> PyResult { + Ok(runtime()?.handle().clone()) +} + #[expect(clippy::disallowed_methods, reason = "this is the gated door")] fn future_into_py(py: Python<'_>, future: F) -> PyResult> where diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index fd826f549e5..bd62f16fed9 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -29,6 +29,13 @@ bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true litellm-error.workspace = true +litellm-circuit-breaker.workspace = true +litellm-redis-client.workspace = true +litellm-redis-connect.workspace = true +litellm-redis-store.workspace = true +litellm-redis-resolve.workspace = true +litellm-redis-types.workspace = true +tracing.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index de35419308a..133ac0991b6 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -12,6 +12,7 @@ mod marshal; mod model_registry; mod preflight; mod python_settings; +mod redis; mod routes; mod secrets; mod tokenizer; @@ -36,6 +37,12 @@ mod _native { #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] + use crate::redis::client::NativeRedisClient; + #[pymodule_export] + use crate::redis::client::native_redis_breaker_states; + #[pymodule_export] + use crate::redis::store::request::NativeRequestRedis; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{acompletion, completion}; @@ -108,6 +115,9 @@ mod tests { "aresponses", "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", + "NativeRedisClient", + "NativeRequestRedis", + "native_redis_breaker_states", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", diff --git a/litellm-rust/crates/python-bridge/src/redis/client/mod.rs b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs new file mode 100644 index 00000000000..f709a688bcb --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs @@ -0,0 +1,199 @@ +use std::{ + collections::BTreeMap, + sync::{Arc, LazyLock, Mutex, PoisonError}, + time::Duration, +}; + +use litellm_circuit_breaker::Config as BreakerConfig; +use litellm_core_utils::settings::ProcessEnvironment; +use litellm_redis_client::{RedisClient, RedisEvents}; +use litellm_redis_connect::{HostServices, SharedEnvironment}; +use litellm_redis_resolve::Kwargs; +use litellm_redis_types::Resolution; +use pyo3::{ + prelude::*, + sync::MutexExt, + types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple}, +}; + +use crate::redis::events::{ServiceLoggingEvents, state_label}; + +/// The kwargs `_merge_redis_kwargs` returns for a `RedisCache`, which the resolver reads the way +/// redis-py's constructors do. A key whose value is a Python object, such as `retry`, reaches the +/// resolver by name only, so it can name the setting it cannot honor. +pub(crate) struct RedisSpec(Kwargs); + +impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec { + type Error = PyErr; + + fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult { + let py = kwargs.py(); + let plain = PyDict::new(py); + let mut host_objects = Vec::new(); + for (name, value) in kwargs.cast::()?.iter() { + match value.is_none() || is_plain(&value) { + true => plain.set_item(&name, value)?, + false => host_objects.push(name.extract::()?), + } + } + let fields: BTreeMap = + litellm_host_python::from_py_argument(plain.as_any())?; + let objects = host_objects + .into_iter() + .map(|name| (name, serde_json::Value::Bool(true))); + Ok(Self(litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)))) + } +} + +/// `litellm.constants`' `REDIS_CIRCUIT_BREAKER_*`, which Python reads from the environment once. +pub(crate) struct RedisBreaker(BreakerConfig); + +impl<'a, 'py> FromPyObject<'a, 'py> for RedisBreaker { + type Error = PyErr; + + fn extract(breaker: Borrowed<'a, 'py, PyAny>) -> PyResult { + let (failure_threshold, recovery_timeout, timeout_min_duration): (u32, f64, f64) = + breaker.extract()?; + let seconds = |seconds: f64| { + Duration::try_from_secs_f64(seconds) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string())) + }; + Ok(Self(BreakerConfig { + failure_threshold, + recovery_timeout: seconds(recovery_timeout)?, + timeout_min_duration: seconds(timeout_min_duration)?, + })) + } +} + +fn is_plain(value: &Bound<'_, PyAny>) -> bool { + value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() +} + +struct Shared { + resolved: Resolution, + client: RedisClient, +} + +static CLIENTS: LazyLock>> = LazyLock::new(|| Mutex::new(Vec::new())); + +/// The one client for a resolved Redis config, so every family in this process shares its pool +/// and its breaker. Like redis-py, which never connects at construction, a config the resolver or +/// the connector cannot honor leaves the family on memory with a warning rather than failing the +/// `Router`. +fn shared_client( + py: Python<'_>, + spec: RedisSpec, + namespace: Option, + breaker: Option, +) -> PyResult> { + let environment: SharedEnvironment = Arc::new(ProcessEnvironment); + let logger = crate::logger::capture(py); + let resolved = match logger.scope(|| litellm_redis_resolve::interpret(&spec.0, None)) { + Ok(resolution) => { + logger.scope(|| { + for ignored in &resolution.ignored { + litellm_tracing::warn!( + "Redis setting {} is not supported by the Rust client and is ignored: {}", + ignored.setting, + ignored.consequence + ); + } + }); + Resolution { + namespace, + ..resolution + } + } + Err(error) => return Ok(memory_only(py, &error)), + }; + let runtime = litellm_host_python::runtime_handle()?; + let events: Arc = Arc::new(ServiceLoggingEvents::new(py, runtime.clone())?); + let host = HostServices { + events: Some(events), + http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?, + environment, + }; + let mut clients = CLIENTS + .lock_py_attached(py) + .unwrap_or_else(PoisonError::into_inner); + if let Some(shared) = clients.iter().find(|shared| shared.resolved == resolved) { + return Ok(Some(shared.client.clone())); + } + let connected = { + let _runtime = runtime.enter(); + litellm_redis_connect::connect( + resolved.clone(), + host, + breaker.map(|breaker| breaker.0), + ) + }; + let client = match connected { + Ok(client) => client, + Err(error) => { + drop(clients); + return Ok(memory_only(py, &error)); + } + }; + clients.push(Shared { + resolved, + client: client.clone(), + }); + Ok(Some(client)) +} + +/// The Rust client for one `RedisCache`, built once and handed to every native store, so a store +/// never sees the redis-py kwargs or the namespace. +#[pyclass(frozen, name = "NativeRedisClient")] +pub(crate) struct NativeRedisClient( + #[expect( + dead_code, + reason = "the native stores built on this client read it, starting in #49525" + )] + pub(crate) RedisClient, +); + +#[pymethods] +impl NativeRedisClient { + /// `None` when Redis is not usable, which leaves every store given it on memory. + #[staticmethod] + #[pyo3(signature = (redis_kwargs, *, namespace = None, breaker = None))] + fn connect( + py: Python<'_>, + redis_kwargs: RedisSpec, + namespace: Option, + breaker: Option, + ) -> PyResult> { + Ok(shared_client(py, redis_kwargs, namespace, breaker)?.map(Self)) + } +} + +/// The breaker state of every native Redis client this process built, so the Prometheus state +/// gauge counts them next to the Python `RedisCache` breakers. +#[pyfunction] +pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> { + let clients = CLIENTS + .lock_py_attached(py) + .unwrap_or_else(PoisonError::into_inner); + clients + .iter() + .filter_map(|shared| shared.client.breaker_state()) + .map(|state| state_label(&state)) + .collect() +} + +fn memory_only(py: Python<'_>, error: &(dyn std::error::Error + 'static)) -> Option { + let report = litellm_error::Report::new(error).to_string(); + crate::logger::capture(py).scope(|| { + litellm_tracing::warn!( + "Redis is not usable from the Rust stores, their state stays in memory: {report}" + ) + }); + None +} diff --git a/litellm-rust/crates/python-bridge/src/redis/events.rs b/litellm-rust/crates/python-bridge/src/redis/events.rs new file mode 100644 index 00000000000..22d22f7f233 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/events.rs @@ -0,0 +1,274 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; +use pyo3::{prelude::*, types::PyModule}; +use tokio::runtime::Handle; + +tokio::task_local! { + static CALL: RedisCall; +} + +/// The Python caller of one store operation. Its context carries `with_service_target`, and its +/// `call_type` and `caller` are the labels `RedisCache` reports, so dashboards keep their series. +#[derive(Clone)] +pub(crate) struct RedisCall { + pub(crate) context: PythonContext, + pub(crate) call_type: &'static str, + pub(crate) caller: Option, +} + +impl RedisCall { + pub(crate) async fn scope(self, operation: F) -> F::Output { + CALL.scope(self, operation).await + } +} + +/// Reports every native Redis call and breaker transition to `ServiceLogging`. +pub(crate) struct ServiceLoggingEvents { + record_call: Arc>, + record_breaker: Arc>, + context: PythonContext, + runtime: Handle, +} + +impl ServiceLoggingEvents { + pub(crate) fn new(py: Python<'_>, runtime: Handle) -> PyResult { + const MODULE: &str = "litellm.rust_bridge.host.redis_events"; + Self::from_module(py, &py.import(MODULE)?, runtime) + } + + fn from_module( + py: Python<'_>, + module: &Bound<'_, PyModule>, + runtime: Handle, + ) -> PyResult { + Ok(Self { + record_call: Arc::new(module.getattr("record_redis_call")?.unbind()), + record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()), + context: PythonContext::capture(py)?, + runtime, + }) + } + + fn report(&self, context: PythonContext, record: Arc>, args: ReportArgs) { + self.runtime.spawn(async move { + let reported = attach_blocking(context, move |py| match args { + ReportArgs::Call(call_type, duration, failure, caller) => record + .bind(py) + .call1((call_type, duration.as_secs_f64(), failure, caller)) + .map(drop), + ReportArgs::Breaker(state) => record.bind(py).call1((state,)).map(drop), + }) + .await + .and_then(|reported| reported); + if let Err(error) = reported { + tracing::warn!(error = %error, "reporting a Redis event to ServiceLogging failed"); + } + }); + } +} + +enum ReportArgs { + Call(&'static str, Duration, Option<&'static str>, Option), + Breaker(&'static str), +} + +fn operation_call_type(operation: Operation) -> &'static str { + match operation { + Operation::Command(name) => name, + Operation::Script => "eval", + Operation::Pipeline { .. } => "pipeline", + } +} + +fn failure_label(outcome: Outcome) -> Option<&'static str> { + match outcome { + Outcome::Succeeded => None, + Outcome::Failed(FailureKind::Timeout) => Some("timeout"), + Outcome::Failed(FailureKind::Transport) => Some("transport"), + Outcome::Failed(FailureKind::Server) => Some("server"), + Outcome::Failed(FailureKind::Decode) => Some("decode"), + Outcome::Failed(FailureKind::BreakerOpen) => Some("breaker_open"), + } +} + +pub(crate) fn state_label(state: &State) -> &'static str { + match state { + State::Closed { .. } => "closed", + State::Open { .. } => "open", + State::HalfOpen => "half_open", + } +} + +impl RedisEvents for ServiceLoggingEvents { + fn call(&self, event: &CallEvent) { + let caller = CALL.try_with(Clone::clone).ok(); + let (context, call_type, caller) = match caller { + Some(call) => (call.context, call.call_type, call.caller), + None => ( + self.context.clone(), + operation_call_type(event.operation), + None, + ), + }; + self.report( + context, + self.record_call.clone(), + ReportArgs::Call( + call_type, + event.duration, + failure_label(event.outcome), + caller, + ), + ); + } + + fn breaker(&self, transition: &Transition) { + self.report( + self.context.clone(), + self.record_breaker.clone(), + ReportArgs::Breaker(state_label(&transition.to)), + ); + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use litellm_host_python::PythonContext; + use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; + use pyo3::{prelude::*, types::PyModule}; + use rstest::rstest; + use tokio::runtime::Handle; + + use super::{RedisCall, ServiceLoggingEvents}; + + const RECORDER: &str = " +import contextvars +target = contextvars.ContextVar('target', default=None) +seen = [] +def record_redis_call(call_type, duration, failure, caller): + seen.append((call_type, round(duration, 3), failure, caller, target.get())) +def record_redis_breaker(state): + seen.append(('breaker', state)) +"; + + fn recorder(name: &str) -> (ServiceLoggingEvents, Py) { + Python::initialize(); + Python::attach(|py| { + let name = std::ffi::CString::new(format!("redis_events_recorder_{name}")).unwrap(); + let module = + PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name) + .unwrap(); + ( + ServiceLoggingEvents::from_module(py, &module, Handle::current()).unwrap(), + module.unbind(), + ) + }) + } + + async fn seen(module: &Py, count: usize) -> String { + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + loop { + let seen: (usize, String) = Python::attach(|py| { + let seen = module.bind(py).getattr("seen").unwrap(); + (seen.len().unwrap(), seen.repr().unwrap().to_string()) + }); + if seen.0 >= count || tokio::time::Instant::now() > deadline { + return seen.1; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn a_scoped_call_reports_the_store_labels_in_the_callers_context() { + let (events, namespace) = recorder("scoped"); + let call = Python::attach(|py| { + namespace + .bind(py) + .getattr("target") + .unwrap() + .call_method1("set", ("spend_counter",)) + .unwrap(); + RedisCall { + context: PythonContext::capture(py).unwrap(), + call_type: "async_increment", + caller: Some("router.py:42".to_owned()), + } + }); + let event = CallEvent { + operation: Operation::Command("INCRBYFLOAT"), + duration: Duration::from_millis(250), + outcome: Outcome::Failed(FailureKind::Timeout), + }; + + call.scope(async { events.call(&event) }).await; + + assert_eq!( + seen(&namespace, 1).await, + "[('async_increment', 0.25, 'timeout', 'router.py:42', 'spend_counter')]" + ); + } + + #[rstest] + #[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")] + #[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")] + #[tokio::test(flavor = "multi_thread")] + async fn breaker_transitions_report_python_state_names( + #[case] to: litellm_circuit_breaker::State, + #[case] expected: &str, + ) { + let (events, namespace) = recorder(expected); + + events.breaker(&litellm_circuit_breaker::Transition { + from: litellm_circuit_breaker::State::HalfOpen, + to, + generation: 1, + }); + + assert_eq!( + seen(&namespace, 1).await, + format!("[('breaker', '{expected}')]") + ); + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn a_transition_on_a_thread_without_a_runtime_is_still_reported() { + let (events, namespace) = recorder("plain_thread"); + + std::thread::spawn(move || { + events.breaker(&litellm_circuit_breaker::Transition { + from: litellm_circuit_breaker::State::HalfOpen, + to: litellm_circuit_breaker::State::HalfOpen, + generation: 1, + }); + }) + .join() + .unwrap(); + + assert_eq!(seen(&namespace, 1).await, "[('breaker', 'half_open')]"); + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn an_unscoped_call_falls_back_to_the_command_name() { + let (events, namespace) = recorder("unscoped"); + + events.call(&CallEvent { + operation: Operation::Pipeline { ops: 3 }, + duration: Duration::from_millis(5), + outcome: Outcome::Succeeded, + }); + + assert_eq!( + seen(&namespace, 1).await, + "[('pipeline', 0.005, None, None, None)]" + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/mod.rs b/litellm-rust/crates/python-bridge/src/redis/mod.rs new file mode 100644 index 00000000000..0a54822bcd7 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod client; +mod events; +pub(crate) mod store; diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs new file mode 100644 index 00000000000..486e8748387 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -0,0 +1,16 @@ +pub(crate) mod request; + +use litellm_host_python::PythonContext; +use pyo3::prelude::*; + +use crate::redis::events::RedisCall; + +/// The labels `RedisCache` reports for this operation, in the caller's context so its +/// `service_target` reaches the service event. +fn call(py: Python<'_>, call_type: &'static str) -> PyResult { + Ok(RedisCall { + context: PythonContext::capture(py)?, + call_type, + caller: None, + }) +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/request.rs b/litellm-rust/crates/python-bridge/src/redis/store/request.rs new file mode 100644 index 00000000000..1ee80bb0467 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/request.rs @@ -0,0 +1,41 @@ +use std::time::Duration; + +use litellm_redis_store::RequestRedis; +use pyo3::prelude::*; + +use super::call; +use crate::execution::run_async_value; + +/// One request's native Redis batches. Python builds one per request scope and passes it to each +/// request-aware store call. +#[pyclass(frozen, name = "NativeRequestRedis")] +pub(crate) struct NativeRequestRedis(pub(crate) RequestRedis); + +#[pymethods] +impl NativeRequestRedis { + #[new] + fn new(post_call_deadline_seconds: f64) -> PyResult { + Ok(Self(RequestRedis::new( + litellm_host_python::runtime_handle()?, + Duration::from_secs_f64(post_call_deadline_seconds.max(0.0)), + ))) + } + + fn flush_pre<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "request_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_pre()).await; + Ok(()) + }) + } + + fn flush_post<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "post_call_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_post()).await; + Ok(()) + }) + } +} diff --git a/litellm-rust/crates/redis-connect/Cargo.toml b/litellm-rust/crates/redis-connect/Cargo.toml index 29b3293b856..72a453d36de 100644 --- a/litellm-rust/crates/redis-connect/Cargo.toml +++ b/litellm-rust/crates/redis-connect/Cargo.toml @@ -17,6 +17,7 @@ litellm-circuit-breaker.workspace = true litellm-redis-types.workspace = true litellm-core-utils.workspace = true litellm-http.workspace = true +rustls.workspace = true litellm-redis-client.workspace = true futures-util.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs index 813ed837639..758cb4acd82 100644 --- a/litellm-rust/crates/redis-connect/src/connect.rs +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -29,6 +29,9 @@ pub fn connect( host: HostServices, breaker: Option, ) -> Result { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } tracing::info!(target = ?resolved.target, namespace = ?resolved.namespace, "redis client configured"); let credentials = credentials(resolved.auth, &host)?; let events = host.events; diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md new file mode 100644 index 00000000000..493297d57d1 --- /dev/null +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -0,0 +1,7 @@ +- One module per Redis family (`router_cooldowns.rs`, `router_usage.rs`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys +- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec +- A read-through family builds its own `MemoryTier` (`tier.rs`, `main`'s `DualCache` over a `RedisCache`) from a `TierPolicy` the host passes in. Its `pub const POLICY` is `main`'s defaults, and the Python host overrides what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `RedisClient` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` +- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not +- `tests/fixtures/router_wire.json` is the wire contract, checked here against real Redis and in `tests/unit/redis/stores/` against Python's own codec. Change both sides together +- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. A family exposes a request-aware variant that declares on it (`declare_active`, `increment_deployment_post_call`) next to its direct method, which runs the same declaration on a batch of its own +- Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages diff --git a/litellm-rust/crates/redis-store/Cargo.toml b/litellm-rust/crates/redis-store/Cargo.toml new file mode 100644 index 00000000000..b5d3c0cf88c --- /dev/null +++ b/litellm-rust/crates/redis-store/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "litellm-redis-store" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-python-compat.workspace = true +litellm-redis-client.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tracing.workspace = true + +[dev-dependencies] +litellm-circuit-breaker.workspace = true +litellm-redis-testing.workspace = true +litellm-secret-value.workspace = true +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs new file mode 100644 index 00000000000..f64ac797cb5 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -0,0 +1,3 @@ +mod request; + +pub use request::RequestRedis; diff --git a/litellm-rust/crates/redis-store/src/request.rs b/litellm-rust/crates/redis-store/src/request.rs new file mode 100644 index 00000000000..dc6153c0da9 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/request.rs @@ -0,0 +1,121 @@ +use std::{ + future::Future, + sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::Duration, +}; + +use futures_util::future::join_all; +use litellm_redis_client::{RedisBatch, RedisClient}; +use tokio::{runtime::Handle, sync::oneshot, task::JoinHandle}; + +/// One request's Redis batches: a pre-call batch for the reads it needs before answering and a +/// post-call batch for the writes nobody waits on, one of each per Redis client. The host that owns +/// the request builds one and passes it down; it knows nothing else about the request. +#[derive(Clone)] +pub struct RequestRedis(Arc); + +struct Shared { + runtime: Handle, + post_call_deadline: Duration, + batches: Mutex>, + follow_ups: Mutex>>, + deadline: Mutex>>, +} + +#[derive(Clone)] +struct ClientBatches { + client: RedisClient, + pre: RedisBatch, + post: RedisBatch, +} + +impl RequestRedis { + /// `runtime` runs the post-call deadline, so declarations may come from any thread, including a + /// Python one that has no Tokio context. + pub fn new(runtime: Handle, post_call_deadline: Duration) -> Self { + Self(Arc::new(Shared { + runtime, + post_call_deadline, + batches: Mutex::new(Vec::new()), + follow_ups: Mutex::new(Vec::new()), + deadline: Mutex::new(None), + })) + } + + pub fn pre(&self, client: &RedisClient) -> RedisBatch { + self.batches_for(client).pre + } + + /// The post-call batch, which goes out on `flush_post` or `post_call_deadline` after the first + /// declaration, whichever comes first. + pub fn post(&self, client: &RedisClient) -> RedisBatch { + self.arm_deadline(); + self.batches_for(client).post + } + + /// Runs work that follows this request's post-call batch, such as mirroring settled totals + /// into memory, on the request's runtime, from any thread. `flush_post` waits for it, so the + /// shutdown drain does not drop a fallback write still in flight. + pub fn spawn(&self, task: impl Future + Send + 'static) { + let follow_up = self.0.runtime.spawn(task); + lock(&self.0.follow_ups).push(follow_up); + } + + pub async fn flush_pre(&self) { + join_all(self.batches().iter().map(|batches| batches.pre.flush())).await; + } + + /// Sends every post-call op declared so far and waits for the follow-ups spawned on them; the + /// shutdown drain calls it too. + pub async fn flush_post(&self) { + lock(&self.0.deadline).take(); + let follow_ups = std::mem::take(&mut *lock(&self.0.follow_ups)); + join_all(self.batches().iter().map(|batches| batches.post.flush())).await; + join_all(follow_ups) + .await + .into_iter() + .filter_map(Result::err) + .for_each(|error| tracing::warn!(%error, "a post-call follow-up failed")); + } + + fn batches(&self) -> Vec { + lock(&self.0.batches).clone() + } + + fn batches_for(&self, client: &RedisClient) -> ClientBatches { + let mut batches = lock(&self.0.batches); + if let Some(existing) = batches + .iter() + .find(|batches| batches.client.same_client(client)) + { + return existing.clone(); + } + let created = ClientBatches { + client: client.clone(), + pre: client.batch(), + post: client.batch(), + }; + batches.push(created.clone()); + created + } + + fn arm_deadline(&self) { + let mut deadline = lock(&self.0.deadline); + if deadline.is_some() { + return; + } + let (disarm, disarmed) = oneshot::channel(); + *deadline = Some(disarm); + let request = self.clone(); + self.0.runtime.spawn(async move { + tokio::select! { + () = tokio::time::sleep(request.0.post_call_deadline) => request.flush_post().await, + _ = disarmed => {} + } + }); + } +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} diff --git a/litellm-rust/crates/redis-store/tests/request.rs b/litellm-rust/crates/redis-store/tests/request.rs new file mode 100644 index 00000000000..d8f439291b0 --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/request.rs @@ -0,0 +1,181 @@ +mod support; + +use std::time::Duration; + +use litellm_redis_client::{Amount, Operation}; +use litellm_redis_store::RequestRedis; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, connect_observed}; +use tokio::{runtime::Handle, sync::oneshot}; + +const TTL: Option = Some(Duration::from_secs(60)); + +// A batch is one pipeline on a standalone Redis and one call per increment on a cluster, like RedisBatch +fn increments(topology: Topology, count: usize) -> Vec { + match topology { + Topology::Standalone => vec![Operation::Pipeline { ops: count }], + Topology::Cluster => vec![Operation::Command("INCRBYFLOAT"); count], + } +} +const LONG_DEADLINE: Duration = Duration::from_secs(60); + +#[rstest] +#[tokio::test] +async fn pre_declarations_on_one_client_share_one_pipeline_and_another_client_gets_its_own( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let (other, other_trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + + let first = request.pre(&client).incr_by_float("a", Amount::Int(1), TTL); + let second = request + .pre(&client.clone()) + .incr_by_float("b", Amount::Int(2), TTL); + let third = request.pre(&other).incr_by_float("c", Amount::Int(3), TTL); + request.flush_pre().await; + + assert_eq!(trips.seen(), increments(topology, 2)); + assert_eq!(other_trips.seen(), increments(topology, 1)); + assert_eq!( + [ + first.resolve().await.unwrap(), + second.resolve().await.unwrap(), + third.resolve().await.unwrap(), + ], + [1.0, 2.0, 3.0] + ); + assert_eq!(trips.seen(), increments(topology, 2)); +} + +#[rstest] +#[tokio::test] +async fn a_post_declaration_waits_for_flush_post_and_flush_pre_leaves_it_alone( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + + let total = tokio::spawn( + request + .post(&client) + .incr_by_float("spend", Amount::Float(2.5), TTL) + .settled(), + ); + request.flush_pre().await; + tokio::time::sleep(Duration::from_millis(200)).await; + + assert!(!total.is_finished()); + assert!(trips.seen().is_empty()); + assert_eq!(raw.get::("spend").await.unwrap(), None); + + request.flush_post().await; + + assert_eq!(total.await.unwrap().unwrap(), 2.5); + assert_eq!(trips.seen(), increments(topology, 1)); +} + +#[rstest] +#[tokio::test] +async fn flush_post_returns_after_the_follow_ups_of_the_ops_it_sent( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&deployment); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + let total = request + .post(&client) + .incr_by_float("spend", Amount::Float(2.5), TTL); + let (mirror, mut mirrored) = oneshot::channel(); + request.spawn(async move { + mirror.send(total.settled().await.unwrap()).unwrap(); + }); + + request.flush_post().await; + + assert_eq!(mirrored.try_recv(), Ok(2.5)); +} + +#[rstest] +#[tokio::test] +async fn a_post_batch_nobody_flushes_goes_out_at_the_deadline( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), Duration::from_millis(100)); + + let first = request + .post(&client) + .incr_by_float("a", Amount::Int(1), TTL); + let second = request + .post(&client) + .incr_by_float("b", Amount::Int(2), TTL); + let settled = tokio::time::timeout(Duration::from_secs(5), async { + (first.settled().await, second.settled().await) + }) + .await + .unwrap(); + + assert_eq!((settled.0.unwrap(), settled.1.unwrap()), (1.0, 2.0)); + assert_eq!(trips.seen(), increments(topology, 2)); +} + +#[rstest] +#[tokio::test] +async fn a_post_declaration_from_a_thread_without_a_runtime_goes_out_at_the_deadline( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), Duration::from_millis(100)); + + let declared = std::thread::spawn({ + let request = request.clone(); + move || { + request + .post(&client) + .incr_by_float("spend", Amount::Float(2.5), TTL) + } + }) + .join() + .expect("a Python thread has no Tokio context, and declaring there must not panic"); + let total = tokio::time::timeout(Duration::from_secs(5), declared.settled()) + .await + .unwrap(); + + assert_eq!(total.unwrap(), 2.5); + assert_eq!(trips.seen(), increments(topology, 1)); +} + +#[rstest] +#[tokio::test] +async fn a_post_declaration_after_a_flush_arms_a_new_deadline( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), Duration::from_millis(300)); + + let early = request + .post(&client) + .incr_by_float("a", Amount::Int(1), TTL); + request.flush_post().await; + tokio::time::sleep(Duration::from_millis(200)).await; + let late = request + .post(&client) + .incr_by_float("a", Amount::Int(1), TTL); + tokio::time::sleep(Duration::from_millis(200)).await; + let flushed_before_its_deadline = trips.seen().len() > 1; + let late = tokio::time::timeout(Duration::from_secs(5), late.settled()) + .await + .unwrap(); + + assert_eq!(early.settled().await.unwrap(), 1.0); + assert!(!flushed_before_its_deadline); + assert_eq!(late.unwrap(), 2.0); +} diff --git a/litellm-rust/crates/redis-store/tests/support/mod.rs b/litellm-rust/crates/redis-store/tests/support/mod.rs new file mode 100644 index 00000000000..f026aadeb7b --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/support/mod.rs @@ -0,0 +1,83 @@ +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_redis_client::{ + CallEvent, Cluster, Credentials, Namespace, Node, Operation, RedisClient, RedisConfig, + RedisEvents, RedisUrl, Target, Timeouts, +}; +use litellm_redis_testing::RedisDeployment; + +pub fn connect(deployment: &RedisDeployment) -> RedisClient { + connect_with(deployment, None, None) +} + +/// Every round trip a client made, as the client reports it. +#[derive(Default)] +pub struct RoundTrips(Mutex>); + +impl RoundTrips { + pub fn seen(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +impl RedisEvents for RoundTrips { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(event.operation); + } + + fn breaker(&self, _transition: &Transition) {} +} + +pub fn connect_observed(deployment: &RedisDeployment) -> (RedisClient, Arc) { + let trips = Arc::new(RoundTrips::default()); + (connect_with(deployment, None, Some(trips.clone())), trips) +} + +pub fn connect_guarded( + deployment: &RedisDeployment, + breaker: litellm_circuit_breaker::Breaker, +) -> RedisClient { + connect_with(deployment, Some(breaker), None) +} + +fn connect_with( + deployment: &RedisDeployment, + breaker: Option, + events: Option>, +) -> RedisClient { + let endpoint = match deployment { + RedisDeployment::Standalone(server) => Target::Url(RedisUrl { + url: litellm_secret_value::SecretValue::new(server.url().to_owned()), + tls: None, + db: None, + username: None, + password: None, + }), + RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster { + nodes: cluster + .nodes() + .map(|(host, port)| Node { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + }; + RedisClient::new(RedisConfig { + target: endpoint, + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(500), None), + breaker, + events, + health_check: None, + }) + .unwrap() +} diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index aa458c9c926..dd886b09067 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -20,7 +20,7 @@ from contextvars import ContextVar, Token from dataclasses import dataclass, field from datetime import timedelta from types import MappingProxyType, TracebackType -from typing import Final, Generic, Protocol, TypeVar +from typing import TYPE_CHECKING, Final, Generic, Protocol, TypeAlias, TypeVar from litellm._internal_context import ( REDIS_FAMILIES_METADATA_KEY, @@ -38,7 +38,11 @@ from litellm.caching.redis_cache import ( from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.types.services import ServiceTypes +if TYPE_CHECKING: + from litellm.rust_bridge._native import NativeRequestRedis + _T = TypeVar("_T") +_T_co: Final = TypeVar("_T_co", covariant=True) _ScriptArg = str | bytes | int | float SettledHook = Callable[[asyncio.Future[_T]], Awaitable[None] | None] POST_CALL_FLUSH_DEADLINE_SECONDS: Final = 1.0 @@ -443,6 +447,37 @@ def _backend_key(redis_cache: RedisCache) -> object: return (type(redis_cache), redis_cache.namespace, settings) +class NativePending(Protocol[_T_co]): + def resolve(self) -> Awaitable[_T_co]: ... + + +class NativeBatchResult(Generic[_T]): + """Awaitable handle for a read declared on the request's native pre-call batch; awaiting it runs the + native flush hooks, flushes that batch, then answers from what it read.""" + + __slots__ = ("_pending", "_request") + + def __init__(self, request: RequestRedisBatches, pending: NativePending[_T]) -> None: + self._request: Final = request + self._pending: Final = pending + + def __await__(self) -> Generator[object, None, _T]: + return self._wait().__await__() + + async def _wait(self) -> _T: + await self._request.flush_native_pre() + return await self._pending.resolve() + + +NativeRequestFactory: TypeAlias = Callable[[float], "NativeRequestRedis"] + + +def _native_request(post_call_deadline: float) -> NativeRequestRedis: + from litellm.rust_bridge._native import NativeRequestRedis + + return NativeRequestRedis(post_call_deadline) + + _open_post_call: Final[weakref.WeakSet[RequestRedisBatches]] = weakref.WeakSet() """Requests whose post-call batch still holds declared ops, so a shutdown can send them before Redis goes away.""" @@ -460,15 +495,25 @@ class RequestRedisBatches: "_batches", "_deadline", "_deadline_flush", + "_native", + "_native_factory", + "_native_flush_hooks", "_post_call", "post_call_deadline", "prefetched", ) - def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None: + def __init__( + self, + post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS, + native_factory: NativeRequestFactory = _native_request, + ) -> None: self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend self.post_call_deadline: Final = post_call_deadline + self._native_factory: Final = native_factory + self._native: NativeRequestRedis | None = None + self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry self._deadline: asyncio.TimerHandle | None = None self._deadline_flush: asyncio.Task[None] | None = None # Reads declared early for a consumer that runs later in the request, keyed by consumer name. @@ -495,20 +540,49 @@ class RequestRedisBatches: _open_post_call.add(self) return batch + def native(self) -> NativeRequestRedis: + """The request's native batches, built on first use so requests that touch no native store pay nothing.""" + if self._native is None: + self._native = self._native_factory(self.post_call_deadline) + return self._native + + def native_post_call(self) -> NativeRequestRedis: + """``native`` for a post-call declaration, so a shutdown drain also flushes it.""" + _open_post_call.add(self) + return self.native() + + def native_result(self, pending: NativePending[_T]) -> NativeBatchResult[_T]: + return NativeBatchResult(self, pending) + + def add_native_flush_hook(self, hook: Callable[[], None]) -> None: + """Called before every native pre-call flush so late readers can declare into the same trip.""" + self._native_flush_hooks.append(hook) + + async def flush_native_pre(self) -> None: + for hook in self._native_flush_hooks: + hook() + if self._native is not None: + await self._native.flush_pre() + def _flush_on_deadline(self) -> None: self._deadline = None self._deadline_flush = asyncio.ensure_future(self.flush_post_call()) async def flush_all(self) -> None: """Send whatever is still declared (write-backs nobody awaits) before the request scope closes.""" - await asyncio.gather(*(batch.flush() for batch in self._batches.values() if batch.pending)) + await asyncio.gather( + *(batch.flush() for batch in self._batches.values() if batch.pending), self.flush_native_pre() + ) async def flush_post_call(self) -> None: """One pipeline per backend for the post-call writes; the deadline is disarmed since this is that flush.""" if self._deadline is not None: self._deadline.cancel() self._deadline = None - await asyncio.gather(*(batch.flush() for batch in self._post_call.values() if batch.pending)) + await asyncio.gather( + *(batch.flush() for batch in self._post_call.values() if batch.pending), + *(() if self._native is None else (self._native.flush_post(),)), + ) if not any(batch.pending for batch in self._post_call.values()): _open_post_call.discard(self) @@ -557,17 +631,24 @@ async def drain_post_call_redis_batches() -> None: class request_redis_batch_scope: """Redis reads declared inside share one pipeline per backend; nested scopes join the outer one.""" - __slots__ = ("_post_call_deadline", "_token") + __slots__ = ("_native_factory", "_post_call_deadline", "_token") - def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None: + def __init__( + self, + post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS, + native_factory: NativeRequestFactory = _native_request, + ) -> None: self._token: Token[RequestRedisBatches | None] | None = None self._post_call_deadline: Final = post_call_deadline + self._native_factory: Final = native_factory def __enter__(self) -> RequestRedisBatches: outer: Final = _active_request_batches.get() if outer is not None: return outer - batches: Final = RequestRedisBatches(post_call_deadline=self._post_call_deadline) + batches: Final = RequestRedisBatches( + post_call_deadline=self._post_call_deadline, native_factory=self._native_factory + ) self._token = _active_request_batches.set(batches) return batches diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index afe6f8a1eab..66c9a1c6771 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -472,8 +472,14 @@ _LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet() def redis_breaker_state_counts() -> Mapping[str, int]: - """How many live Redis circuit breakers sit in each state, for the Prometheus services logger.""" - return MappingProxyType(dict(Counter(breaker.state for breaker in _LIVE_BREAKERS))) + """How many live Redis circuit breakers sit in each state, for the Prometheus services logger. + + Counts the Python breakers of every live ``RedisCache`` and the breaker of every native Redis + client the router stores built, since each opens and closes on its own.""" + from litellm.rust_bridge._native import native_redis_breaker_states + + python_states: Final = Counter(breaker.state for breaker in _LIVE_BREAKERS) + return MappingProxyType(dict(python_states + Counter(native_redis_breaker_states()))) class RedisCircuitBreakerOpenError(Exception): diff --git a/litellm/redis/__init__.py b/litellm/redis/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/redis/stores/__init__.py b/litellm/redis/stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/redis/stores/_connection.py b/litellm/redis/stores/_connection.py new file mode 100644 index 00000000000..c6cace43b9f --- /dev/null +++ b/litellm/redis/stores/_connection.py @@ -0,0 +1,40 @@ +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter + +from litellm._redis import _merge_redis_kwargs +from litellm.caching.redis_cache import RedisCache +from litellm.constants import ( + REDIS_CIRCUIT_BREAKER_ENABLED, + REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, + REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, + REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, +) +from litellm.rust_bridge._native import NativeRedisClient + +_KWARGS: Final = TypeAdapter(dict[str, object]) + + +def _redis_kwargs(connection: RedisCache) -> Mapping[str, object]: + return _merge_redis_kwargs( + _KWARGS.validate_python( + connection.redis_kwargs # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped legacy dict + ) + ) + + +def _breaker() -> tuple[int, float, float] | None: + if not REDIS_CIRCUIT_BREAKER_ENABLED: + return None + return ( + REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, + float(REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT), + REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, + ) + + +def native_client(connection: RedisCache | None) -> NativeRedisClient | None: + if connection is None: + return None + return NativeRedisClient.connect(_redis_kwargs(connection), namespace=connection.namespace, breaker=_breaker()) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 954815f6a4f..b35752cfa34 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -15,6 +15,7 @@ class ProcessReservedForForking(RuntimeError): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... +def native_redis_breaker_states() -> list[str]: ... @final class NativeClickHouseSpendConfig: @@ -26,6 +27,23 @@ class NativeClickHouseSpendStorage: def ensure_schema(self) -> Future[None]: ... def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... +@final +class NativeRedisClient: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @staticmethod + def connect( + redis_kwargs: Mapping[str, object], + *, + namespace: str | None = None, + breaker: tuple[int, float, float] | None = None, + ) -> NativeRedisClient | None: ... + +@final +class NativeRequestRedis: + def __new__(cls, post_call_deadline_seconds: float) -> NativeRequestRedis: ... + def flush_pre(self) -> Future[None]: ... + def flush_post(self) -> Future[None]: ... + @final class NativeDatabase: @staticmethod @@ -216,6 +234,8 @@ __all__ = [ "NativeClickHouseSpendStorage", "NativeDatabase", "NativeDiagnosticProcessor", + "NativeRedisClient", + "NativeRequestRedis", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", @@ -233,6 +253,7 @@ __all__ = [ "embedding", "gil_stats", "messages", + "native_redis_breaker_states", "ocr", "ocr_health_check_document", "ocr_passthrough_response", diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index 9b751a1ada2..d4cb25cb4fd 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -2,7 +2,7 @@ Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name -- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds) +- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`) - A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here - Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers - The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green diff --git a/litellm/rust_bridge/host/redis_events.py b/litellm/rust_bridge/host/redis_events.py new file mode 100644 index 00000000000..202d43cd4fd --- /dev/null +++ b/litellm/rust_bridge/host/redis_events.py @@ -0,0 +1,46 @@ +"""Redis call and circuit breaker events from the native Redis client, reported the way ``RedisCache`` reports them.""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Final, Literal, TypeAlias + +from litellm._service_logger import ServiceLogging +from litellm.types.services import ServiceTypes + +Failure: TypeAlias = Literal["timeout", "transport", "server", "decode", "breaker_open"] +BreakerState: TypeAlias = Literal["closed", "open", "half_open"] + +_HEALTH_FAILURE_CLASSES: Final = {"timeout": "timeout", "transport": "connectivity"} + + +class RedisCallError(Exception): + """The failure class of a native Redis call, the only detail the Rust client hands over.""" + + +@lru_cache(maxsize=1) +def _service_logging() -> ServiceLogging: + return ServiceLogging() + + +def record_redis_call(call_type: str, duration: float, failure: Failure | None, caller: str | None) -> None: + if failure == "breaker_open": + return + logger: Final = _service_logging() + if failure is None: + logger.service_success_hook(service=ServiceTypes.REDIS, duration=duration, call_type=call_type, caller=caller) + return + logger.service_failure_hook( + service=ServiceTypes.REDIS, + duration=duration, + error=RedisCallError(failure), + call_type=call_type, + caller=caller, + ) + failure_class: Final = _HEALTH_FAILURE_CLASSES.get(failure) + if failure_class is not None: + logger.redis_breaker_hook(transitions_to=None, failure_class=failure_class) + + +def record_redis_breaker(state: BreakerState) -> None: + _service_logging().redis_breaker_hook(transitions_to=state, failure_class=None) diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index cd270035416..7745bafe1b5 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -6,8 +6,9 @@ import asyncio import hashlib import json from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass, field from datetime import timedelta -from typing import Any +from typing import Any, Generic, TypeVar import pytest from redis.exceptions import NoScriptError @@ -28,6 +29,7 @@ from litellm.caching.redis_cache import ( from litellm.caching.redis_cluster_cache import RedisClusterCache SCRIPT = "return redis.call('GET', KEYS[1])" +_Answer = TypeVar("_Answer") SHA = hashlib.sha1(SCRIPT.encode()).hexdigest() # noqa: S324 @@ -93,6 +95,7 @@ class FakeRedisCache(RedisCache): def __init__(self, client: FakeClient, namespace: str | None = None) -> None: # super().__init__ needs a server self.client = client self.namespace = namespace + self.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2, "db": id(self) % 1_000_000} self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) self.service_logger_obj = ServiceLogging() self.default_ttl = None @@ -147,6 +150,43 @@ def replies(command: tuple[Any, ...]) -> Any: raise AssertionError(command) +@dataclass +class FakeNativeRequest: + """Stands in for ``NativeRequestRedis``: records each flush in ``events`` next to what the test's pendings log.""" + + post_call_deadline: float + events: list[str] = field(default_factory=list) + + async def flush_pre(self) -> None: + self.events.append("flush_pre") + + async def flush_post(self) -> None: + self.events.append("flush_post") + + +@dataclass +class FakeNatives: + built: list[FakeNativeRequest] = field(default_factory=list) + + def __call__(self, post_call_deadline: float) -> FakeNativeRequest: + self.built.append(FakeNativeRequest(post_call_deadline)) + return self.built[-1] + + +@dataclass +class FakeNativePending(Generic[_Answer]): + """A read declared on a ``FakeNativeRequest``; it only answers once that request's pre-call batch was flushed.""" + + request: FakeNativeRequest + answer: Callable[[], _Answer] + + async def resolve(self) -> _Answer: + if "flush_pre" not in self.request.events: + raise ConnectionError("resolved before its batch was flushed") + self.request.events.append("resolve") + return self.answer() + + def make(fail: Exception | None = None, namespace: str | None = None) -> tuple[FakeRedisCache, FakeClient]: client = FakeClient(replies, fail) return FakeRedisCache(client, namespace), client diff --git a/tests/unit/redis/__init__.py b/tests/unit/redis/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/redis/stores/__init__.py b/tests/unit/redis/stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/redis/stores/test_connection.py b/tests/unit/redis/stores/test_connection.py new file mode 100644 index 00000000000..9eb7707b18d --- /dev/null +++ b/tests/unit/redis/stores/test_connection.py @@ -0,0 +1,28 @@ +import uuid +from concurrent.futures import ThreadPoolExecutor +from typing import Final + + +from litellm.rust_bridge._native import NativeRedisClient, native_redis_breaker_states + + +def test_concurrent_first_connects_share_one_client_without_deadlocking() -> None: + before: Final = len(native_redis_breaker_states()) + namespace: Final = f"concurrent-connect-{uuid.uuid4()}" + pool: Final = ThreadPoolExecutor(max_workers=8) + futures: Final = tuple( + pool.submit( + NativeRedisClient.connect, + {"host": "127.0.0.1", "port": 1}, + namespace=namespace, + breaker=(5, 60.0, 5.0), + ) + for _ in range(8) + ) + try: + clients: Final = tuple(future.result(timeout=30) for future in futures) + finally: + pool.shutdown(wait=False) + + assert all(client is not None for client in clients) + assert len(native_redis_breaker_states()) == before + 1 diff --git a/tests/unit/rust_bridge/host/test_redis_events.py b/tests/unit/rust_bridge/host/test_redis_events.py new file mode 100644 index 00000000000..6403d0b43c4 --- /dev/null +++ b/tests/unit/rust_bridge/host/test_redis_events.py @@ -0,0 +1,81 @@ +from typing import Final + +import pytest +from prometheus_client import REGISTRY + +import litellm +from litellm.rust_bridge.host import redis_events + + +def _sample(name: str, labels: dict[str, str]) -> float: + return REGISTRY.get_sample_value(name, labels) or 0.0 + + +@pytest.fixture +def prometheus(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + redis_events._service_logging.cache_clear() + yield + redis_events._service_logging.cache_clear() + + +@pytest.mark.parametrize( + ("failure", "failure_class"), + [("timeout", "timeout"), ("transport", "connectivity")], +) +def test_a_health_failure_counts_as_a_failed_request_and_a_breaker_failure( + prometheus: None, failure: redis_events.Failure, failure_class: str +): + failed: Final = ( + "litellm_redis_failed_requests_total", + {"redis": "redis", "error_class": "RedisCallError", "function_name": "async_increment"}, + ) + breaker: Final = ("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class}) + before: Final = (_sample(*failed), _sample(*breaker)) + + redis_events.record_redis_call("async_increment", 0.25, failure, "caller.py") + + assert (_sample(*failed), _sample(*breaker)) == (before[0] + 1, before[1] + 1) + + +@pytest.mark.parametrize("failure", ["server", "decode"]) +def test_a_rejected_command_is_a_failed_request_but_not_a_breaker_failure( + prometheus: None, failure: redis_events.Failure +): + breaker: Final = { + failure_class: _sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class}) + for failure_class in ("timeout", "connectivity") + } + failed: Final = ( + "litellm_redis_failed_requests_total", + {"redis": "redis", "error_class": "RedisCallError", "function_name": "async_increment"}, + ) + before: Final = _sample(*failed) + + redis_events.record_redis_call("async_increment", 0.25, failure, None) + + assert _sample(*failed) == before + 1 + assert { + failure_class: _sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class}) + for failure_class in ("timeout", "connectivity") + } == breaker + + +def test_a_success_is_counted_and_a_breaker_rejection_is_not(prometheus: None): + total: Final = ("litellm_redis_total_requests_total", {"redis": "redis"}) + before: Final = _sample(*total) + + redis_events.record_redis_call("async_get_cache", 0.01, None, None) + redis_events.record_redis_call("async_get_cache", 0.0, "breaker_open", None) + + assert _sample(*total) == before + 1 + + +@pytest.mark.parametrize("state", ["open", "half_open", "closed"]) +def test_breaker_transitions_are_counted_by_state(prometheus: None, state: redis_events.BreakerState): + labels: Final = {"state": state} + before: Final = _sample("litellm_redis_circuit_breaker_transitions_total", labels) + + redis_events.record_redis_breaker(state) + + assert _sample("litellm_redis_circuit_breaker_transitions_total", labels) == before + 1