diff --git a/litellm-rust/crates/gateway/src/redis/mod.rs b/litellm-rust/crates/gateway/src/redis/mod.rs index 5bf8ea96250..2eff0c77cbf 100644 --- a/litellm-rust/crates/gateway/src/redis/mod.rs +++ b/litellm-rust/crates/gateway/src/redis/mod.rs @@ -261,7 +261,7 @@ mod tests { use std::{collections::HashMap, time::Duration}; use litellm_config::Config; - use litellm_redis_client::{Standalone, StaticAuth}; + use litellm_redis_client::{Standalone, StaticAuth, commands}; use litellm_redis_testing::RedisServer; use litellm_redis_types::{Auth, HealthCheck, Passthrough}; use rstest::rstest; @@ -400,16 +400,16 @@ mod tests { .unwrap(); assert_eq!( - cache.get::("rate-limit").await.unwrap().as_deref(), + cache.send(commands::get::("rate-limit")).await.unwrap().as_deref(), Some("1") ); assert_eq!( - cache.get::("spend").await.unwrap().as_deref(), + cache.send(commands::get::("spend")).await.unwrap().as_deref(), Some("2") ); assert_eq!( unnamespaced - .get::("responses:rate-limit") + .send(commands::get::("responses:rate-limit")) .await .unwrap() .as_deref(), @@ -417,7 +417,7 @@ mod tests { ); assert_eq!( unnamespaced - .get::("responses:spend") + .send(commands::get::("responses:spend")) .await .unwrap() .as_deref(), @@ -425,12 +425,12 @@ mod tests { ); assert!( unnamespaced - .get::("rate-limit") + .send(commands::get::("rate-limit")) .await .unwrap() .is_none() ); - assert!(unnamespaced.get::("spend").await.unwrap().is_none()); + assert!(unnamespaced.send(commands::get::("spend")).await.unwrap().is_none()); } #[rstest] diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 79ab9aa5b40..71afcdfc753 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -45,7 +45,7 @@ mod _native { #[pymodule_export] use crate::redis::store::rate_limits::NativeRateLimits; #[pymodule_export] - use crate::redis::store::request::NativeRequestRedis; + use crate::redis::store::request::{NativeRequestRedis, NativeSend}; #[pymodule_export] use crate::redis::store::router_usage::NativeRouterUsage; #[pymodule_export] @@ -129,6 +129,7 @@ mod tests { "NativeDiagnosticProcessor", "NativeRedisClient", "NativeRequestRedis", + "NativeSend", "NativeAuthCache", "NativeRateLimits", "NativePending", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs index 4bc800b4a23..153ccaecec7 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs @@ -1,26 +1,15 @@ -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; use litellm_python_compat::Value; -use litellm_redis_store::auth_cache::{AuthCache, DeclaredWrite, Error, Settings}; +use litellm_redis_store::auth_cache::{AuthCache, Settings}; use pyo3::prelude::*; use super::{ - NativePending, PyNumber, kind_error, redis_async, redis_sync, request::NativeRequestRedis, - store_error, ttl, value::PyValue, + NativePending, PyNumber, kind_error, pending, redis_async, redis_sync, remote, request::NativeSend, + ttl, value::PyValue, }; use crate::redis::client::NativeRedisClient; -fn auth_cache_error(error: Error) -> PyErr { - kind_error(error.kind(), &error) -} - -/// A `set_pre_call` or `delete_pre_call`, settled once and raising when Redis refused it. -fn declared_write(declared: DeclaredWrite) -> NativePending { - NativePending::new("async_set_cache", auth_cache_error, async move { - declared.resolve().await - }) -} - fn values(values: Vec>) -> Vec> { values.into_iter().map(|value| value.map(PyValue)).collect() } @@ -41,11 +30,10 @@ impl NativeAuthCache { redis_ttl_seconds: u64, ) -> Self { Self(Arc::new(AuthCache::new( - client.map(|client| client.0.clone()), + client.map(|client| remote(&client, redis_ttl_seconds)), Settings { memory_entries, memory_ttl: ttl(memory_ttl_seconds), - redis_ttl: Duration::from_secs(redis_ttl_seconds), }, ))) } @@ -57,14 +45,14 @@ impl NativeAuthCache { fn get<'py>(&self, py: Python<'py>, key: String) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_get_cache", store_error, async move { + redis_async(py, "async_get_cache", kind_error, async move { store.get(&key).await.map(|value| value.map(PyValue)) }) } fn get_sync(&self, py: Python<'_>, key: String) -> PyResult> { let store = self.0.clone(); - redis_sync(py, "get_cache", store_error, async move { + redis_sync(py, "get_cache", kind_error, async move { store.get(&key).await.map(|value| value.map(PyValue)) }) } @@ -78,7 +66,7 @@ impl NativeAuthCache { ttl_seconds: Option, ) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_set_cache", store_error, async move { + redis_async(py, "async_set_cache", kind_error, async move { store.set(&key, value.0, ttl(ttl_seconds)).await }) } @@ -92,37 +80,30 @@ impl NativeAuthCache { ttl_seconds: Option, ) -> PyResult<()> { let store = self.0.clone(); - redis_sync(py, "set_cache", store_error, async move { + redis_sync(py, "set_cache", kind_error, async move { store.set_str(&key, value.0, ttl(ttl_seconds)).await }) } - /// Memory now and the `SET` on the request's pre-call batch; `None` without Redis. - #[pyo3(signature = (request, key, value, ttl_seconds = None))] - fn set_pre_call( + /// Memory now and the `SET` on the batch `send` picks. + #[pyo3(signature = (key, value, ttl_seconds, send))] + fn declare_set( &self, - request: &NativeRequestRedis, key: String, value: PyValue, ttl_seconds: Option, - ) -> PyResult> { - Ok(self + send: &NativeSend, + ) -> PyResult { + let written = self .0 - .set_pre_call(&request.0, &key, value.0, ttl(ttl_seconds)) - .map_err(store_error)? - .map(declared_write)) + .declare_set(send.send(), &key, value.0, ttl(ttl_seconds)) + .map_err(kind_error)?; + Ok(pending("async_set_cache", kind_error, written, std::convert::identity)) } - fn delete_pre_call( - &self, - request: &NativeRequestRedis, - key: String, - ) -> PyResult> { - Ok(self - .0 - .delete_pre_call(&request.0, &key) - .map_err(store_error)? - .map(declared_write)) + fn declare_delete(&self, key: String, send: &NativeSend) -> PyResult { + let deleted = self.0.declare_delete(send.send(), &key).map_err(kind_error)?; + Ok(pending("async_set_cache", kind_error, deleted, std::convert::identity)) } #[pyo3(signature = (entries, ttl_seconds = None))] @@ -137,70 +118,45 @@ impl NativeAuthCache { .into_iter() .map(|(key, value)| (key, value.0)) .collect(); - redis_async(py, "async_set_cache_pipeline", store_error, async move { + redis_async(py, "async_set_cache_pipeline", kind_error, async move { store.set_many(entries, ttl(ttl_seconds)).await }) } - /// On the request's pre-call batch when `request` is given, returning before anything is sent. + /// On a request batch it returns before anything is sent, and `Now` is a pipeline of its own. fn write_back<'py>( &self, py: Python<'py>, - request: Option<&NativeRequestRedis>, entries: Vec<(String, PyValue, Option)>, + send: &NativeSend, ) -> PyResult> { - let store = self.0.clone(); - let request = request.map(|request| request.0.clone()); let entries = entries .into_iter() .map(|(key, value, seconds)| (key, value.0, ttl(seconds))) .collect(); - redis_async( - py, - "async_set_cache_pipeline_with_ttls", - auth_cache_error, - async move { store.write_back(request.as_ref(), entries).await }, - ) + let written = self.0.write_back(send.send(), entries).map_err(kind_error)?; + redis_async(py, "async_set_cache_pipeline_with_ttls", kind_error, written) } fn delete<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_delete_cache", auth_cache_error, async move { + redis_async(py, "async_delete_cache", kind_error, async move { store.delete(&keys).await }) } fn delete_sync(&self, py: Python<'_>, keys: Vec) -> PyResult<()> { let store = self.0.clone(); - redis_sync(py, "delete_cache", auth_cache_error, async move { + redis_sync(py, "delete_cache", kind_error, async move { store.delete(&keys).await }) } - /// Reads of keys the caller found missing in memory, on the request's pre-call batch; `None` + /// One `MGET` of keys the caller found missing in memory, on the batch `send` picks; `None` /// without Redis. - fn declare_get_many( - &self, - request: &NativeRequestRedis, - keys: Vec, - ) -> Option { - self.0.declare_get_many(&request.0, &keys).map(|declared| { - NativePending::new("async_batch_get_cache", auth_cache_error, async move { - declared.resolve().await.map(values) - }) - }) - } - - /// One round trip for keys missing in memory; resolves to `None` without Redis. - fn get_many<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { - let store = self.0.clone(); - redis_async(py, "async_batch_get_cache", auth_cache_error, async move { - store - .get_many(&keys) - .await - .transpose() - .map(|read| read.map(values)) - }) + fn get_many(&self, keys: Vec, send: &NativeSend) -> Option { + let read = self.0.get_many(send.send(), &keys)?; + Some(pending("async_batch_get_cache", kind_error, read, values)) } fn increment_if_present<'py>( @@ -213,32 +169,32 @@ impl NativeAuthCache { redis_async( py, "async_increment_if_present", - auth_cache_error, + kind_error, async move { store.increment_if_present(&key, amount.0).await }, ) } fn get_local(&self, key: &str) -> PyResult> { - Ok(self.0.get_local(key).map_err(store_error)?.map(PyValue)) + Ok(self.0.get_local(key).map_err(kind_error)?.map(PyValue)) } #[pyo3(signature = (key, value, ttl_seconds = None))] fn set_local(&self, key: &str, value: PyValue, ttl_seconds: Option) -> PyResult<()> { self.0 .set_local(key, value.0, ttl(ttl_seconds)) - .map_err(store_error) + .map_err(kind_error) } fn delete_local(&self, key: &str) -> PyResult<()> { - self.0.delete_local(key).map_err(store_error) + self.0.delete_local(key).map_err(kind_error) } fn flush_local(&self) -> PyResult<()> { - self.0.flush_local().map_err(store_error) + self.0.flush_local().map_err(kind_error) } /// `(management objects, user-key objects)` held in memory. fn len_local(&self) -> PyResult<(usize, usize)> { - self.0.len_local().map_err(store_error) + self.0.len_local().map_err(kind_error) } } 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 0dcd65dee17..a35731931ed 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -7,6 +7,8 @@ pub(crate) mod router_usage; pub(crate) mod spend_counters; mod value; +use value::PyValue; + use std::{ pin::Pin, sync::{Mutex, PoisonError}, @@ -14,16 +16,16 @@ use std::{ }; use litellm_host_python::PythonContext; -use litellm_redis_store::{ErrorKind, Number}; +use litellm_redis_store::{Error, FromValue, IntoValue, Number, Pending, Remote}; use pyo3::{ IntoPyObjectExt, - exceptions::{PyRuntimeError, PyTypeError}, + exceptions::{PyRuntimeError, PyTypeError, PyValueError}, prelude::*, types::{PyBool, PyFloat, PyInt}, }; use crate::execution::{run_async_value, run_sync_value}; -use crate::redis::events::RedisCall; +use crate::redis::{client::NativeRedisClient, events::RedisCall}; /// The labels `RedisCache` reports for this operation, in the caller's context so its /// `service_target` reaches the service event. @@ -73,6 +75,14 @@ fn ttl(seconds: Option) -> Option { seconds.map(|seconds| Duration::from_secs_f64(seconds.max(0.0))) } +/// A store's Redis with the `RedisCache.default_ttl` Python attached it with. +fn remote(client: &NativeRedisClient, redis_ttl_seconds: u64) -> Remote { + Remote { + client: client.0.clone(), + default_ttl: Duration::from_secs(redis_ttl_seconds), + } +} + type Resolve = Pin>> + Send>>; /// A Redis call declared on a request's batch. `resolve` flushes that batch if it has not gone @@ -84,7 +94,7 @@ pub(crate) struct NativePending { } impl NativePending { - fn new( + pub(crate) fn new( call_type: &'static str, error: fn(E) -> PyErr, reply: impl Future> + Send + 'static, @@ -123,15 +133,31 @@ pyo3::create_exception!( "A native Redis store operation failed. Args are (kind, message), kind spelled as litellm-redis-store's ErrorKind." ); -fn kind_error(kind: ErrorKind, error: &(dyn std::error::Error + 'static)) -> PyErr { - let kind: &'static str = kind.into(); - NativeRedisStoreError::new_err((kind, litellm_error::Report::new(error).to_string())) +/// What the spend counter and auth cache stores raise. +fn kind_error(error: Error) -> PyErr { + let kind: &'static str = error.kind().into(); + NativeRedisStoreError::new_err((kind, litellm_error::Report::new(&error).to_string())) } -fn store_error(error: litellm_cache::Error) -> PyErr { +/// What the router stores and the rate limiter raise. +fn runtime_error(error: Error) -> PyErr { PyRuntimeError::new_err(litellm_error::Report::new(&error).to_string()) } +/// A store op declared on the batch the caller's `NativeSend` picked, answered as `reply` shapes it. +fn pending( + call_type: &'static str, + error: fn(Error) -> PyErr, + declared: Pending, + reply: fn(T) -> R, +) -> NativePending +where + T: Send + 'static, + R: for<'a> IntoPyObject<'a> + Send + 'static, +{ + NativePending::new(call_type, error, async move { declared.await.map(reply) }) +} + /// A Python `int` or `float` across the bridge, keeping which one it is. pub(crate) struct PyNumber(Number); @@ -161,3 +187,32 @@ impl<'py> IntoPyObject<'py> for PyNumber { }) } } + +/// A store type whose Python side is plain data, converted through its `FromValue` and `IntoValue` +/// so the bridge never names a field. A shape the store refuses is `ValueError`. +pub(crate) struct PyData(pub(crate) T); + +impl<'a, 'py, T: FromValue> FromPyObject<'a, 'py> for PyData { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let PyValue(value) = value.extract()?; + T::from_value(&value) + .map(Self) + .map_err(|error| PyValueError::new_err(error.to_string())) + } +} + +impl<'py, T: IntoValue> IntoPyObject<'py> for PyData { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + PyValue(self.0.into_value()).into_pyobject(py) + } +} + +fn unwrap_all(values: Vec>) -> Vec { + values.into_iter().map(|value| value.0).collect() +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs index 758d114fa2e..409e949b11b 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs @@ -1,17 +1,23 @@ use std::{sync::Arc, time::Duration}; -use litellm_redis_store::rate_limits::{ - AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Error, Gauge, - GuardedIncrement, GuardedReply, Increment, Local, Pending, RateLimits, Scope, WindowCount, +use litellm_redis_store::{ + Pending, + rate_limits::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Gauge, + GuardedIncrement, GuardedReply, Increment, Local, RateLimits, Scope, WindowCount, + }, }; use pyo3::{ - exceptions::{PyRuntimeError, PyTypeError, PyValueError}, + exceptions::PyTypeError, prelude::*, types::{PyDict, PyString}, }; -use super::{NativePending, PyNumber, redis_async, request::NativeRequestRedis, store_error, ttl}; -use crate::execution::run_async_value; +use super::{ + NativePending, PyData, PyNumber, redis_async, remote, + request::{NativeRequestRedis, NativeSend}, + runtime_error, ttl, unwrap_all, +}; use crate::redis::client::NativeRedisClient; /// The v3 rate limiter's keyspace. Memory reads and writes are synchronous like `InMemoryCache`; @@ -19,91 +25,6 @@ use crate::redis::client::NativeRedisClient; #[pyclass(frozen, name = "NativeRateLimits")] pub(crate) struct NativeRateLimits(Arc); -/// `(descriptor_key, descriptor_value)`. -struct PyScope(Scope); - -impl<'a, 'py> FromPyObject<'a, 'py> for PyScope { - type Error = PyErr; - - fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { - let (descriptor, value): (String, String) = value.extract()?; - Ok(Self(Scope { descriptor, value })) - } -} - -/// `(descriptor_key, descriptor_value, kind)`, the kind spelled as its key suffix. -struct PyCounter(Counter); - -impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter { - type Error = PyErr; - - fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { - let (descriptor, value, kind): (String, String, String) = value.extract()?; - Ok(Self(Counter { - scope: Scope { descriptor, value }, - kind: kind.parse().map_err(|_| unknown_suffix(&kind))?, - })) - } -} - -fn unknown_suffix(suffix: &str) -> PyErr { - PyValueError::new_err(format!("unknown rate-limit key suffix {suffix:?}")) -} - -/// `(key_suffix, descriptor_key, descriptor_value)`, read by `Entry::new`. -struct PyEntry(Entry); - -impl<'a, 'py> FromPyObject<'a, 'py> for PyEntry { - type Error = PyErr; - - fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { - let (suffix, descriptor, value): (String, String, String) = value.extract()?; - Entry::new(&suffix, Scope { descriptor, value }) - .map(Self) - .map_err(|_| unknown_suffix(&suffix)) - } -} - -/// `(counter, amount, ttl_seconds | None)`. -struct PyIncrement(Increment); - -impl<'a, 'py> FromPyObject<'a, 'py> for PyIncrement { - type Error = PyErr; - - fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { - let (counter, amount, ttl): (PyCounter, PyNumber, Option) = value.extract()?; - Ok(Self(Increment { - counter: counter.0, - amount: amount.0, - ttl: ttl.map(Duration::from_secs), - })) - } -} - -/// `(counter, limit, increment, ttl_seconds, window_size_seconds)`. -struct PyAtomicCounter(AtomicCounter); - -impl<'a, 'py> FromPyObject<'a, 'py> for PyAtomicCounter { - type Error = PyErr; - - fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { - let (counter, limit, increment, ttl_seconds, window_size_seconds): ( - PyCounter, - i64, - i64, - i64, - i64, - ) = value.extract()?; - Ok(Self(AtomicCounter { - counter: counter.0, - limit, - increment, - ttl_seconds, - window_size_seconds, - })) - } -} - /// An `int` or `float` counter, a `str` window start, or a `{slot_id: acquired_at}` registry. struct PyLocal(Local); @@ -201,48 +122,13 @@ fn acquire(reply: AcquireReply) -> (OverLimit, Vec) { } } -/// One script call declared on a batch: the request's pre-call batch, or a batch of its own when -/// no request was passed. Resolving it flushes that batch if it has not gone out. -fn pending( - call_type: &'static str, - pending: Result, Error>, - reply: fn(T) -> R, -) -> PyResult +/// One script call declared on the batch the caller's `NativeSend` picked; `None` without Redis. +fn script(call_type: &'static str, declared: Option>, reply: fn(T) -> R) -> Option where T: Send + 'static, R: for<'a> IntoPyObject<'a> + Send + 'static, { - let pending = pending.map_err(rate_limit_error)?; - Ok(NativePending::new( - call_type, - rate_limit_error, - async move { pending.resolve().await.map(reply) }, - )) -} - -fn rate_limit_error(error: Error) -> PyErr { - PyRuntimeError::new_err(litellm_error::Report::new(&error).to_string()) -} - -fn scopes(scopes: Vec) -> Vec { - scopes.into_iter().map(|scope| scope.0).collect() -} - -fn counters(counters: Vec) -> Vec { - counters.into_iter().map(|counter| counter.0).collect() -} - -fn increments(increments: Vec) -> Vec { - increments - .into_iter() - .map(|increment| increment.0) - .collect() -} - -fn request_of( - request: Option>, -) -> Option { - request.map(|request| request.0.clone()) + Some(super::pending(call_type, runtime_error, declared?, reply)) } #[pymethods] @@ -254,51 +140,47 @@ impl NativeRateLimits { backfill_ttl_seconds: Option, redis_ttl_seconds: u64, ) -> Self { - let backfill_ttl = ttl(backfill_ttl_seconds); Self(Arc::new(RateLimits::new( - client.map(|client| client.0.clone()), - backfill_ttl, - Duration::from_secs(redis_ttl_seconds), + client.map(|client| remote(&client, redis_ttl_seconds)), + ttl(backfill_ttl_seconds), ))) } + #[getter] fn has_redis(&self) -> bool { self.0.has_redis() } - fn slot_groups(&self, scopes_: Vec, force: bool) -> Vec> { - self.0.slot_groups(&scopes(scopes_), force) + fn slot_groups(&self, scopes: Vec>, force: bool) -> Vec> { + self.0.slot_groups(&unwrap_all(scopes), force) } - fn local(&self, entries: Vec) -> PyResult>> { - entries - .iter() - .map(|entry| { - self.0 - .local(&entry.0) - .map(|local| local.map(PyLocal)) - .map_err(store_error) - }) - .collect() + fn get_local(&self, entry: PyData) -> PyResult> { + Ok(self.0.get_local(&entry.0).map_err(runtime_error)?.map(PyLocal)) } #[pyo3(signature = (entry, value, ttl_seconds))] - fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option) -> PyResult<()> { + fn set_local( + &self, + entry: PyData, + value: PyLocal, + ttl_seconds: Option, + ) -> PyResult<()> { self.0 .set_local(&entry.0, value.0, ttl(ttl_seconds)) - .map_err(store_error) + .map_err(runtime_error) } #[pyo3(signature = (entry, *, backfill_ttl_seconds = None))] fn get<'py>( &self, py: Python<'py>, - entry: PyEntry, + entry: PyData, backfill_ttl_seconds: Option, ) -> PyResult> { let store = self.0.clone(); let backfill_ttl = ttl(backfill_ttl_seconds); - redis_async(py, "async_get_cache", store_error, async move { + redis_async(py, "async_get_cache", runtime_error, async move { store .get(&entry.0, backfill_ttl) .await @@ -306,99 +188,79 @@ impl NativeRateLimits { }) } - #[pyo3(signature = (counters_, now, window_size, *, request = None))] fn sliding_window( &self, - counters_: Vec, + counters: Vec>, now: i64, window_size: i64, - request: Option>, - ) -> PyResult { - let request = request_of(request); - pending( + send: &NativeSend, + ) -> Option { + script( "batch_rate_limiter_script", self.0 - .sliding_window(request.as_ref(), &counters(counters_), now, window_size), + .sliding_window(send.send(), &unwrap_all(counters), now, window_size), window_counts, ) } - #[pyo3(signature = (counters_, *, request = None))] fn read_counters( &self, - counters_: Vec, - request: Option>, - ) -> PyResult { - let request = request_of(request); - pending( + counters: Vec>, + send: &NativeSend, + ) -> Option { + script( "batch_counter_read_script", - self.0.read_counters(request.as_ref(), &counters(counters_)), + self.0.read_counters(send.send(), &unwrap_all(counters)), std::convert::identity, ) } - #[pyo3(signature = (counters_, *, request = None))] fn check_and_increment( &self, - counters_: Vec, - request: Option>, - ) -> PyResult { - let request = request_of(request); - let counters_ = counters_ - .into_iter() - .map(|counter| counter.0) - .collect::>(); - pending( + counters: Vec>, + send: &NativeSend, + ) -> Option { + script( "check_and_increment_by_n_script", - self.0.check_and_increment(request.as_ref(), &counters_), + self.0 + .check_and_increment(send.send(), &unwrap_all(counters)), atomic, ) } - #[pyo3(signature = (counter, expected_window_start, amount, ttl_seconds, *, request = None))] fn guarded_increment( &self, - counter: PyCounter, + counter: PyData, expected_window_start: String, amount: PyNumber, ttl_seconds: Option, - request: Option>, - ) -> PyResult { - let request = request_of(request); + send: &NativeSend, + ) -> Option { let increment = GuardedIncrement { counter: counter.0, expected_window_start, amount: amount.0, ttl: ttl_seconds.map(Duration::from_secs), }; - pending( + script( "window_guarded_token_increment_script", - self.0.guarded_increment(request.as_ref(), &increment), + self.0.guarded_increment(send.send(), &increment), guarded, ) } - #[pyo3(signature = (gauges, slot_ttl_seconds, slot_id, *, request = None))] fn parallel_acquire( &self, - gauges: Vec<(PyScope, i64)>, + gauges: Vec>, slot_ttl_seconds: u64, slot_id: &str, - request: Option>, - ) -> PyResult { - let request = request_of(request); - let gauges = gauges - .into_iter() - .map(|(scope, limit)| Gauge { - scope: scope.0, - limit, - }) - .collect::>(); - pending( + send: &NativeSend, + ) -> Option { + script( "parallel_acquire_script", self.0.parallel_acquire( - request.as_ref(), - &gauges, + send.send(), + &unwrap_all(gauges), Duration::from_secs(slot_ttl_seconds), slot_id, ), @@ -406,96 +268,76 @@ impl NativeRateLimits { ) } - #[pyo3(signature = (gauges, slot_id, *, request = None))] fn parallel_release( &self, - gauges: Vec, + gauges: Vec>, slot_id: &str, - request: Option>, - ) -> PyResult { - let request = request_of(request); - pending( + send: &NativeSend, + ) -> Option { + script( "parallel_release_script", self.0 - .parallel_release(request.as_ref(), &scopes(gauges), slot_id), + .parallel_release(send.send(), &unwrap_all(gauges), slot_id), std::convert::identity, ) } - #[pyo3(signature = (gauges, slot_ttl_seconds, *, request = None))] fn parallel_count( &self, - gauges: Vec, + gauges: Vec>, slot_ttl_seconds: u64, - request: Option>, - ) -> PyResult { - let request = request_of(request); - pending( + send: &NativeSend, + ) -> Option { + script( "parallel_count_script", self.0.parallel_count( - request.as_ref(), - &scopes(gauges), + send.send(), + &unwrap_all(gauges), Duration::from_secs(slot_ttl_seconds), ), std::convert::identity, ) } - #[pyo3(signature = (increments_, *, request = None))] fn token_increment( &self, - increments_: Vec, - request: Option>, - ) -> PyResult { - let request = request_of(request); - pending( + increments: Vec>, + send: &NativeSend, + ) -> Option { + script( "token_increment_script", - self.0 - .token_increment(request.as_ref(), &increments(increments_)), + self.0.token_increment(send.send(), &unwrap_all(increments)), std::convert::identity, ) } - fn token_increment_post_call<'py>( + /// `False` without Redis. + fn token_increment_post_call( &self, - py: Python<'py>, - request: PyRef<'_, NativeRequestRedis>, - increments_: Vec, - ) -> PyResult> { - let store = self.0.clone(); - let request = request.0.clone(); - let increments_ = increments(increments_); - run_async_value(py, async move { - store - .token_increment_post_call(&request, increments_) - .map_err(rate_limit_error) - }) + request: &NativeRequestRedis, + increments: Vec>, + ) -> bool { + self.0.token_increment_post_call(&request.0, unwrap_all(increments)) } - fn parallel_release_post_call<'py>( + /// `False` without Redis. + fn parallel_release_post_call( &self, - py: Python<'py>, - request: PyRef<'_, NativeRequestRedis>, - gauges: Vec, - slot_id: String, - ) -> PyResult> { - let store = self.0.clone(); - let request = request.0.clone(); - let gauges = scopes(gauges); - run_async_value(py, async move { - store - .parallel_release_post_call(&request, &gauges, &slot_id) - .map_err(rate_limit_error) - }) + request: &NativeRequestRedis, + gauges: Vec>, + slot_id: &str, + ) -> bool { + self.0 + .parallel_release_post_call(&request.0, &unwrap_all(gauges), slot_id) } fn increment<'py>( &self, py: Python<'py>, - increment: PyIncrement, + increment: PyData, ) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_increment", store_error, async move { + redis_async(py, "async_increment", runtime_error, async move { store.increment(&increment.0).await.map(PyNumber) }) } @@ -503,28 +345,22 @@ impl NativeRateLimits { fn increment_many<'py>( &self, py: Python<'py>, - increments_: Vec, + increments: Vec>, ) -> PyResult> { let store = self.0.clone(); - let increments_ = increments(increments_); - redis_async(py, "async_increment_pipeline", store_error, async move { - store.increment_many(&increments_).await + let increments = unwrap_all(increments); + redis_async(py, "async_increment_pipeline", runtime_error, async move { + store.increment_many(&increments).await }) } - fn increment_many_post_call<'py>( + fn increment_many_post_call( &self, - py: Python<'py>, - request: PyRef<'_, NativeRequestRedis>, - increments_: Vec, - ) -> PyResult> { - let store = self.0.clone(); - let request = request.0.clone(); - let increments_ = increments(increments_); - run_async_value(py, async move { - store - .increment_many_post_call(&request, &increments_) - .map_err(store_error) - }) + request: &NativeRequestRedis, + increments: Vec>, + ) -> PyResult<()> { + self.0 + .increment_many_post_call(&request.0, &unwrap_all(increments)) + .map_err(runtime_error) } } diff --git a/litellm-rust/crates/python-bridge/src/redis/store/request.rs b/litellm-rust/crates/python-bridge/src/redis/store/request.rs index 1ee80bb0467..6860cf8dcbb 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/request.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/request.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_redis_store::RequestRedis; +use litellm_redis_store::{RequestRedis, Send}; use pyo3::prelude::*; use super::call; @@ -39,3 +39,22 @@ impl NativeRequestRedis { }) } } + +/// The Python face of `Send`: which batch a store op rides. Each variant carries exactly what it +/// needs, so there is no phase string and no phase without a request. +#[pyclass(frozen, name = "NativeSend")] +pub(crate) enum NativeSend { + Now(), + PreCall(Py), + PostCall(Py), +} + +impl NativeSend { + pub(crate) fn send(&self) -> Send<'_> { + match self { + Self::Now() => Send::Now, + Self::PreCall(request) => Send::PreCall(&request.get().0), + Self::PostCall(request) => Send::PostCall(&request.get().0), + } + } +} 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 044c64dcf98..77b04bc7a26 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 @@ -1,15 +1,17 @@ use std::sync::Arc; -use litellm_redis_store::router_cooldowns::{Cooldown, RouterCooldowns}; +use litellm_redis_store::{ + Send, + 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, value::PyValue, + NativePending, PyData, PyNumber, pending, policy::router_tier_policies, redis_sync, remote, + request::NativeSend, runtime_error, ttl, }; use crate::redis::client::NativeRedisClient; @@ -18,30 +20,6 @@ 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 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 = PyAny; - type Output = Bound<'py, PyAny>; - type Error = PyErr; - - fn into_pyobject(self, py: Python<'py>) -> PyResult { - PyValue(self.0.to_value()).into_pyobject(py) - } -} - /// `[(model_id, cooldown), ...]` the way `RedisCache` hands the router its batch read. struct PyActive(Vec<(String, Cooldown)>); @@ -68,7 +46,7 @@ fn active_to_py( py, [ model_id.into_pyobject(py)?.into_any(), - PyCooldown(cooldown).into_pyobject(py)?.into_any(), + PyData(cooldown).into_pyobject(py)?, ], ) }) @@ -78,22 +56,29 @@ fn active_to_py( #[pymethods] impl NativeRouterCooldowns { #[new] - fn new(py: Python<'_>, client: Option>) -> PyResult { - let client = client.map(|client| client.0.clone()); + #[pyo3(signature = (client, redis_ttl_seconds = 60))] + fn new(py: Python<'_>, client: Option>, redis_ttl_seconds: u64) -> PyResult { + let remote = client.map(|client| remote(&client, redis_ttl_seconds)); let policy = router_tier_policies(py)?.cooldowns; - Ok(Self(Arc::new(RouterCooldowns::new(client, policy)))) + Ok(Self(Arc::new(RouterCooldowns::new(remote, policy)))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + #[pyo3(signature = (client, redis_ttl_seconds = 60))] + fn attach(&self, client: PyRef<'_, NativeRedisClient>, redis_ttl_seconds: u64) -> bool { + self.0.attach(remote(&client, redis_ttl_seconds)) } fn put( &self, py: Python<'_>, model_id: String, - cooldown: PyCooldown, + cooldown: PyData, ttl_seconds: f64, ) -> PyResult<()> { let store = self.0.clone(); let ttl = ttl(Some(ttl_seconds)).unwrap_or_default(); - redis_sync(py, "set_cache", store_error, async move { + redis_sync(py, "set_cache", runtime_error, async move { store.put(&model_id, &cooldown.0, ttl).await }) } @@ -107,52 +92,33 @@ impl NativeRouterCooldowns { ) -> PyResult { let store = self.0.clone(); let ttl = ttl(Some(ttl_seconds)).unwrap_or_default(); - redis_sync(py, "increment_cache", store_error, async move { + redis_sync(py, "increment_cache", runtime_error, async move { store .increment_allowed_fails(&model_id, suffix.as_deref(), ttl) .await }) } - fn active<'py>( + /// The deployments still cooling down, read on the batch `send` picks. + fn active(&self, model_ids: Vec, send: &NativeSend) -> PyResult { + let active = self.0.active(send.send(), &model_ids).map_err(runtime_error)?; + Ok(pending("async_batch_get_cache", runtime_error, active, PyActive)) + } + + /// `active` on its own with the GIL released, for the router's sync completion path. + fn active_sync<'py>( &self, py: Python<'py>, model_ids: Vec, ) -> PyResult>> { - let store = self.0.clone(); - let active = redis_sync(py, "batch_get_cache", store_error, async move { - store.active(&model_ids).await - })?; + let active = self.0.active(Send::Now, &model_ids).map_err(runtime_error)?; + let active = redis_sync(py, "batch_get_cache", runtime_error, active)?; active_to_py(py, active) } - fn aactive<'py>(&self, py: Python<'py>, model_ids: Vec) -> PyResult> { - let store = self.0.clone(); - redis_async(py, "async_batch_get_cache", store_error, async move { - store.active(&model_ids).await.map(PyActive) - }) - } - - fn declare_active( - &self, - request: &NativeRequestRedis, - model_ids: Vec, - ) -> PyResult { - let declared = self - .0 - .declare_active(&request.0, &model_ids) - .map_err(store_error)?; - let store = self.0.clone(); - Ok(NativePending::new( - "async_batch_get_cache", - store_error, - async move { store.resolve_active(declared).await.map(PyActive) }, - )) - } - fn minimum(&self, py: Python<'_>, model_ids: Vec) -> PyResult> { let store = self.0.clone(); - redis_sync(py, "batch_get_cache", store_error, async move { + redis_sync(py, "batch_get_cache", runtime_error, async move { store .minimum(&model_ids) .await @@ -160,18 +126,18 @@ impl NativeRouterCooldowns { }) } - fn evict_local(&self, model_id: &str) -> PyResult<()> { - self.0.evict_local(model_id).map_err(store_error) + fn delete_local(&self, model_id: &str) -> PyResult<()> { + self.0.delete_local(model_id).map_err(runtime_error) } fn flush_local_counters(&self) -> PyResult<()> { - self.0.flush_local_counters().map_err(store_error) + self.0.flush_local_counters().map_err(runtime_error) } fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { let snapshot = PyDict::new(py); - for (key, cooldown) in self.0.snapshot().map_err(store_error)? { - snapshot.set_item(key, PyCooldown(cooldown))?; + for (key, cooldown) in self.0.snapshot().map_err(runtime_error)? { + snapshot.set_item(key, PyData(cooldown))?; } Ok(snapshot) } 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 26597a16731..30aea85c676 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,54 +1,40 @@ use std::{sync::Arc, time::Duration}; use litellm_redis_store::router_usage::{Counter, RouterUsage}; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; -use super::{PyNumber, policy::router_tier_policies, redis_async, store_error}; +use super::{PyData, PyNumber, policy::router_tier_policies, redis_async, remote, runtime_error, unwrap_all}; use crate::redis::client::NativeRedisClient; /// One router's global TPM/RPM counters. Every caller on `main` is async. #[pyclass(frozen, name = "NativeRouterUsage")] pub(crate) struct NativeRouterUsage(Arc); -/// `(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 { - type Error = PyErr; - - fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { - let (deployment_id, model, minute, kind): (String, Option, String, String) = - value.extract()?; - let kind = kind - .parse() - .map_err(|_| PyValueError::new_err(format!("unknown usage counter kind {kind:?}")))?; - Ok(Self(Counter { - deployment_id, - model, - minute, - kind, - })) - } -} - #[pymethods] impl NativeRouterUsage { #[new] - fn new(py: Python<'_>, client: Option>) -> PyResult { - let client = client.map(|client| client.0.clone()); + #[pyo3(signature = (client, redis_ttl_seconds = 60))] + fn new(py: Python<'_>, client: Option>, redis_ttl_seconds: u64) -> PyResult { + let remote = client.map(|client| remote(&client, redis_ttl_seconds)); let policy = router_tier_policies(py)?.usage; - Ok(Self(Arc::new(RouterUsage::new(client, policy)))) + Ok(Self(Arc::new(RouterUsage::new(remote, policy)))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + #[pyo3(signature = (client, redis_ttl_seconds = 60))] + fn attach(&self, client: PyRef<'_, NativeRedisClient>, redis_ttl_seconds: u64) -> bool { + self.0.attach(remote(&client, redis_ttl_seconds)) } fn increment<'py>( &self, py: Python<'py>, - counter: PyCounter, + counter: PyData, value: PyNumber, ttl_seconds: u64, ) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_increment", store_error, async move { + redis_async(py, "async_increment", runtime_error, async move { store .increment(&counter.0, value.0, Duration::from_secs(ttl_seconds)) .await @@ -67,7 +53,7 @@ impl NativeRouterUsage { requests: i64, ) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_increment_pipeline", store_error, async move { + redis_async(py, "async_increment_pipeline", runtime_error, async move { store .increment_deployment( &deployment_id, @@ -83,14 +69,11 @@ impl NativeRouterUsage { fn get_many<'py>( &self, py: Python<'py>, - counters: Vec, + counters: Vec>, ) -> PyResult> { let store = self.0.clone(); - let counters = counters - .into_iter() - .map(|counter| counter.0) - .collect::>(); - redis_async(py, "async_batch_get_cache", store_error, async move { + let counters = unwrap_all(counters); + redis_async(py, "async_batch_get_cache", runtime_error, async move { let read = store.get_many(&counters).await?; Ok(read .into_iter() @@ -100,12 +83,12 @@ impl NativeRouterUsage { } fn flush_local_counters(&self) -> PyResult<()> { - self.0.flush_local_counters().map_err(store_error) + self.0.flush_local_counters().map_err(runtime_error) } fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { let snapshot = PyDict::new(py); - for (key, count) in self.0.snapshot().map_err(store_error)? { + for (key, count) in self.0.snapshot().map_err(runtime_error)? { snapshot.set_item(key, PyNumber(count))?; } Ok(snapshot) diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs index 90a3b8dc4a3..3c57aa82212 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -1,53 +1,27 @@ -use std::{convert::Infallible, sync::Arc, time::Duration}; +use std::sync::Arc; -use litellm_redis_store::spend_counters::{Error, LostSeed, SpendCounter, SpendCounters}; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use litellm_redis_store::{ + Number, + spend_counters::{LostSeed, SpendCounter, SpendCounters}, +}; +use pyo3::{prelude::*, types::PyDict}; use super::{ - NativePending, PyNumber, kind_error, redis_async, request::NativeRequestRedis, store_error, ttl, + NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote, + request::{NativeRequestRedis, NativeSend}, + ttl, unwrap_all, }; use crate::redis::client::NativeRedisClient; -fn spend_error(error: Error) -> PyErr { - kind_error(error.kind(), &error) -} - -fn infallible(error: Infallible) -> PyErr { - match error {} -} - -/// A `spend:*` key, refused when it names anything outside the spend counter keyspace. -pub(crate) struct PySpendCounter(SpendCounter); - -impl<'a, 'py> FromPyObject<'a, 'py> for PySpendCounter { - type Error = PyErr; - - fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { - let key: String = value.extract()?; - SpendCounter::parse(&key) - .map(Self) - .map_err(|error| PyValueError::new_err(error.to_string())) - } -} - -fn counters(counters: Vec) -> Vec { - counters.into_iter().map(|counter| counter.0).collect() -} - -fn increments( - increments: Vec<(PySpendCounter, PyNumber)>, -) -> Vec<(SpendCounter, litellm_redis_store::Number)> { +fn increments(increments: Vec<(PyData, PyNumber)>) -> Vec<(SpendCounter, Number)> { increments .into_iter() .map(|(counter, amount)| (counter.0, amount.0)) .collect() } -fn numbers(values: Vec>) -> Vec> { - values - .into_iter() - .map(|value| value.map(PyNumber)) - .collect() +fn numbers(values: Option>>) -> Option>> { + values.map(|values| values.into_iter().map(|value| value.map(PyNumber)).collect()) } /// The proxy's spend counters, Redis first with this worker's memory copy behind it. @@ -60,55 +34,41 @@ impl NativeSpendCounters { #[pyo3(signature = (client, redis_ttl_seconds = 60))] fn new(client: Option>, redis_ttl_seconds: u64) -> Self { Self(Arc::new(SpendCounters::new( - client.map(|client| client.0.clone()), - Duration::from_secs(redis_ttl_seconds), + client.map(|client| remote(&client, redis_ttl_seconds)), ))) } + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + #[pyo3(signature = (client, redis_ttl_seconds = 60))] + fn attach(&self, client: PyRef<'_, NativeRedisClient>, redis_ttl_seconds: u64) -> bool { + self.0.attach(remote(&client, redis_ttl_seconds)) + } + /// `("redis" | "memory", value)`: where the answer came from, a Redis miss included. - fn read<'py>(&self, py: Python<'py>, counter: PySpendCounter) -> PyResult> { + fn read<'py>(&self, py: Python<'py>, counter: PyData) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_get_cache", store_error, async move { + redis_async(py, "async_get_cache", kind_error, async move { let reading = store.read(&counter.0).await?; let source: &'static str = reading.into(); Ok((source, reading.value().map(PyNumber))) }) } - fn read_many<'py>( - &self, - py: Python<'py>, - counters: Vec, - ) -> PyResult> { - let store = self.0.clone(); - let counters = self::counters(counters); - redis_async(py, "async_batch_get_cache", infallible, async move { - Ok(store.read_many(&counters).await.map(numbers)) - }) - } - - fn declare_read_many( - &self, - request: &NativeRequestRedis, - counters: Vec, - ) -> Option { - self.0 - .declare_read_many_pre_call(&request.0, &self::counters(counters)) - .map(|declared| { - NativePending::new("async_batch_get_cache", infallible, async move { - Ok(declared.resolve().await.map(numbers)) - }) - }) + /// One `MGET` on the batch `send` picks; `None` without Redis. It resolves to `None` when the + /// read failed, so each caller reads on its own. + fn read_many(&self, counters: Vec>, send: &NativeSend) -> Option { + let read = self.0.read_many(send.send(), &unwrap_all(counters))?; + Some(pending("async_batch_get_cache", kind_error, read, numbers)) } fn increment<'py>( &self, py: Python<'py>, - counter: PySpendCounter, + counter: PyData, amount: PyNumber, ) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_increment", spend_error, async move { + redis_async(py, "async_increment", kind_error, async move { store.increment(&counter.0, amount.0).await.map(PyNumber) }) } @@ -116,11 +76,11 @@ impl NativeSpendCounters { fn increment_many<'py>( &self, py: Python<'py>, - increments: Vec<(PySpendCounter, PyNumber)>, + increments: Vec<(PyData, PyNumber)>, ) -> PyResult> { let store = self.0.clone(); let increments = self::increments(increments); - redis_async(py, "async_increment_pipeline", spend_error, async move { + redis_async(py, "async_increment_pipeline", kind_error, async move { let totals = store.increment_many(&increments).await?; Ok(totals.into_iter().map(PyNumber).collect::>()) }) @@ -130,7 +90,7 @@ impl NativeSpendCounters { fn increment_many_post_call( &self, request: &NativeRequestRedis, - increments: Vec<(PySpendCounter, PyNumber)>, + increments: Vec<(PyData, PyNumber)>, ) -> bool { self.0 .increment_many_post_call(&request.0, &self::increments(increments)) @@ -139,11 +99,11 @@ impl NativeSpendCounters { fn set_max<'py>( &self, py: Python<'py>, - counter: PySpendCounter, + counter: PyData, spend: PyNumber, ) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_set_max", spend_error, async move { + redis_async(py, "async_set_max", kind_error, async move { store .set_max(&counter.0, spend.0) .await @@ -155,7 +115,7 @@ impl NativeSpendCounters { fn seed<'py>( &self, py: Python<'py>, - counter: PySpendCounter, + counter: PyData, spend: PyNumber, increment_if_lost: bool, ) -> PyResult> { @@ -165,7 +125,7 @@ impl NativeSpendCounters { } else { LostSeed::Answer }; - redis_async(py, "async_set_cache", spend_error, async move { + redis_async(py, "async_set_cache", kind_error, async move { store.seed(&counter.0, spend.0, lost).await.map(PyNumber) }) } @@ -174,12 +134,12 @@ impl NativeSpendCounters { fn set<'py>( &self, py: Python<'py>, - counter: PySpendCounter, + counter: PyData, spend: PyNumber, ttl_seconds: Option, ) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_set_cache", spend_error, async move { + redis_async(py, "async_set_cache", kind_error, async move { store.set(&counter.0, spend.0, ttl(ttl_seconds)).await }) } @@ -187,11 +147,11 @@ impl NativeSpendCounters { fn delete<'py>( &self, py: Python<'py>, - counters: Vec, + counters: Vec>, ) -> PyResult> { let store = self.0.clone(); - let counters = self::counters(counters); - redis_async(py, "async_delete_cache", spend_error, async move { + let counters = unwrap_all(counters); + redis_async(py, "async_delete_cache", kind_error, async move { store.delete(&counters).await }) } @@ -199,52 +159,52 @@ impl NativeSpendCounters { fn refresh_ttl<'py>( &self, py: Python<'py>, - counter: PySpendCounter, + counter: PyData, ) -> PyResult> { let store = self.0.clone(); - redis_async(py, "async_refresh_ttl", spend_error, async move { + redis_async(py, "async_refresh_ttl", kind_error, async move { store.refresh_ttl(&counter.0).await }) } - fn get_local(&self, counter: PySpendCounter) -> PyResult> { + fn get_local(&self, counter: PyData) -> PyResult> { Ok(self .0 .get_local(&counter.0) - .map_err(store_error)? + .map_err(kind_error)? .map(PyNumber)) } #[pyo3(signature = (counter, spend, ttl_seconds = None))] fn set_local( &self, - counter: PySpendCounter, + counter: PyData, spend: PyNumber, ttl_seconds: Option, ) -> PyResult<()> { self.0 .set_local(&counter.0, spend.0, ttl(ttl_seconds)) - .map_err(store_error) + .map_err(kind_error) } - fn delete_local(&self, counter: PySpendCounter) -> PyResult<()> { - self.0.delete_local(&counter.0).map_err(store_error) + fn delete_local(&self, counter: PyData) -> PyResult<()> { + self.0.delete_local(&counter.0).map_err(kind_error) } - fn increment_local(&self, counter: PySpendCounter, amount: PyNumber) -> PyResult { + fn increment_local(&self, counter: PyData, amount: PyNumber) -> PyResult { self.0 .increment_local(&counter.0, amount.0) .map(PyNumber) - .map_err(store_error) + .map_err(kind_error) } fn flush_local(&self) -> PyResult<()> { - self.0.flush_local().map_err(store_error) + self.0.flush_local().map_err(kind_error) } fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { let snapshot = PyDict::new(py); - for (key, spend) in self.0.snapshot().map_err(store_error)? { + for (key, spend) in self.0.snapshot().map_err(kind_error)? { snapshot.set_item(key, PyNumber(spend))?; } Ok(snapshot) diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index 2af2dcf20b8..6cbfd8a827e 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -23,21 +23,23 @@ - The rule never changes, because Python and Rust pods share keys in the same Redis ## Commands -- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- A command is a value, one constructor in `commands.rs` returning `Command`: its alone and pipelined forms and its reply decode. `RedisClient::send`, `Pipeline::add` and `RedisBatch::add` are the only entry points, and each applies its namespace to the keys once. A new command is one constructor, never a method on an executor. `ping` and `subscribe` stay methods. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- A store attaches its value codec with `Command::map`, never a keyspace codec inside this crate - On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test - Every command is the one `RedisCache` sends, argument for argument, so `MONITOR` on a mixed fleet shows one shape. A difference from main's Python is a defect even when the result is the same - TTLs are `Duration` and are sent in whole seconds (`SET ... EX`, `EXPIRE`, `TTL`), like redis-py's `int(ttl)`. `set_nx` sends `EX` before `NX`, the order redis-py builds - Numbers go out as `Amount`, written the way redis-py encodes them: an `int` with `str()` and a `float` with `repr()` from `python-compat` -- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. In a pipeline `incr_by_float` is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when a TTL is given +- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` alone is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. Pipelined it is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when an expiry is given, `KeepExisting` lowered to `Refresh`. `Plan` keeps only the cases whose alone and pipelined forms differ ## Scripts and pipelines - A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production - `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline a script is sent as `EVALSHA`, like `RedisBatch`. A `NOSCRIPT` reply means that op ran nothing, so it alone runs again after the pipeline through `invoke_async` (`EVALSHA`, `SCRIPT LOAD`, `EVALSHA`), like `RedisBatch`'s `_settle_alone`. It then lands after the ops around it, as it does in Python -- `Entry` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice +- `Plan::Script` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice +- `RedisBatch::pipelined` is one `Pipeline` whose ops answer as typed `Declared`s. A store that sends its own pipeline uses it rather than decoding `run`'s raw `Value`s - A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip - redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key and runs the slots concurrently, returning results in input order. Within a slot, runs of idempotent ops (GET, SET, DEL, EXPIRE) share a request and every script op gets its own, sent in order. redis-rs retries a cluster request whole, so a script op batched with others would be applied twice when only a neighbour failed - Pipeline ops take one key each. A multi-key op would span slots inside one group. An op can send more than one command for that key (`INCRBYFLOAT` then `EXPIRE`), and settles from the first reply, failing if any of its replies failed -- Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once +- Command constructors live in `commands.rs` and every executor sends them, so a command is written once ## Pub/sub - Classic `PUBLISH` / `SUBSCRIBE`, never sharded, since redis-py publishes and subscribes through one plain client and cluster broadcasts classic messages to every node. Payloads are raw bytes, and the codec belongs to the store that owns the channel @@ -50,6 +52,7 @@ ## Batches - `RedisBatch` is litellm's `RedisBatch` (`litellm/caching/redis_batch.py`). It keeps each declared op as data and, on a standalone or Sentinel Redis, sends them later as one `Pipeline`, so it inherits every pipeline rule above. On a cluster it never pipelines and runs each op on its own, concurrently, the way `_settle_alone` calls the `RedisCache` method behind it. An increment there keeps a TTL the key already has, as `async_increment` does, where the pipelined one refreshes it like `_Increment` - `mget` is one `MGET` of the keys in order with repeats read once, like `_MGet`. Use it instead of one `get` per key whenever Python reads several keys in one call It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local +- `RedisBatch::alone` runs every op on its own like the cluster path, so a fresh alone batch is "send it now" and `Declared` is the one reply handle on every route - Each op returns a `Declared` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once - An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis - `settled` waits for whichever flush carries the op and never sends one itself, for follow-ups that must not force a flush. `RedisClient::same_client` tells clones of one client apart from a second connection, so a host can keep one batch per client diff --git a/litellm-rust/crates/redis-client/src/batch.rs b/litellm-rust/crates/redis-client/src/batch.rs index ab008f95a34..84fc1643424 100644 --- a/litellm-rust/crates/redis-client/src/batch.rs +++ b/litellm-rust/crates/redis-client/src/batch.rs @@ -1,160 +1,78 @@ use std::{ io, sync::{Arc, Mutex, MutexGuard}, - time::Duration, }; use futures_util::future::join_all; -use redis::{FromRedisValue, RedisError, ToRedisArgs, ToSingleRedisArg, Value}; +use redis::{RedisError, Value}; use tokio::sync::oneshot::{self, error::TryRecvError}; -use crate::{Amount, Error, Expiry, Pipeline, RedisClient, Script}; +use crate::{Command, Error, Pipeline, RedisClient, commands::Op}; /// Ops declared by many callers and sent together, like litellm's `RedisBatch`: one pipeline on a -/// standalone or Sentinel Redis, and each op on its own on a cluster. +/// standalone or Sentinel Redis, and each op on its own on a cluster or in an alone batch. #[derive(Clone)] pub struct RedisBatch(Arc); +#[derive(Clone, Copy)] +enum Mode { + /// litellm's `RedisBatch`: one pipeline, or each op on its own on a cluster. + Batch, + /// Each op on its own, the way the `RedisCache` method behind it sends it. + Alone, + /// One redis-py pipeline on every topology, split by slot on a cluster. + Pipeline, +} + struct Shared { client: RedisClient, + mode: Mode, pending: Mutex, flushing: tokio::sync::Mutex<()>, } #[derive(Default)] struct Pending { - ops: Vec, + ops: Vec, replies: Vec>>, } -/// One declared op, kept as data so a flush can pipeline it or run it alone. -enum BatchOp { - Get(String), - MGet(Vec), - Set { - key: String, - value: Vec, - ttl: Duration, - }, - Del(String), - Expire { - key: String, - ttl: Duration, - }, - IncrByFloat { - key: String, - amount: Amount, - ttl: Option, - }, - Eval { - script: Script, - keys: Vec, - args: Vec>, - }, -} +type Decode = Box Result + Send>; #[must_use] pub struct Declared { batch: RedisBatch, reply: oneshot::Receiver>, - decode: fn(Value) -> Result, + decode: Decode, } impl RedisBatch { pub(crate) fn new(client: RedisClient) -> Self { + Self::build(client, Mode::Batch) + } + + /// A batch whose every op goes out on its own, the way the `RedisCache` method behind it sends + /// it, so a fresh one is "send it now": the first `resolve` sends it. + pub fn alone(client: RedisClient) -> Self { + Self::build(client, Mode::Alone) + } + + /// A batch sent as one redis-py pipeline, like `async_increment_pipeline`, on every topology. + pub fn pipelined(client: RedisClient) -> Self { + Self::build(client, Mode::Pipeline) + } + + fn build(client: RedisClient, mode: Mode) -> Self { Self(Arc::new(Shared { client, + mode, pending: Mutex::new(Pending::default()), flushing: tokio::sync::Mutex::new(()), })) } - pub fn get(&self, key: &str) -> Declared> { - self.declare(BatchOp::Get(key.to_owned()), decode) - } - - /// One `MGET` of `keys` in order, repeated keys read once, like `RedisBatch.mget`. - pub fn mget(&self, keys: &[&str]) -> Declared>> { - let mut unique: Vec = Vec::with_capacity(keys.len()); - for key in keys { - if !unique.iter().any(|seen| seen == key) { - unique.push((*key).to_owned()); - } - } - self.declare(BatchOp::MGet(unique), decode) - } - - pub fn set(&self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Declared<()> { - let value = value.to_redis_args().concat(); - self.declare( - BatchOp::Set { - key: key.to_owned(), - value, - ttl, - }, - decode, - ) - } - - pub fn del(&self, key: &str) -> Declared { - self.declare(BatchOp::Del(key.to_owned()), decode) - } - - pub fn expire(&self, key: &str, ttl: Duration) -> Declared { - self.declare( - BatchOp::Expire { - key: key.to_owned(), - ttl, - }, - decode, - ) - } - - /// `RedisBatch.increment`: `INCRBYFLOAT`, then `EXPIRE` when a TTL is given. - pub fn incr_by_float(&self, key: &str, amount: Amount, ttl: Option) -> Declared { - self.declare( - BatchOp::IncrByFloat { - key: key.to_owned(), - amount, - ttl, - }, - decode, - ) - } - - pub fn eval( - &self, - script: &Script, - keys: &[&str], - args: impl ToRedisArgs, - ) -> Declared { - self.declare( - BatchOp::Eval { - script: script.clone(), - keys: keys.iter().map(|key| (*key).to_owned()).collect(), - args: args.to_redis_args(), - }, - decode, - ) - } - - pub async fn flush(&self) { - let _flushing = self.0.flushing.lock().await; - let Pending { ops, replies } = std::mem::take(&mut *self.pending()); - if replies.is_empty() { - return; - } - let client = &self.0.client; - let results = match client.is_cluster() { - true => join_all(ops.iter().map(|op| op.run_alone(client))).await, - false => client.run(&pipeline(client, &ops)).await, - }; - for (reply, result) in replies.into_iter().zip(results) { - let _ = reply.send(result); - } - } - - fn declare(&self, op: BatchOp, decode: fn(Value) -> Result) -> Declared { + pub fn add(&self, command: Command) -> Declared { + let (op, decode) = command.prepare(self.0.client.namespace()); let (sender, reply) = oneshot::channel(); let mut pending = self.pending(); pending.ops.push(op); @@ -166,6 +84,31 @@ impl RedisBatch { } } + pub async fn flush(&self) { + let _flushing = self.0.flushing.lock().await; + let Pending { ops, replies } = std::mem::take(&mut *self.pending()); + if replies.is_empty() { + return; + } + let client = &self.0.client; + let alone = match self.0.mode { + Mode::Alone => true, + Mode::Batch => client.is_cluster(), + Mode::Pipeline => false, + }; + let results = match alone { + true => join_all(ops.into_iter().map(|op| client.run_alone(op))).await, + false => { + client + .run(&Pipeline::of(client.namespace().clone(), ops)) + .await + } + }; + for (reply, result) in replies.into_iter().zip(results) { + let _ = reply.send(result); + } + } + fn pending(&self) -> MutexGuard<'_, Pending> { self.0 .pending @@ -174,74 +117,6 @@ impl RedisBatch { } } -fn pipeline(client: &RedisClient, ops: &[BatchOp]) -> Pipeline { - ops.iter().fold(client.pipeline(), |pipeline, op| match op { - BatchOp::Get(key) => pipeline.get(key), - BatchOp::MGet(keys) => { - let keys: Vec<&str> = keys.iter().map(String::as_str).collect(); - pipeline.mget(&keys) - } - BatchOp::Set { key, value, ttl } => pipeline.set(key, value.as_slice(), *ttl), - BatchOp::Del(key) => pipeline.del(key), - BatchOp::Expire { key, ttl } => pipeline.expire(key, *ttl), - BatchOp::IncrByFloat { key, amount, ttl } => pipeline.incr_by_float(key, *amount, *ttl), - BatchOp::Eval { script, keys, args } => { - let keys: Vec<&str> = keys.iter().map(String::as_str).collect(); - pipeline.eval(script, &keys, args.as_slice()) - } - }) -} - -impl BatchOp { - /// What `RedisBatch` runs for an op on a cluster, where it never pipelines: the `RedisCache` - /// method behind it, so an increment there keeps a TTL it already has like `async_increment`. - async fn run_alone(&self, client: &RedisClient) -> Result { - match self { - Self::Get(key) => client - .get::(key) - .await - .map(|value| value.unwrap_or(Value::Nil)), - Self::MGet(keys) => { - let keys: Vec<&str> = keys.iter().map(String::as_str).collect(); - let values = client.mget::(&keys).await?; - Ok(Value::Array( - values - .into_iter() - .map(|value| value.unwrap_or(Value::Nil)) - .collect(), - )) - } - Self::Set { key, value, ttl } => client - .set(key, value.as_slice(), *ttl) - .await - .map(|()| Value::Okay), - Self::Del(key) => client - .del(&[key]) - .await - .map(|deleted| Value::Int(i64::try_from(deleted).unwrap_or(i64::MAX))), - Self::Expire { key, ttl } => client - .expire(key, *ttl) - .await - .map(|set| Value::Int(i64::from(set))), - Self::IncrByFloat { key, amount, ttl } => { - let value = match ttl { - Some(ttl) => { - client - .incr_by_float(key, *amount, Expiry::KeepExisting(*ttl)) - .await? - } - None => client.incr_by_float_without_ttl(key, *amount).await?, - }; - Ok(Value::Double(value)) - } - Self::Eval { script, keys, args } => { - let keys: Vec<&str> = keys.iter().map(String::as_str).collect(); - client.eval::(script, &keys, args.as_slice()).await - } - } - } -} - impl Declared { pub async fn resolve(mut self) -> Result { match self.reply.try_recv() { @@ -257,10 +132,11 @@ impl Declared { let reply = self.reply.await.unwrap_or_else(|_| Err(abandoned())); reply.and_then(self.decode) } -} -fn decode(value: Value) -> Result { - T::from_redis_value(value).map_err(|error| Error::Decode(error.into())) + /// Every reply in order, flushing their batch once if nobody has yet. + pub async fn all(declared: impl IntoIterator) -> Vec> { + join_all(declared.into_iter().map(Self::resolve)).await + } } fn abandoned() -> Error { diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index a9b4d9deb1e..25367fd63ba 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -6,11 +6,13 @@ use std::{ use futures_util::stream::BoxStream; use litellm_circuit_breaker::{Admission, Breaker, Failure, State}; -use redis::{Cmd, FromRedisValue, RedisError, ToRedisArgs, ToSingleRedisArg, Value}; +use redis::{Cmd, RedisError, Value}; use crate::{ - Amount, CallEvent, ConnectError, Error, Expiry, FailureKind, Namespace, Operation, Outcome, - Pipeline, RedisBatch, RedisConfig, RedisEvents, Script, commands, connection::Connection, + CallEvent, Command, ConnectError, Error, FailureKind, Namespace, Operation, Outcome, Pipeline, + RedisBatch, RedisConfig, RedisEvents, + commands::{self, Op, Plan, Trips}, + connection::Connection, pubsub::Subscriber, }; @@ -74,126 +76,70 @@ impl RedisClient { } pub async fn ping(&self) -> Result<(), Error> { - self.query("PING", &commands::ping()).await - } - - pub async fn get(&self, key: &str) -> Result, Error> { - self.query("GET", &commands::get(self.key(key))).await - } - - pub async fn mget(&self, keys: &[&str]) -> Result>, Error> { - if keys.is_empty() { - return Ok(Vec::new()); - } - self.query("MGET", &commands::mget(self.keys(keys))).await - } - - pub async fn set( - &self, - key: &str, - value: impl ToSingleRedisArg, - ttl: Duration, - ) -> Result<(), Error> { - self.query("SET", &commands::set(self.key(key), value, ttl)) - .await - } - - pub async fn set_nx( - &self, - key: &str, - value: impl ToSingleRedisArg, - ttl: Duration, - ) -> Result { - let reply: Option = self - .query("SET", &commands::set_nx(self.key(key), value, ttl)) - .await?; - Ok(reply.is_some()) - } - - pub async fn del(&self, keys: &[&str]) -> Result { - if keys.is_empty() { - return Ok(0); - } - self.query("DEL", &commands::del(self.keys(keys))).await - } - - pub async fn expire(&self, key: &str, ttl: Duration) -> Result { - self.query("EXPIRE", &commands::expire(self.key(key), ttl)) - .await - } - - pub async fn ttl(&self, key: &str) -> Result { - let reply: i64 = self.query("TTL", &commands::ttl(self.key(key))).await?; - Ok(match reply { - -2 => Ttl::Missing, - -1 => Ttl::Persistent, - seconds => Ttl::Expires(Duration::from_secs(seconds.unsigned_abs())), - }) - } - - /// `RedisCache.increment_cache`: `INCRBY`, then `TTL`, then `EXPIRE` only when the key has no - /// TTL, each its own round trip. - pub async fn incr_by(&self, key: &str, delta: i64, ttl: Duration) -> Result { - let key = self.key(key); - self.guarded(Operation::Command("INCRBY"), async { - let mut connection = self.0.connection.clone(); - let value: i64 = commands::incr_by(key.clone(), delta) - .query_async(&mut connection) - .await?; - let current: i64 = commands::ttl(key.clone()) - .query_async(&mut connection) - .await?; - if current == -1 { - commands::expire(key, ttl) - .exec_async(&mut connection) - .await?; - } - Ok(value) + self.guarded(Operation::Command("PING"), async { + commands::ping() + .query_async(&mut self.0.connection.clone()) + .await }) .await } - /// `INCRBYFLOAT` alone, `_incrbyfloat_with_ttl` without a TTL. - pub(crate) async fn incr_by_float_without_ttl( - &self, - key: &str, - amount: Amount, - ) -> Result { - self.query( - "INCRBYFLOAT", - &commands::incr_by_float(self.key(key), amount), - ) - .await + /// Sends one command on its own, the way the `RedisCache` method behind it sends it. + pub async fn send(&self, command: Command) -> Result { + let (op, decode) = command.prepare(&self.0.namespace); + self.run_alone(op).await.and_then(decode) } - /// `RedisCache.async_increment`: `INCRBYFLOAT` and its TTL command in one round trip, and for - /// [`Expiry::KeepExisting`] a third, `EXPIRE`, only when the key has no TTL. - pub async fn incr_by_float( - &self, - key: &str, - amount: Amount, - expiry: Expiry, - ) -> Result { - let key = self.key(key); - self.guarded(Operation::Command("INCRBYFLOAT"), async { + pub(crate) async fn run_alone(&self, op: Op) -> Result { + let Op { + operation, plan, .. + } = op; + self.guarded(operation, async { let mut connection = self.0.connection.clone(); - let mut pipeline = redis::pipe(); - pipeline.add_command(commands::incr_by_float(key.clone(), amount)); - match expiry { - Expiry::Refresh(ttl) => { - let (value, _): (f64, bool) = pipeline - .add_command(commands::expire(key, ttl)) - .query_async(&mut connection) - .await?; - Ok(value) + match plan { + Plan::Reply(value) => Ok(value), + Plan::Commands(commands) => match <[Cmd; 1]>::try_from(commands) { + Ok([command]) => command.query_async(&mut connection).await, + Err(commands) => { + let mut pipeline = redis::pipe(); + for command in commands { + pipeline.add_command(command); + } + let replies: Vec = pipeline.query_async(&mut connection).await?; + Ok(replies.into_iter().next().unwrap_or(Value::Nil)) + } + }, + Plan::Script(call) => { + call.script + .invocation(&call.keys, &call.args) + .invoke_async(&mut connection) + .await } - Expiry::KeepExisting(ttl) => { - let (value, current): (f64, i64) = pipeline - .add_command(commands::ttl(key.clone())) - .query_async(&mut connection) - .await?; + Plan::KeepingTtl { + increment, + key, + ttl, + trips, + } => { + let (value, current): (Value, i64) = match trips { + Trips::Shared => { + redis::pipe() + .add_command(increment) + .add_command(commands::ttl_cmd(&key)) + .query_async(&mut connection) + .await? + } + Trips::Separate => ( + increment.query_async(&mut connection).await?, + commands::ttl_cmd(&key) + .query_async(&mut connection) + .await?, + ), + }; if current == -1 { - commands::expire(key, ttl) + redis::cmd("EXPIRE") + .arg(&key) + .arg(ttl.as_secs()) .exec_async(&mut connection) .await?; } @@ -204,60 +150,20 @@ impl RedisClient { .await } - pub async fn eval( - &self, - script: &Script, - keys: &[&str], - args: impl ToRedisArgs, - ) -> Result { - let keys = self.keys(keys); - let args = args.to_redis_args(); - self.guarded(Operation::Script, async { - script - .invocation(&keys, &args) - .invoke_async(&mut self.0.connection.clone()) - .await - }) - .await - } - - /// The channel is sent as given, without the namespace, like redis-py's `publish`. - pub async fn publish(&self, channel: &str, payload: &[u8]) -> Result<(), Error> { - self.query("PUBLISH", &commands::publish(channel, payload)) - .await - } - /// Messages published on `channel`, on a connection of its own. After the connection is lost /// it reconnects and resubscribes, and messages published meanwhile are not delivered. pub fn subscribe(&self, channel: &str) -> BoxStream<'static, Vec> { self.0.subscriber.clone().messages(channel.to_owned()) } - /// A script sent with its source as `EVAL` every time, like redis-py's `Redis.eval`, for the - /// scripts `RedisCache` runs that way rather than through `register_script`. - pub async fn eval_source( - &self, - source: &str, - keys: &[&str], - args: impl ToRedisArgs, - ) -> Result { - let keys = self.keys(keys); - let command = redis::cmd("EVAL") - .arg(source) - .arg(keys.len()) - .arg(keys) - .arg(args) - .to_owned(); - self.guarded(Operation::Script, async { - command.query_async(&mut self.0.connection.clone()).await - }) - .await - } - pub fn pipeline(&self) -> Pipeline { Pipeline::new(self.0.namespace.clone()) } + pub(crate) fn namespace(&self) -> &Namespace { + &self.0.namespace + } + pub fn is_cluster(&self) -> bool { matches!(self.0.connection, Connection::Cluster(_)) } @@ -305,25 +211,6 @@ impl RedisClient { results } - fn key(&self, key: &str) -> String { - self.0.namespace.apply(key) - } - - fn keys(&self, keys: &[&str]) -> Vec { - keys.iter().map(|key| self.key(key)).collect() - } - - async fn query( - &self, - name: &'static str, - command: &Cmd, - ) -> Result { - self.guarded(Operation::Command(name), async { - command.query_async(&mut self.0.connection.clone()).await - }) - .await - } - async fn guarded( &self, operation: Operation, diff --git a/litellm-rust/crates/redis-client/src/commands.rs b/litellm-rust/crates/redis-client/src/commands.rs index 6d7485c7bdb..232c1666b81 100644 --- a/litellm-rust/crates/redis-client/src/commands.rs +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -1,7 +1,11 @@ use std::time::Duration; use litellm_python_compat::repr::float_repr; -use redis::{Cmd, RedisWrite, ToRedisArgs, ToSingleRedisArg}; +use redis::{ + Cmd, ErrorKind, FromRedisValue, RedisError, RedisWrite, ToRedisArgs, ToSingleRedisArg, Value, +}; + +use crate::{Error, Namespace, Operation, Script, Ttl, script::ScriptCall}; /// How an increment treats the key's TTL, the two ways `RedisCache` increments with one. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -9,6 +13,7 @@ pub enum Expiry { /// `async_increment(refresh_ttl=True)` and every pipelined increment: `EXPIRE` after each write. Refresh(Duration), /// `async_increment(refresh_ttl=False)`: read the `TTL` and `EXPIRE` only a key that has none. + /// Pipelined it is lowered to [`Expiry::Refresh`], like `RedisBatch.increment`. KeepExisting(Duration), } @@ -30,58 +35,357 @@ impl ToRedisArgs for Amount { impl ToSingleRedisArg for Amount {} -pub(crate) fn ping() -> Cmd { - redis::cmd("PING") +type Build = Box) -> Plan + Send>; +type Decode = Box Result + Send>; + +/// One Redis command as a value: what it sends alone and pipelined, and how its reply decodes. Every +/// executor (`RedisClient::send`, `Pipeline::add`, `RedisBatch::add`) applies its namespace to the +/// keys once. +#[must_use] +pub struct Command { + operation: Operation, + keys: Vec, + idempotent: bool, + build: Build, + decode: Decode, } -pub(crate) fn publish(channel: &str, payload: &[u8]) -> Cmd { - redis::cmd("PUBLISH").arg(channel).arg(payload).to_owned() +/// What a command sends once its keys carry the namespace. Only the cases whose alone and pipelined +/// forms differ get a variant of their own. +pub(crate) enum Plan { + /// Answered without a round trip, like an empty `MGET` or `DEL`. + Reply(Value), + /// Sent together; the reply is the first command's, failing if any failed. + Commands(Vec), + /// Alone through `invoke_async`, pipelined as `EVALSHA` resent once on `NOSCRIPT`. + Script(ScriptCall), + /// Alone the increment reads the `TTL` and sends `EXPIRE` only to a key without one, pipelined it + /// is the increment then `EXPIRE`. + KeepingTtl { + increment: Cmd, + key: String, + ttl: Duration, + trips: Trips, + }, } -pub(crate) fn get(key: String) -> Cmd { - redis::cmd("GET").arg(key).to_owned() +#[derive(Clone, Copy)] +pub(crate) enum Trips { + /// `async_increment`: the increment and its `TTL` in one round trip. + Shared, + /// `increment_cache`: the increment and its `TTL` each on its own round trip. + Separate, } -pub(crate) fn mget(keys: Vec) -> Cmd { - redis::cmd("MGET").arg(keys).to_owned() +/// A command with the namespace applied, ready for an executor to send. +pub(crate) struct Op { + pub(crate) operation: Operation, + pub(crate) key: String, + pub(crate) idempotent: bool, + pub(crate) plan: Plan, } -// redis-py's `set(ex=...)` takes whole seconds, as `int(ttl)` or a `timedelta`'s `int(total_seconds())` -pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { - redis::cmd("SET") - .arg(key) - .arg(value) - .arg("EX") - .arg(ttl.as_secs()) - .to_owned() +impl Plan { + /// What a pipeline sends for this op. + pub(crate) fn pipelined(&self) -> Vec { + match self { + Self::Reply(_) => Vec::new(), + Self::Commands(commands) => commands.clone(), + Self::Script(call) => vec![call.script.evalsha(&call.keys, &call.args)], + Self::KeepingTtl { + increment, key, ttl, .. + } => vec![increment.clone(), expire_cmd(key, *ttl)], + } + } } -pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { - redis::cmd("SET") - .arg(key) - .arg(value) - .arg("EX") - .arg(ttl.as_secs()) - .arg("NX") - .to_owned() +impl Command { + fn new( + operation: Operation, + keys: Vec, + idempotent: bool, + build: impl FnOnce(Vec) -> Plan + Send + 'static, + decode: impl FnOnce(Value) -> Result + Send + 'static, + ) -> Self { + Self { + operation, + keys, + idempotent, + build: Box::new(build), + decode: Box::new(decode), + } + } + + /// Attaches a codec over the decoded reply, such as a store's value format. + pub fn map(self, decode: fn(T) -> Result) -> Command + where + T: 'static, + { + let inner = self.decode; + Command { + operation: self.operation, + keys: self.keys, + idempotent: self.idempotent, + build: self.build, + decode: Box::new(move |value| inner(value).and_then(decode)), + } + } + + pub(crate) fn prepare(self, namespace: &Namespace) -> (Op, Decode) { + let keys: Vec = self.keys.iter().map(|key| namespace.apply(key)).collect(); + let key = keys.first().cloned().unwrap_or_default(); + let op = Op { + operation: self.operation, + key, + idempotent: self.idempotent, + plan: (self.build)(keys), + }; + (op, self.decode) + } } -pub(crate) fn del(keys: Vec) -> Cmd { - redis::cmd("DEL").arg(keys).to_owned() +fn decode(value: Value) -> Result { + T::from_redis_value(value).map_err(|error| Error::Decode(error.into())) } -pub(crate) fn expire(key: String, ttl: Duration) -> Cmd { +fn owned(keys: &[&str]) -> Vec { + keys.iter().map(|key| (*key).to_owned()).collect() +} + +fn single(keys: Vec) -> String { + keys.into_iter().next().unwrap_or_default() +} + +fn expire_cmd(key: &str, ttl: Duration) -> Cmd { redis::cmd("EXPIRE").arg(key).arg(ttl.as_secs()).to_owned() } -pub(crate) fn ttl(key: String) -> Cmd { +pub(crate) fn ttl_cmd(key: &str) -> Cmd { redis::cmd("TTL").arg(key).to_owned() } -pub(crate) fn incr_by(key: String, delta: i64) -> Cmd { - redis::cmd("INCRBY").arg(key).arg(delta).to_owned() +/// A reply whose shape a codec over it refuses, as `what` describes. +pub fn unexpected_reply(what: &'static str) -> Error { + Error::Decode(RedisError::from((ErrorKind::UnexpectedReturnType, what))) } -pub(crate) fn incr_by_float(key: String, amount: Amount) -> Cmd { - redis::cmd("INCRBYFLOAT").arg(key).arg(amount).to_owned() +pub fn get(key: &str) -> Command> { + Command::new( + Operation::Command("GET"), + vec![key.to_owned()], + true, + |keys| Plan::Commands(vec![redis::cmd("GET").arg(single(keys)).to_owned()]), + decode, + ) +} + +/// One `MGET` of `keys` in order, repeated keys read once, like `RedisBatch.mget`. No keys is no +/// round trip, since Redis rejects an empty `MGET`. Pipelined, every key must share a hash slot on a +/// cluster, and a batch never pipelines there. +pub fn mget(keys: &[&str]) -> Command>> { + let unique = keys.iter().fold(Vec::::new(), |mut unique, key| { + if !unique.iter().any(|seen| seen == key) { + unique.push((*key).to_owned()); + } + unique + }); + Command::new( + Operation::Command("MGET"), + unique, + true, + |keys| match keys.is_empty() { + true => Plan::Reply(Value::Array(Vec::new())), + false => Plan::Commands(vec![redis::cmd("MGET").arg(keys).to_owned()]), + }, + decode, + ) +} + +// redis-py's `set(ex=...)` takes whole seconds, as `int(ttl)` or a `timedelta`'s `int(total_seconds())` +pub fn set(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command<()> { + let value = value.to_redis_args().concat(); + Command::new( + Operation::Command("SET"), + vec![key.to_owned()], + true, + move |keys| { + Plan::Commands(vec![ + redis::cmd("SET") + .arg(single(keys)) + .arg(value) + .arg("EX") + .arg(ttl.as_secs()) + .to_owned(), + ]) + }, + decode, + ) +} + +/// `SET ... EX ... NX`, `EX` before `NX` as redis-py builds it. `true` when this call wrote the key. +pub fn set_nx(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command { + let value = value.to_redis_args().concat(); + Command::new( + Operation::Command("SET"), + vec![key.to_owned()], + false, + move |keys| { + Plan::Commands(vec![ + redis::cmd("SET") + .arg(single(keys)) + .arg(value) + .arg("EX") + .arg(ttl.as_secs()) + .arg("NX") + .to_owned(), + ]) + }, + |value| decode::>(value).map(|reply| reply.is_some()), + ) +} + +/// One multi-key `DEL`, no round trip for no keys. Pipelined, every key must share a hash slot. +pub fn del(keys: &[&str]) -> Command { + Command::new( + Operation::Command("DEL"), + owned(keys), + true, + |keys| match keys.is_empty() { + true => Plan::Reply(Value::Int(0)), + false => Plan::Commands(vec![redis::cmd("DEL").arg(keys).to_owned()]), + }, + decode, + ) +} + +pub fn expire(key: &str, ttl: Duration) -> Command { + Command::new( + Operation::Command("EXPIRE"), + vec![key.to_owned()], + true, + move |keys| Plan::Commands(vec![expire_cmd(&single(keys), ttl)]), + decode, + ) +} + +pub fn ttl(key: &str) -> Command { + Command::new( + Operation::Command("TTL"), + vec![key.to_owned()], + true, + |keys| Plan::Commands(vec![ttl_cmd(&single(keys))]), + |value| { + Ok(match decode::(value)? { + -2 => Ttl::Missing, + -1 => Ttl::Persistent, + seconds => Ttl::Expires(Duration::from_secs(seconds.unsigned_abs())), + }) + }, + ) +} + +/// `RedisCache.increment_cache`: `INCRBY`, then `TTL`, then `EXPIRE` only when the key has no TTL, +/// each its own round trip. +pub fn incr_by(key: &str, delta: i64, ttl: Duration) -> Command { + Command::new( + Operation::Command("INCRBY"), + vec![key.to_owned()], + false, + move |keys| { + let key = single(keys); + Plan::KeepingTtl { + increment: redis::cmd("INCRBY").arg(&key).arg(delta).to_owned(), + key, + ttl, + trips: Trips::Separate, + } + }, + decode, + ) +} + +/// `INCRBYFLOAT`, with the TTL handled as `expiry` says, or left alone without one. Alone it is +/// `RedisCache.async_increment`, pipelined `RedisBatch.increment`. +pub fn incr_by_float(key: &str, amount: Amount, expiry: Option) -> Command { + Command::new( + Operation::Command("INCRBYFLOAT"), + vec![key.to_owned()], + false, + move |keys| { + let key = single(keys); + let increment = redis::cmd("INCRBYFLOAT").arg(&key).arg(amount).to_owned(); + match expiry { + None => Plan::Commands(vec![increment]), + Some(Expiry::Refresh(ttl)) => { + Plan::Commands(vec![increment, expire_cmd(&key, ttl)]) + } + Some(Expiry::KeepExisting(ttl)) => Plan::KeepingTtl { + increment, + key, + ttl, + trips: Trips::Shared, + }, + } + }, + decode, + ) +} + +/// A registered script, like redis-py's `Script`. Every key it touches must be in `keys`. +pub fn eval( + script: &Script, + keys: &[&str], + args: impl ToRedisArgs, +) -> Command { + let script = script.clone(); + let args = args.to_redis_args(); + Command::new( + Operation::Script, + owned(keys), + false, + move |keys| Plan::Script(ScriptCall { script, keys, args }), + decode, + ) +} + +/// A script sent with its source as `EVAL` every time, like redis-py's `Redis.eval`, for the scripts +/// `RedisCache` runs that way rather than through `register_script`. +pub fn eval_source( + source: &'static str, + keys: &[&str], + args: impl ToRedisArgs, +) -> Command { + let args = args.to_redis_args(); + Command::new( + Operation::Script, + owned(keys), + false, + move |keys| { + Plan::Commands(vec![ + redis::cmd("EVAL") + .arg(source) + .arg(keys.len()) + .arg(keys) + .arg(args) + .to_owned(), + ]) + }, + decode, + ) +} + +/// The channel is sent as given, without the namespace, like redis-py's `publish`. +pub fn publish(channel: &str, payload: &[u8]) -> Command<()> { + let command = redis::cmd("PUBLISH").arg(channel).arg(payload).to_owned(); + Command::new( + Operation::Command("PUBLISH"), + Vec::new(), + false, + move |_| Plan::Commands(vec![command]), + |_| Ok(()), + ) +} + +pub(crate) fn ping() -> Cmd { + redis::cmd("PING") } diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 7eaa3e21b57..55a7ba79362 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -1,6 +1,6 @@ mod batch; mod client; -mod commands; +pub mod commands; mod config; mod connection; mod error; @@ -12,7 +12,7 @@ mod script; pub use batch::{Declared, RedisBatch}; pub use client::{RedisClient, Ttl, WeakRedisClient}; -pub use commands::{Amount, Expiry}; +pub use commands::{Amount, Command, Expiry}; pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts}; pub use error::{ConnectError, Error, KeyError}; pub use events::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; diff --git a/litellm-rust/crates/redis-client/src/pipeline.rs b/litellm-rust/crates/redis-client/src/pipeline.rs index 84dd12fe914..7f3de875136 100644 --- a/litellm-rust/crates/redis-client/src/pipeline.rs +++ b/litellm-rust/crates/redis-client/src/pipeline.rs @@ -1,13 +1,13 @@ -use std::{collections::HashMap, time::Duration}; +use std::collections::HashMap; use futures_util::future::join_all; -use redis::{ - Cmd, ErrorKind, RedisError, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value, - cluster_routing::Slot, -}; +use redis::{ErrorKind, RedisError, RedisResult, ServerErrorKind, Value, cluster_routing::Slot}; use crate::{ - Amount, Error, Namespace, Script, commands, connection::Connection, script::ScriptCall, + Command, Error, Namespace, + commands::{Op, Plan}, + connection::Connection, + script::ScriptCall, }; pub struct Pipeline { @@ -15,29 +15,6 @@ pub struct Pipeline { ops: Vec, } -struct Op { - entry: Entry, - key: String, - idempotent: bool, -} - -/// What a pipeline sends for one op. Only a [`ScriptCall`] can be sent again, so a retry can never -/// repeat a command that already ran. -enum Entry { - /// The op's reply is the first command's, like `RedisBatch`'s `INCRBYFLOAT` plus `EXPIRE`. - Commands(Vec), - Script(ScriptCall), -} - -impl Entry { - fn commands(&self) -> Vec { - match self { - Self::Commands(commands) => commands.clone(), - Self::Script(call) => vec![call.script.evalsha(&call.keys, &call.args)], - } - } -} - type Settled = (usize, Result); impl Pipeline { @@ -48,67 +25,14 @@ impl Pipeline { } } - pub fn get(self, key: &str) -> Self { - let key = self.namespace.apply(key); - self.commands(vec![commands::get(key.clone())], key, true) + pub(crate) fn of(namespace: Namespace, ops: Vec) -> Self { + Self { namespace, ops } } - /// One `MGET`, so every key must share a hash slot on a cluster. Only `RedisBatch` sends it, - /// and a batch never pipelines on a cluster. - pub(crate) fn mget(self, keys: &[&str]) -> Self { - let keys: Vec = keys.iter().map(|key| self.namespace.apply(key)).collect(); - let key = keys.first().cloned().unwrap_or_default(); - self.commands(vec![commands::mget(keys)], key, true) - } - - pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self { - let key = self.namespace.apply(key); - self.commands(vec![commands::set(key.clone(), value, ttl)], key, true) - } - - pub fn del(self, key: &str) -> Self { - let key = self.namespace.apply(key); - self.commands(vec![commands::del(vec![key.clone()])], key, true) - } - - pub fn expire(self, key: &str, ttl: Duration) -> Self { - let key = self.namespace.apply(key); - self.commands(vec![commands::expire(key.clone(), ttl)], key, true) - } - - /// `INCRBYFLOAT`, then `EXPIRE` when a TTL is given, like `RedisBatch` and - /// `async_increment_pipeline` send it. - pub fn incr_by_float(self, key: &str, amount: Amount, ttl: Option) -> Self { - let key = self.namespace.apply(key); - let increment = commands::incr_by_float(key.clone(), amount); - let commands = match ttl { - Some(ttl) => vec![increment, commands::expire(key.clone(), ttl)], - None => vec![increment], - }; - self.commands(commands, key, false) - } - - pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self { - let keys: Vec = keys.iter().map(|key| self.namespace.apply(key)).collect(); - let key = keys.first().cloned().unwrap_or_default(); - let call = ScriptCall { - script: script.clone(), - keys, - args: args.to_redis_args(), - }; - self.push(Entry::Script(call), key, false) - } - - fn commands(self, commands: Vec, key: String, idempotent: bool) -> Self { - self.push(Entry::Commands(commands), key, idempotent) - } - - fn push(mut self, entry: Entry, key: String, idempotent: bool) -> Self { - self.ops.push(Op { - entry, - key, - idempotent, - }); + /// Appends one command in its pipelined form; `RedisClient::run` answers it in order, undecoded. + pub fn add(mut self, command: Command) -> Self { + let (op, _) = command.prepare(&self.namespace); + self.ops.push(op); self } @@ -175,8 +99,8 @@ impl Pipeline { index: usize, reply: RedisResult, ) -> (usize, RedisResult) { - match (&self.ops[index].entry, reply) { - (Entry::Script(call), Err(error)) if is_missing_script(&error) => { + match (&self.ops[index].plan, reply) { + (Plan::Script(call), Err(error)) if is_missing_script(&error) => { (index, rerun_missing_script(connection, call).await) } (_, reply) => (index, reply), @@ -190,7 +114,7 @@ impl Pipeline { let widths: Vec = indices .iter() .map(|&index| { - let commands = self.ops[index].entry.commands(); + let commands = self.ops[index].plan.pipelined(); let width = commands.len(); for command in commands { pipeline.add_command(command); @@ -198,8 +122,10 @@ impl Pipeline { width }) .collect(); - let replies: RedisResult>> = - pipeline.query_async(&mut connection).await; + let replies: RedisResult>> = match pipeline.is_empty() { + true => Ok(Vec::new()), + false => pipeline.query_async(&mut connection).await, + }; let replies = match replies { Ok(replies) => replies, Err(error) => { @@ -213,7 +139,10 @@ impl Pipeline { let per_op: Vec<(usize, RedisResult)> = indices .iter() .zip(widths) - .map(|(&index, width)| (index, first_reply(replies.by_ref().take(width)))) + .map(|(&index, width)| match &self.ops[index].plan { + Plan::Reply(value) => (index, Ok(value.clone())), + _ => (index, first_reply(replies.by_ref().take(width))), + }) .collect(); join_all( per_op diff --git a/litellm-rust/crates/redis-client/tests/batch.rs b/litellm-rust/crates/redis-client/tests/batch.rs index 034b1580b57..0ed02178a30 100644 --- a/litellm-rust/crates/redis-client/tests/batch.rs +++ b/litellm-rust/crates/redis-client/tests/batch.rs @@ -7,14 +7,15 @@ use std::{ use litellm_circuit_breaker::{State, Transition}; use litellm_redis_client::{ - Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Error, Namespace, Operation, RedisClient, - RedisConfig, RedisEvents, Script, Timeouts, Ttl, + Amount, CallEvent, Credentials, Declared, DEFAULT_TIMEOUT, Error, Expiry, Namespace, Operation, RedisClient, + RedisBatch, RedisConfig, RedisEvents, Script, Timeouts, Ttl, }; +use litellm_redis_client::commands; use litellm_redis_testing::{RedisDeployment, Topology}; use rstest::rstest; use support::{connect, endpoint}; -const MINUTE: Option = Some(Duration::from_secs(60)); +const MINUTE: Option = Some(Expiry::KeepExisting(Duration::from_secs(60))); // RedisBatch pipelines a standalone flush and runs each op on its own on a cluster fn expected_trips(topology: Topology, alone: &[Operation]) -> Vec { @@ -64,13 +65,13 @@ async fn declared_ops_wait_for_the_flush_and_share_one_round_trip( let server = RedisDeployment::start(topology).await.unwrap(); let (client, trips) = observed(&server); let batch = client.batch(); - let counter = batch.incr_by_float("counter", Amount::Int(3), MINUTE); - let total = batch.incr_by_float("total", Amount::Float(1.5), MINUTE); - let written = batch.set("other", "v", Duration::from_secs(60)); - let read = batch.get::("value"); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(3), MINUTE)); + let total = batch.add(commands::incr_by_float("total", Amount::Float(1.5), MINUTE)); + let written = batch.add(commands::set("other", "v", Duration::from_secs(60))); + let read = batch.add(commands::get::("value")); let outside = connect(&server, "ns", DEFAULT_TIMEOUT); - assert_eq!(outside.get::("counter").await.unwrap(), None); + assert_eq!(outside.send(commands::get::("counter")).await.unwrap(), None); assert!(trips.seen().is_empty()); batch.flush().await; @@ -104,12 +105,12 @@ async fn resolving_before_a_flush_sends_everything_declared_so_far( let server = RedisDeployment::start(topology).await.unwrap(); let (client, trips) = observed(&server); let batch = client.batch(); - let first = batch.incr_by_float("a", Amount::Int(1), MINUTE); - let second = batch.incr_by_float("b", Amount::Int(2), MINUTE); + let first = batch.add(commands::incr_by_float("a", Amount::Int(1), MINUTE)); + let second = batch.add(commands::incr_by_float("b", Amount::Int(2), MINUTE)); let increment = Operation::Command("INCRBYFLOAT"); assert_eq!(second.resolve().await.unwrap(), 2.0); - let late = batch.incr_by_float("a", Amount::Int(10), MINUTE); + let late = batch.add(commands::incr_by_float("a", Amount::Int(10), MINUTE)); assert_eq!(first.resolve().await.unwrap(), 1.0); assert_eq!( trips.seen(), @@ -126,6 +127,22 @@ async fn resolving_before_a_flush_sends_everything_declared_so_far( assert_eq!(trips.seen().len(), both.len()); } +#[rstest] +#[tokio::test] +async fn a_pipelined_batch_is_one_pipeline_on_every_topology( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = RedisBatch::pipelined(client); + let declared = ["a", "b", "c"].map(|key| batch.add(commands::incr_by_float(key, Amount::Int(2), MINUTE))); + + let totals = Declared::all(declared).await; + + assert_eq!(totals.into_iter().map(Result::unwrap).collect::>(), [2.0; 3]); + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 3 }]); +} + #[rstest] #[tokio::test] async fn each_declared_op_gets_its_own_typed_result( @@ -134,16 +151,19 @@ async fn each_declared_op_gets_its_own_typed_result( let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", DEFAULT_TIMEOUT); client - .set("text", "not a number", Duration::from_secs(60)) + .send(commands::set("text", "not a number", Duration::from_secs(60))) .await .unwrap(); let batch = client.batch(); - let rejected = batch.incr_by_float("text", Amount::Int(1), MINUTE); - let undecodable = batch.get::("text"); - let missing = batch.get::("missing"); - let script = - batch.eval::>(&Script::new("return {KEYS[1], ARGV[1]}"), &["k"], "arg"); - let removed = batch.del("text"); + let rejected = batch.add(commands::incr_by_float("text", Amount::Int(1), MINUTE)); + let undecodable = batch.add(commands::get::("text")); + let missing = batch.add(commands::get::("missing")); + let script = batch.add(commands::eval::>( + &Script::new("return {KEYS[1], ARGV[1]}"), + &["k"], + "arg", + )); + let removed = batch.add(commands::del(&["text"])); batch.flush().await; @@ -162,7 +182,7 @@ async fn an_op_whose_flush_was_cancelled_resolves_as_transport( let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", Duration::from_secs(5)); let batch = client.batch(); - let counter = batch.incr_by_float("counter", Amount::Int(1), MINUTE); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE)); server.freeze().await.unwrap(); let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await; @@ -182,7 +202,7 @@ async fn a_settled_op_waits_for_someone_else_to_flush_and_sends_nothing_itself( let batch = client.batch(); let counter = tokio::spawn( batch - .incr_by_float("counter", Amount::Int(4), MINUTE) + .add(commands::incr_by_float("counter", Amount::Int(4), MINUTE)) .settled(), ); @@ -205,7 +225,7 @@ async fn a_settled_op_whose_flush_was_cancelled_resolves_as_transport() { let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); let client = connect(&server, "", Duration::from_secs(5)); let batch = client.batch(); - let counter = batch.incr_by_float("counter", Amount::Int(1), MINUTE); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE)); server.freeze().await.unwrap(); let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await; @@ -233,15 +253,15 @@ async fn mget_reads_each_key_once_in_order_in_one_command( ) { let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", DEFAULT_TIMEOUT); - client.set("a", "1", Duration::from_secs(60)).await.unwrap(); - client.set("c", "3", Duration::from_secs(60)).await.unwrap(); + client.send(commands::set("a", "1", Duration::from_secs(60))).await.unwrap(); + client.send(commands::set("c", "3", Duration::from_secs(60))).await.unwrap(); server .cli_on_every_node(&["CONFIG", "RESETSTAT"]) .await .unwrap(); let batch = client.batch(); - let values = batch.mget::(&["c", "missing", "a", "c"]); + let values = batch.add(commands::mget::(&["c", "missing", "a", "c"])); batch.flush().await; assert_eq!( @@ -266,16 +286,16 @@ async fn a_batched_increment_treats_the_ttl_like_redis_batch( let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", DEFAULT_TIMEOUT); client - .set("counter", "1", Duration::from_secs(30)) + .send(commands::set("counter", "1", Duration::from_secs(30))) .await .unwrap(); let batch = client.batch(); - let counter = batch.incr_by_float("counter", Amount::Int(1), Some(Duration::from_secs(600))); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), Some(Expiry::KeepExisting(Duration::from_secs(600))))); batch.flush().await; assert_eq!(counter.resolve().await.unwrap(), 2.0); - let Ttl::Expires(left) = client.ttl("counter").await.unwrap() else { + let Ttl::Expires(left) = client.send(commands::ttl("counter")).await.unwrap() else { panic!("expected an expiring counter") }; assert_eq!(left <= Duration::from_secs(30), keeps_the_ttl, "{left:?}"); diff --git a/litellm-rust/crates/redis-client/tests/client.rs b/litellm-rust/crates/redis-client/tests/client.rs index 13105dde6de..e77cad858fb 100644 --- a/litellm-rust/crates/redis-client/tests/client.rs +++ b/litellm-rust/crates/redis-client/tests/client.rs @@ -13,15 +13,16 @@ use litellm_circuit_breaker::{Breaker, Config, State}; use litellm_redis_client::{ Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Error, Expiry, FailureKind, Namespace, Operation, Outcome, RedisClient, RedisConfig, RedisEvents, RedisUrl, Script, Standalone, - Target, Timeouts, Ttl, + Target, Timeouts, Ttl, Value, }; +use litellm_redis_client::commands; use litellm_secret_value::SecretValue; use support::{KEEP, connect, connect_guarded, connections, endpoint}; async fn persist(client: &RedisClient, key: &str) { let set = Script::new("return redis.call('SET', KEYS[1], '0')"); client - .eval::<()>(&set, &[key], Vec::::new()) + .send(commands::eval::<()>(&set, &[key], Vec::::new())) .await .unwrap(); } @@ -37,12 +38,12 @@ const MINUTE: Duration = Duration::from_secs(60); async fn write(client: &RedisClient, counter: Counter, key: &str, expiry: Expiry) { match (counter, expiry) { (Counter::IncrBy, Expiry::KeepExisting(ttl)) => { - client.incr_by(key, 1, ttl).await.unwrap(); + client.send(commands::incr_by(key, 1, ttl)).await.unwrap(); } (Counter::IncrBy, Expiry::Refresh(_)) => unreachable!("increment_cache never refreshes"), (Counter::IncrByFloat, expiry) => { client - .incr_by_float(key, Amount::Float(0.5), expiry) + .send(commands::incr_by_float(key, Amount::Float(0.5), Some(expiry))) .await .unwrap(); } @@ -52,7 +53,7 @@ async fn write(client: &RedisClient, counter: Counter, key: &str, expiry: Expiry async fn persist_as(client: &RedisClient, key: &str) { let source = "return redis.call('SET', KEYS[1], '1')"; client - .eval::(&Script::new(source), &[key], Vec::::new()) + .send(commands::eval::(&Script::new(source), &[key], Vec::::new())) .await .unwrap(); } @@ -79,7 +80,7 @@ async fn a_counter_write_expires_a_new_key( write(&client, counter, "fresh", expiry).await; - let left = remaining(client.ttl("fresh").await.unwrap()); + let left = remaining(client.send(commands::ttl("fresh")).await.unwrap()); assert!(left > Duration::from_secs(50) && left <= MINUTE, "{left:?}"); } @@ -124,9 +125,9 @@ async fn keep_existing_leaves_a_running_ttl_and_refresh_resets_it( ) .await; - assert!(remaining(client.ttl("kept").await.unwrap()) <= Duration::from_secs(30)); - assert!(remaining(client.ttl("kept_int").await.unwrap()) <= Duration::from_secs(30)); - assert!(remaining(client.ttl("refreshed").await.unwrap()) > Duration::from_secs(500)); + assert!(remaining(client.send(commands::ttl("kept")).await.unwrap()) <= Duration::from_secs(30)); + assert!(remaining(client.send(commands::ttl("kept_int")).await.unwrap()) <= Duration::from_secs(30)); + assert!(remaining(client.send(commands::ttl("refreshed")).await.unwrap()) > Duration::from_secs(500)); } #[rstest] @@ -141,7 +142,7 @@ async fn keep_existing_expires_a_key_that_had_no_ttl( write(&client, counter, "persistent", KEEP).await; - let left = remaining(client.ttl("persistent").await.unwrap()); + let left = remaining(client.send(commands::ttl("persistent")).await.unwrap()); assert!(left > Duration::from_secs(50), "{left:?}"); } @@ -157,21 +158,21 @@ async fn counters_send_the_commands_redis_cache_sends( .await .unwrap(); - client.incr_by("int", 1, MINUTE).await.unwrap(); + client.send(commands::incr_by("int", 1, MINUTE)).await.unwrap(); client - .incr_by_float("kept", Amount::Int(1), KEEP) + .send(commands::incr_by_float("kept", Amount::Int(1), Some(KEEP))) .await .unwrap(); client - .incr_by_float("refreshed", Amount::Float(0.5), Expiry::Refresh(MINUTE)) + .send(commands::incr_by_float("refreshed", Amount::Float(0.5), Some(Expiry::Refresh(MINUTE)))) .await .unwrap(); client .run( &client .pipeline() - .incr_by_float("piped", Amount::Float(0.25), Some(MINUTE)) - .incr_by_float("bare", Amount::Float(0.25), None), + .add(commands::incr_by_float("piped", Amount::Float(0.25), Some(Expiry::Refresh(MINUTE)))) + .add(commands::incr_by_float("bare", Amount::Float(0.25), None)), ) .await; @@ -204,13 +205,13 @@ async fn integer_increments_preserve_every_bit( let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", DEFAULT_TIMEOUT); client - .set("integer", initial, Duration::from_secs(60)) + .send(commands::set("integer", initial, Duration::from_secs(60))) .await .unwrap(); let expected = initial + delta; - let result = client.incr_by("integer", delta, MINUTE).await.unwrap(); + let result = client.send(commands::incr_by("integer", delta, MINUTE)).await.unwrap(); assert_eq!(result, expected); - assert_eq!(client.get::("integer").await.unwrap(), Some(expected)); + assert_eq!(client.send(commands::get::("integer")).await.unwrap(), Some(expected)); } #[rstest] @@ -240,12 +241,12 @@ async fn outage_events_report_the_failure_kind( } let result = if pipeline { client - .run(&client.pipeline().get("missing")) + .run(&client.pipeline().add(commands::get::("missing"))) .await .remove(0) .map(drop) } else { - client.get::("missing").await.map(drop) + client.send(commands::get::("missing")).await.map(drop) }; if timeout { server.thaw().await.unwrap(); @@ -290,19 +291,19 @@ async fn every_call_reports_its_operation_and_outcome( }) .unwrap(); client - .set("list", "item", Duration::from_secs(60)) + .send(commands::set("list", "item", Duration::from_secs(60))) .await .unwrap(); - let _ = client.incr_by("list", 1, Duration::from_secs(60)).await; - let _ = client.get::("missing").await; - let _ = client.run(&client.pipeline().get("a").get("b")).await; + let _ = client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await; + let _ = client.send(commands::get::("missing")).await; + let _ = client.run(&client.pipeline().add(commands::get::("a")).add(commands::get::("b"))).await; let failures = client - .run(&client.pipeline().incr_by_float( + .run(&client.pipeline().add(commands::incr_by_float( "list", Amount::Int(1), - Some(Duration::from_secs(60)), - )) + Some(Expiry::Refresh(Duration::from_secs(60))), + ))) .await; assert!(matches!(failures[0], Err(Error::Server(_)))); @@ -341,15 +342,15 @@ async fn namespaces_isolate_keys_on_the_server( let b = connect(&server, "b", DEFAULT_TIMEOUT); let bare = connect(&server, "", DEFAULT_TIMEOUT); - a.set("x", "1", Duration::from_secs(60)).await.unwrap(); + a.send(commands::set("x", "1", Duration::from_secs(60))).await.unwrap(); - assert_eq!(a.get::("x").await.unwrap().as_deref(), Some("1")); - assert_eq!(a.get::("a:x").await.unwrap().as_deref(), Some("1")); + assert_eq!(a.send(commands::get::("x")).await.unwrap().as_deref(), Some("1")); + assert_eq!(a.send(commands::get::("a:x")).await.unwrap().as_deref(), Some("1")); assert_eq!( - bare.get::("a:x").await.unwrap().as_deref(), + bare.send(commands::get::("a:x")).await.unwrap().as_deref(), Some("1") ); - assert_eq!(b.get::("x").await.unwrap(), None); + assert_eq!(b.send(commands::get::("x")).await.unwrap(), None); } #[rstest] @@ -360,16 +361,16 @@ async fn a_rejected_command_keeps_the_connection( let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", DEFAULT_TIMEOUT); client - .set("list", "item", Duration::from_secs(60)) + .send(commands::set("list", "item", Duration::from_secs(60))) .await .unwrap(); let before = connections(&server).await; - let rejected = client.incr_by("list", 1, Duration::from_secs(60)).await; + let rejected = client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await; assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); assert_eq!( - client.get::("list").await.unwrap().as_deref(), + client.send(commands::get::("list")).await.unwrap().as_deref(), Some("item") ); assert_eq!(connections(&server).await, before); @@ -383,14 +384,14 @@ async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", DEFAULT_TIMEOUT); client - .set("float", "15.0", Duration::from_secs(60)) + .send(commands::set("float", "15.0", Duration::from_secs(60))) .await .unwrap(); - let decoded = client.get::("float").await; + let decoded = client.send(commands::get::("float")).await; assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); - assert_eq!(client.get::("missing").await.unwrap(), None); + assert_eq!(client.send(commands::get::("missing")).await.unwrap(), None); } #[rstest] @@ -401,16 +402,16 @@ async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", Duration::from_millis(500)); client - .set("first", "1", Duration::from_secs(60)) + .send(commands::set("first", "1", Duration::from_secs(60))) .await .unwrap(); client - .set("second", "2", Duration::from_secs(60)) + .send(commands::set("second", "2", Duration::from_secs(60))) .await .unwrap(); server.freeze().await.unwrap(); - let timed_out = client.get::("first").await; + let timed_out = client.send(commands::get::("first")).await; server.thaw().await.unwrap(); assert!( @@ -418,7 +419,7 @@ async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( "{timed_out:?}" ); assert_eq!( - client.get::("second").await.unwrap().as_deref(), + client.send(commands::get::("second")).await.unwrap().as_deref(), Some("2") ); } @@ -435,7 +436,7 @@ async fn the_same_client_recovers_after_redis_restarts( let recovered = eventually(async || { client - .incr_by("after_restart", 1, Duration::from_secs(60)) + .send(commands::incr_by("after_restart", 1, Duration::from_secs(60))) .await .ok() }) @@ -476,10 +477,10 @@ async fn empty_batches_succeed_without_asking_redis( let client = connect(&server, "", DEFAULT_TIMEOUT); assert_eq!( - client.mget::(&[]).await.unwrap(), + client.send(commands::mget::(&[])).await.unwrap(), Vec::>::new() ); - assert_eq!(client.del(&[]).await.unwrap(), 0); + assert_eq!(client.send(commands::del(&[])).await.unwrap(), 0); } #[rstest] @@ -489,11 +490,11 @@ async fn mget_and_del_span_slots_in_input_order( ) { let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", DEFAULT_TIMEOUT); - client.set("a", "1", Duration::from_secs(60)).await.unwrap(); - client.set("c", "3", Duration::from_secs(60)).await.unwrap(); + client.send(commands::set("a", "1", Duration::from_secs(60))).await.unwrap(); + client.send(commands::set("c", "3", Duration::from_secs(60))).await.unwrap(); - let values = client.mget::(&["c", "b", "a"]).await.unwrap(); - let deleted = client.del(&["a", "b", "c"]).await.unwrap(); + let values = client.send(commands::mget::(&["c", "b", "a"])).await.unwrap(); + let deleted = client.send(commands::del(&["a", "b", "c"])).await.unwrap(); assert_eq!( values, @@ -501,7 +502,7 @@ async fn mget_and_del_span_slots_in_input_order( ); assert_eq!(deleted, 2); assert_eq!( - client.mget::(&["a", "c"]).await.unwrap(), + client.send(commands::mget::(&["a", "c"])).await.unwrap(), vec![None, None] ); } @@ -515,18 +516,18 @@ async fn set_nx_only_writes_a_missing_key( let client = connect(&server, "", DEFAULT_TIMEOUT); let first = client - .set_nx("lock", "owner-1", Duration::from_secs(60)) + .send(commands::set_nx("lock", "owner-1", Duration::from_secs(60))) .await .unwrap(); let second = client - .set_nx("lock", "owner-2", Duration::from_secs(60)) + .send(commands::set_nx("lock", "owner-2", Duration::from_secs(60))) .await .unwrap(); assert!(first); assert!(!second); assert_eq!( - client.get::("lock").await.unwrap().as_deref(), + client.send(commands::get::("lock")).await.unwrap().as_deref(), Some("owner-1") ); } @@ -540,33 +541,33 @@ async fn ttl_reports_missing_persistent_and_remaining_time( let client = connect(&server, "", DEFAULT_TIMEOUT); persist(&client, "persistent").await; client - .set("expiring", "1", Duration::from_secs(60)) + .send(commands::set("expiring", "1", Duration::from_secs(60))) .await .unwrap(); client - .incr_by_float("expired_later", Amount::Float(0.5), KEEP) + .send(commands::incr_by_float("expired_later", Amount::Float(0.5), Some(KEEP))) .await .unwrap(); assert!( client - .expire("expired_later", Duration::from_secs(30)) + .send(commands::expire("expired_later", Duration::from_secs(30))) .await .unwrap() ); - assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing); - assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent); + assert_eq!(client.send(commands::ttl("missing")).await.unwrap(), Ttl::Missing); + assert_eq!(client.send(commands::ttl("persistent")).await.unwrap(), Ttl::Persistent); assert!(matches!( - client.ttl("expiring").await.unwrap(), + client.send(commands::ttl("expiring")).await.unwrap(), Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) )); assert!(matches!( - client.ttl("expired_later").await.unwrap(), + client.send(commands::ttl("expired_later")).await.unwrap(), Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) )); assert!( !client - .expire("missing", Duration::from_secs(30)) + .send(commands::expire("missing", Duration::from_secs(30))) .await .unwrap() ); @@ -581,13 +582,13 @@ async fn incr_by_float_accumulates_fractions( let client = connect(&server, "", DEFAULT_TIMEOUT); client - .incr_by_float("spend", Amount::Float(0.25), KEEP) + .send(commands::incr_by_float("spend", Amount::Float(0.25), Some(KEEP))) .await .unwrap(); assert_eq!( client - .incr_by_float("spend", Amount::Float(0.5), KEEP) + .send(commands::incr_by_float("spend", Amount::Float(0.5), Some(KEEP))) .await .unwrap(), 0.75 @@ -608,7 +609,7 @@ async fn concurrent_commands_on_one_cloned_client_all_apply( tokio::spawn(async move { for _ in 0..10 { client - .incr_by("shared", 1, Duration::from_secs(60)) + .send(commands::incr_by("shared", 1, Duration::from_secs(60))) .await .unwrap(); } @@ -619,7 +620,7 @@ async fn concurrent_commands_on_one_cloned_client_all_apply( task.await.unwrap(); } - assert_eq!(client.get::("shared").await.unwrap(), Some(1000)); + assert_eq!(client.send(commands::get::("shared")).await.unwrap(), Some(1000)); } fn breaker(failure_threshold: u32) -> Breaker { @@ -648,10 +649,10 @@ async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( ); server.freeze().await.unwrap(); - let first = client.get::("k").await; - let second = client.get::("k").await; + let first = client.send(commands::get::("k")).await; + let second = client.send(commands::get::("k")).await; let started = Instant::now(); - let third = client.get::("k").await; + let third = client.send(commands::get::("k")).await; let rejected_in = started.elapsed(); server.thaw().await.unwrap(); @@ -660,7 +661,7 @@ async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); - let recovered = eventually(async || client.get::("k").await.ok()).await; + let recovered = eventually(async || client.send(commands::get::("k")).await.ok()).await; assert_eq!(recovered, None); assert!( matches!(breaker.state(), State::Closed { .. }), @@ -678,14 +679,14 @@ async fn rejected_commands_never_open_the_breaker( let breaker = breaker(2); let client = connect_guarded(&server, "", DEFAULT_TIMEOUT, Some(breaker.clone())); client - .set("list", "item", Duration::from_secs(60)) + .send(commands::set("list", "item", Duration::from_secs(60))) .await .unwrap(); let rejected = [ - client.incr_by("list", 1, Duration::from_secs(60)).await, - client.incr_by("list", 1, Duration::from_secs(60)).await, - client.incr_by("list", 1, Duration::from_secs(60)).await, + client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await, + client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await, + client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await, ]; assert!( @@ -709,8 +710,8 @@ async fn an_open_breaker_fails_every_pipeline_op_without_sending( let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker(1))); let pipeline = client .pipeline() - .incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60))) - .incr_by_float("c", Amount::Int(1), Some(Duration::from_secs(60))); + .add(commands::incr_by_float("a", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60))))) + .add(commands::incr_by_float("c", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60))))); server.freeze().await.unwrap(); let failed = client.run(&pipeline).await; @@ -770,13 +771,13 @@ async fn a_host_endpoint_connects_to_its_database() { let via_url = connect(url(format!("{}/3", server.url()))).unwrap(); in_db_3 - .set("k", "3", Duration::from_secs(60)) + .send(commands::set("k", "3", Duration::from_secs(60))) .await .unwrap(); - assert_eq!(in_db_0.get::("k").await.unwrap(), None); + assert_eq!(in_db_0.send(commands::get::("k")).await.unwrap(), None); assert_eq!( - via_url.get::("k").await.unwrap().as_deref(), + via_url.send(commands::get::("k")).await.unwrap().as_deref(), Some("3") ); } @@ -826,13 +827,13 @@ async fn an_idle_client_pings_every_node_before_the_next_command_like_redis_py() unreachable!("started a cluster") }; let client = health_checked(&server, Duration::from_secs(3), DEFAULT_TIMEOUT, None, None); - client.set("k", "1", Duration::from_secs(60)).await.unwrap(); + client.send(commands::set("k", "1", Duration::from_secs(60))).await.unwrap(); let before = ping_calls(cluster).await; - client.set("k", "2", Duration::from_secs(60)).await.unwrap(); + client.send(commands::set("k", "2", Duration::from_secs(60))).await.unwrap(); let busy = ping_calls(cluster).await; tokio::time::sleep(Duration::from_millis(3500)).await; - client.set("k", "3", Duration::from_secs(60)).await.unwrap(); + client.send(commands::set("k", "3", Duration::from_secs(60))).await.unwrap(); let idle = ping_calls(cluster).await; assert_eq!(busy, before); @@ -856,11 +857,11 @@ async fn a_failed_health_check_is_neither_a_call_nor_a_breaker_failure() { Some(breaker.clone()), Some(recorded.clone()), ); - client.set("k", "1", Duration::from_secs(60)).await.unwrap(); + client.send(commands::set("k", "1", Duration::from_secs(60))).await.unwrap(); server.freeze().await.unwrap(); tokio::time::sleep(Duration::from_millis(300)).await; - let failed = client.get::("k").await; + let failed = client.send(commands::get::("k")).await; server.thaw().await.unwrap(); assert!(matches!(failed, Err(Error::Transport(_))), "{failed:?}"); diff --git a/litellm-rust/crates/redis-client/tests/connection.rs b/litellm-rust/crates/redis-client/tests/connection.rs index b1bce245b8c..4081ed6409c 100644 --- a/litellm-rust/crates/redis-client/tests/connection.rs +++ b/litellm-rust/crates/redis-client/tests/connection.rs @@ -12,6 +12,7 @@ use litellm_redis_client::{ Credentials, CredentialsProvider, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, Standalone, Target, Timeouts, }; +use litellm_redis_client::commands; use litellm_secret_value::SecretValue; use support::{Rotating, connect_as, connections}; @@ -38,7 +39,7 @@ async fn a_pushed_credential_reauthenticates_the_live_connection_and_a_reconnect &server, Credentials::Provider(CredentialsProvider(provider.clone())), ); - client.set("k", "1", Duration::from_secs(60)).await.unwrap(); + client.send(commands::set("k", "1", Duration::from_secs(60))).await.unwrap(); let before = connections(&server).await; provider.rotate("bob", "second"); @@ -55,7 +56,7 @@ async fn a_pushed_credential_reauthenticates_the_live_connection_and_a_reconnect .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) .await .unwrap(); - let after_reconnect = eventually(async || client.get::("k").await.ok().flatten()).await; + let after_reconnect = eventually(async || client.send(commands::get::("k")).await.ok().flatten()).await; let reconnected = connections(&server).await; @@ -157,7 +158,7 @@ async fn sentinel_writes_follow_the_master_through_a_failover() { let sentinel = RedisSentinel::start().await.unwrap(); let client = sentinel_client(&sentinel); client - .set("before", "1", Duration::from_secs(600)) + .send(commands::set("before", "1", Duration::from_secs(600))) .await .unwrap(); eventually(async || { @@ -176,7 +177,7 @@ async fn sentinel_writes_follow_the_master_through_a_failover() { let promoted = sentinel.failover().await.unwrap(); eventually(async || { client - .set("after", "2", Duration::from_secs(600)) + .send(commands::set("after", "2", Duration::from_secs(600))) .await .ok() }) @@ -184,7 +185,7 @@ async fn sentinel_writes_follow_the_master_through_a_failover() { assert_eq!(promoted, sentinel.replica_port()); assert_eq!( - client.get::("before").await.unwrap().as_deref(), + client.send(commands::get::("before")).await.unwrap().as_deref(), Some("1") ); let promoted = redis::Client::open(format!("redis://127.0.0.1:{promoted}")).unwrap(); diff --git a/litellm-rust/crates/redis-client/tests/pipeline.rs b/litellm-rust/crates/redis-client/tests/pipeline.rs index 89d30d35d50..06b531fc24d 100644 --- a/litellm-rust/crates/redis-client/tests/pipeline.rs +++ b/litellm-rust/crates/redis-client/tests/pipeline.rs @@ -3,10 +3,11 @@ mod support; use std::time::{Duration, Instant}; use litellm_redis_client::Value; +use litellm_redis_client::commands; use litellm_redis_testing::{RedisDeployment, Topology}; use rstest::rstest; -use litellm_redis_client::{Amount, DEFAULT_TIMEOUT, Error, Script}; +use litellm_redis_client::{Amount, DEFAULT_TIMEOUT, Error, Expiry, Script}; use support::connect; async fn begin_migration(cluster: &litellm_redis_testing::RedisCluster, key: &str) { @@ -90,7 +91,7 @@ async fn a_retrying_neighbor_cannot_reapply_a_cluster_counter() { let client = connect(&server, "", DEFAULT_TIMEOUT); let counter = "{migration}:counter"; client - .set(counter, 0, Duration::from_secs(60)) + .send(commands::set(counter, 0, Duration::from_secs(60))) .await .unwrap(); let RedisDeployment::Cluster(cluster) = &server else { @@ -103,16 +104,16 @@ async fn a_retrying_neighbor_cannot_reapply_a_cluster_counter() { .run( &client .pipeline() - .incr_by_float(counter, Amount::Int(1), Some(Duration::from_secs(60))) - .eval( + .add(commands::incr_by_float(counter, Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60))))) + .add(commands::eval::( &both_keys, &[counter, "{migration}:missing"], Vec::::new(), - ), + )), ) .await; - assert_eq!(client.get::(counter).await.unwrap(), Some(1)); + assert_eq!(client.send(commands::get::(counter)).await.unwrap(), Some(1)); assert_eq!( results[0].as_ref().unwrap(), &Value::BulkString(b"1".to_vec()) @@ -132,18 +133,18 @@ async fn each_op_settles_on_its_own_in_input_order( let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", DEFAULT_TIMEOUT); client - .set("text", "not a number", Duration::from_secs(60)) + .send(commands::set("text", "not a number", Duration::from_secs(60))) .await .unwrap(); let pipeline = client .pipeline() - .incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60))) - .incr_by_float("text", Amount::Int(1), Some(Duration::from_secs(60))) - .incr_by_float("c", Amount::Int(5), Some(Duration::from_secs(60))) - .set("b", "2", Duration::from_secs(60)) - .get("b") - .incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60))) - .get("missing"); + .add(commands::incr_by_float("a", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60))))) + .add(commands::incr_by_float("text", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60))))) + .add(commands::incr_by_float("c", Amount::Int(5), Some(Expiry::Refresh(Duration::from_secs(60))))) + .add(commands::set("b", "2", Duration::from_secs(60))) + .add(commands::get::("b")) + .add(commands::incr_by_float("a", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60))))) + .add(commands::get::("missing")); let results = client.run(&pipeline).await; @@ -184,7 +185,7 @@ async fn a_lost_script_reruns_after_the_rest_of_its_pipeline_like_redis_py( let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); let client = connect(&server, "ns", DEFAULT_TIMEOUT); let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.eval::(&incr, &["warmup"], 1).await.unwrap(); + client.send(commands::eval::(&incr, &["warmup"], 1)).await.unwrap(); if !cached { server .cli_on_every_node(&["SCRIPT", "FLUSH"]) @@ -193,14 +194,14 @@ async fn a_lost_script_reruns_after_the_rest_of_its_pipeline_like_redis_py( } let pipeline = client .pipeline() - .eval(&incr, &["counter"], 2) - .get("counter"); + .add(commands::eval::(&incr, &["counter"], 2)) + .add(commands::get::("counter")); let results = client.run(&pipeline).await; assert_eq!(results[0].as_ref().unwrap(), &Value::Int(2)); assert_eq!(results[1].as_ref().unwrap(), &read_in_the_same_pipeline); - assert_eq!(client.get::("counter").await.unwrap(), Some(2)); + assert_eq!(client.send(commands::get::("counter")).await.unwrap(), Some(2)); } #[rstest] @@ -212,8 +213,8 @@ async fn an_unreachable_server_fails_every_op_as_transport( let client = connect(&server, "", Duration::from_millis(500)); let pipeline = client .pipeline() - .incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60))) - .incr_by_float("c", Amount::Int(1), Some(Duration::from_secs(60))); + .add(commands::incr_by_float("a", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60))))) + .add(commands::incr_by_float("c", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60))))); server.freeze().await.unwrap(); let results = client.run(&pipeline).await; @@ -271,20 +272,20 @@ async fn a_pipeline_sends_scripts_by_sha_and_reruns_only_the_ones_the_server_los .run( &client .pipeline() - .incr_by_float( + .add(commands::incr_by_float( "{a}spend", Amount::Float(0.5), - Some(Duration::from_secs(60)), - ) - .eval(&count, &["{a}"], Vec::::new()) - .eval(&count, &["b"], Vec::::new()), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .add(commands::eval::(&count, &["{a}"], Vec::::new())) + .add(commands::eval::(&count, &["b"], Vec::::new())), ) .await; let second = client .run( &client .pipeline() - .eval(&count, &["{a}"], Vec::::new()), + .add(commands::eval::(&count, &["{a}"], Vec::::new())), ) .await; @@ -302,7 +303,7 @@ async fn a_pipeline_sends_scripts_by_sha_and_reruns_only_the_ones_the_server_los Value::Int(2) ] ); - assert_eq!(client.get::("{a}spend").await.unwrap(), Some(0.5)); + assert_eq!(client.send(commands::get::("{a}spend")).await.unwrap(), Some(0.5)); let stats = support::command_stats(&deployment).await; assert!(support::calls(&stats, "evalsha") > 0, "{stats}"); assert_eq!(support::calls(&stats, "eval"), 0, "{stats}"); diff --git a/litellm-rust/crates/redis-client/tests/pubsub.rs b/litellm-rust/crates/redis-client/tests/pubsub.rs index c899bfbc4bf..caf6cfea1d0 100644 --- a/litellm-rust/crates/redis-client/tests/pubsub.rs +++ b/litellm-rust/crates/redis-client/tests/pubsub.rs @@ -8,6 +8,7 @@ use litellm_redis_client::{ Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, Target, Timeouts, }; +use litellm_redis_client::commands; use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually}; use rstest::rstest; use support::{connect, connect_guarded}; @@ -35,7 +36,7 @@ async fn delivered( payload: &[u8], ) { eventually(async || { - publisher.publish(channel, payload).await.ok()?; + publisher.send(commands::publish(channel, payload)).await.ok()?; let deadline = tokio::time::Instant::now() + Duration::from_millis(200); while let Ok(Some(message)) = tokio::time::timeout_at(deadline, received.recv()).await { if message == payload { @@ -147,7 +148,7 @@ async fn an_open_breaker_rejects_a_publish_without_sending_it( server.freeze().await.unwrap(); let opened = [client.ping().await, client.ping().await]; let started = Instant::now(); - let rejected = client.publish(CHANNEL, b"payload").await; + let rejected = client.send(commands::publish(CHANNEL, b"payload")).await; let rejected_in = started.elapsed(); server.thaw().await.unwrap(); diff --git a/litellm-rust/crates/redis-client/tests/script.rs b/litellm-rust/crates/redis-client/tests/script.rs index 23c637989bf..2b59b6beaf1 100644 --- a/litellm-rust/crates/redis-client/tests/script.rs +++ b/litellm-rust/crates/redis-client/tests/script.rs @@ -6,6 +6,7 @@ use litellm_redis_testing::{RedisDeployment, Topology}; use rstest::rstest; use litellm_redis_client::Script; +use litellm_redis_client::commands; use litellm_redis_client::{DEFAULT_TIMEOUT, Error, Ttl}; use support::connect; @@ -18,10 +19,10 @@ async fn script_keys_are_namespaced( let client = connect(&server, "a", DEFAULT_TIMEOUT); let bare = connect(&server, "", DEFAULT_TIMEOUT); let read = Script::new("return redis.call('GET', KEYS[1])"); - bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap(); + bare.send(commands::set("a:x", "1", Duration::from_secs(60))).await.unwrap(); let value: Option = client - .eval(&read, &["x"], Vec::::new()) + .send(commands::eval(&read, &["x"], Vec::::new())) .await .unwrap(); @@ -37,14 +38,14 @@ async fn a_failing_script_is_a_server_error_and_keeps_the_connection( let client = connect(&server, "", DEFAULT_TIMEOUT); let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); client - .set("list", "item", Duration::from_secs(60)) + .send(commands::set("list", "item", Duration::from_secs(60))) .await .unwrap(); - let failed = client.eval::(&incr, &["list"], 1).await; + let failed = client.send(commands::eval::(&incr, &["list"], 1)).await; assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}"); - assert_eq!(client.eval::(&incr, &["counter"], 3).await.unwrap(), 3); + assert_eq!(client.send(commands::eval::(&incr, &["counter"], 3)).await.unwrap(), 3); } #[rstest] @@ -55,14 +56,14 @@ async fn a_flushed_script_is_loaded_again( let server = RedisDeployment::start(topology).await.unwrap(); let client = connect(&server, "", DEFAULT_TIMEOUT); let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.eval::(&incr, &["counter"], 1).await.unwrap(); + client.send(commands::eval::(&incr, &["counter"], 1)).await.unwrap(); server .cli_on_every_node(&["SCRIPT", "FLUSH"]) .await .unwrap(); - assert_eq!(client.eval::(&incr, &["counter"], 1).await.unwrap(), 2); + assert_eq!(client.send(commands::eval::(&incr, &["counter"], 1)).await.unwrap(), 2); } const INCREMENT_WITH_FLOOR: &str = "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) \ @@ -86,21 +87,21 @@ async fn pythons_increment_with_floor_script_floors_at_zero_and_expires_in_one_c let count = match pipelined { false => client - .eval::(&floor, &["slots"], (-5, 60)) + .send(commands::eval::(&floor, &["slots"], (-5, 60))) .await .unwrap(), true => { let results = client - .run(&client.pipeline().eval(&floor, &["slots"], (-5, 60))) + .run(&client.pipeline().add(commands::eval::(&floor, &["slots"], (-5, 60)))) .await; redis::from_redis_value(results[0].as_ref().unwrap().clone()).unwrap() } }; assert_eq!(count, 0); - assert_eq!(client.get::("slots").await.unwrap(), Some(0)); + assert_eq!(client.send(commands::get::("slots")).await.unwrap(), Some(0)); assert!(matches!( - client.ttl("slots").await.unwrap(), + client.send(commands::ttl("slots")).await.unwrap(), Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) )); } diff --git a/litellm-rust/crates/redis-connect/src/refresh.rs b/litellm-rust/crates/redis-connect/src/refresh.rs index 7d50e947d74..5014eac5fe3 100644 --- a/litellm-rust/crates/redis-connect/src/refresh.rs +++ b/litellm-rust/crates/redis-connect/src/refresh.rs @@ -236,6 +236,7 @@ fn refresh_after(expires_at: Option, now: SystemTime) -> Duration { #[cfg(test)] mod tests { + use litellm_redis_client::commands; use std::{ collections::BTreeSet, sync::{ @@ -636,7 +637,7 @@ mod tests { health_check: None, }) .unwrap(); - client.set("k", "1", Duration::from_secs(60)).await.unwrap(); + client.send(commands::set("k", "1", Duration::from_secs(60))).await.unwrap(); let before = server.clients().await.unwrap(); let after = eventually(async || { @@ -654,14 +655,14 @@ mod tests { assert_eq!(connection_users(&before), BTreeSet::from(["alice"])); assert_eq!(connection_ids(&after), connection_ids(&before)); assert_eq!( - client.get::("k").await.unwrap().as_deref(), + client.send(commands::get::("k")).await.unwrap().as_deref(), Some("1") ); server .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) .await .unwrap(); - let value = eventually(async || client.get::("k").await.ok().flatten()).await; + let value = eventually(async || client.send(commands::get::("k")).await.ok().flatten()).await; let reconnected = server.clients().await.unwrap(); assert_eq!(value, "1"); @@ -715,7 +716,7 @@ mod tests { }) .unwrap(); client - .set("before", "1", Duration::from_secs(60)) + .send(commands::set("before", "1", Duration::from_secs(60))) .await .unwrap(); let master = sentinel.master_port().await.unwrap(); @@ -737,7 +738,7 @@ mod tests { .unwrap(); let promoted = sentinel.failover().await.unwrap(); - eventually(async || client.set("after", "2", Duration::from_secs(60)).await.ok()).await; + eventually(async || client.send(commands::set("after", "2", Duration::from_secs(60))).await.ok()).await; let on_promoted = sentinel.clients_on(promoted).await.unwrap(); assert_eq!(promoted, sentinel.replica_port()); @@ -747,7 +748,7 @@ mod tests { "{on_promoted:?}" ); assert_eq!( - client.get::("after").await.unwrap().as_deref(), + client.send(commands::get::("after")).await.unwrap().as_deref(), Some("2") ); } diff --git a/litellm-rust/crates/redis-connect/tests/connect.rs b/litellm-rust/crates/redis-connect/tests/connect.rs index 30754a1912c..e38b7a4b5c5 100644 --- a/litellm-rust/crates/redis-connect/tests/connect.rs +++ b/litellm-rust/crates/redis-connect/tests/connect.rs @@ -1,3 +1,4 @@ +use litellm_redis_client::commands; use std::{collections::BTreeMap, sync::Arc, time::Duration}; use litellm_redis_connect::{ @@ -122,14 +123,14 @@ async fn python_sentinel_kwargs_follow_the_master_through_a_failover() { }); let client = connect(resolved(kwargs), services(), None).unwrap(); client - .set("before", "1", Duration::from_secs(600)) + .send(commands::set("before", "1", Duration::from_secs(600))) .await .unwrap(); let promoted = sentinel.failover().await.unwrap(); eventually(async || { client - .set("after", "2", Duration::from_secs(600)) + .send(commands::set("after", "2", Duration::from_secs(600))) .await .ok() }) @@ -138,7 +139,7 @@ async fn python_sentinel_kwargs_follow_the_master_through_a_failover() { assert_eq!(promoted, sentinel.replica_port()); assert_eq!(sentinel.master_port().await.unwrap(), promoted); assert_eq!( - client.get::("after").await.unwrap().as_deref(), + client.send(commands::get::("after")).await.unwrap().as_deref(), Some("2") ); } diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md index 9f628485eea..8a377cf5f44 100644 --- a/litellm-rust/crates/redis-store/AGENTS.md +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -1,9 +1,13 @@ -- One module per Redis family (`router_cooldowns.rs`, `router_usage.rs`, `rate_limits.rs`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys +- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (Python-facing values with their `FromValue` / `IntoValue`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` with its attach-once slot, the `DualCache` read-through and chunked deletes (`remote.rs`) and the Python codecs (`codec.rs`, `value.rs`) +- An op declared on a `Send` returns `Pending`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on +- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s - 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` +- 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 `Remote` (or none), never a cache, and a family owns one `RemoteSlot` that its memory partitions and `MemoryTier` reads share. Redis-only families are plain functions over `&RedisClient` +- Every family works with no Redis, answering from memory. `attach` adds the Redis tier at most once and keeps memory, as `DualCache.attach_redis_cache` does, so a host builds the store at startup and never rebuilds it to connect - 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 +- Key suffixes and wire spellings derive from strum on the family's enums, and a type whose Python side is plain data implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`, so the bridge never names a field or a spelling - Wire contracts are hand-written cases with expected values, inlined as the same literals in this crate's test and in `tests/unit/redis/stores/`, checked here against real Redis and there against Python's own codec. Change both sides together - `rate_limits.rs` keeps the v3 limiter's Lua byte for byte (`rate_limits/lua.rs`) so pods on either stack share one script SHA; the admission decisions stay in the Python limiter -- `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`) next to its direct method, which runs the same declaration on a batch of its own. A direct script call goes out alone instead, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here +- `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. Each op exists once and takes a `Send` (`Now`, `PreCall`, `PostCall`) picked by the caller, and the store calls `send.batch(&client)`. A store never encodes a phase in a method name. A post-call method exists only where a follow-up runs after the write settles (`token_increment_post_call`), as a thin wrapper over the shared op. `Send::Now` is an alone batch, so a script call outside a request goes out alone, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here - Every store sends the commands main's `RedisCache` sends for that operation: `MGET` where Python reads several keys at once, `SET ... EX`, and `INCRBYFLOAT` with `EXPIRE` or `TTL`. A `MONITOR` diff against main is a parity check, and any difference is a defect - 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/src/auth_cache.rs b/litellm-rust/crates/redis-store/src/auth_cache.rs deleted file mode 100644 index 07e8348dcf7..00000000000 --- a/litellm-rust/crates/redis-store/src/auth_cache.rs +++ /dev/null @@ -1,474 +0,0 @@ -use std::time::Duration; - -use litellm_cache::Error as MemoryError; -use litellm_cache_memory::InMemoryCache; -use litellm_python_compat::{Value, json, repr::to_str}; -use litellm_redis_client::{Declared, RedisBatch, RedisClient}; - -use crate::tier::{Clock, system_clock}; -use crate::{DecodeError, ErrorKind, Number, RequestRedis, codec::python_value}; - -/// `InMemoryCache`'s TTL for a write that names none and has no `default_in_memory_ttl`. -pub const MEMORY_TTL: Duration = Duration::from_secs(600); -/// `BaseCache.default_ttl`, which `RedisCache` keeps when `litellm.default_redis_ttl` is unset. -pub const REDIS_TTL: Duration = Duration::from_secs(60); -pub const MEMORY_ENTRIES: usize = 200; - -/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it. -const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"; -/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, the chunk `DualCache.async_delete_cache_keys` deletes at once. -const DELETE_CHUNK: usize = 1000; - -/// `UserApiKeyCache`'s configuration: each memory partition's size, `default_in_memory_ttl`, and the -/// `RedisCache.default_ttl` of the Redis it was attached to. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Settings { - pub memory_entries: usize, - pub memory_ttl: Option, - pub redis_ttl: Duration, -} - -impl Default for Settings { - fn default() -> Self { - Self { - memory_entries: MEMORY_ENTRIES, - memory_ttl: None, - redis_ttl: REDIS_TTL, - } - } -} - -#[derive(Debug, thiserror::Error)] -pub enum Error { - #[error(transparent)] - Redis(#[from] litellm_redis_client::Error), - #[error(transparent)] - Memory(#[from] MemoryError), - #[error(transparent)] - Decode(#[from] DecodeError), - #[error("the value is not JSON serializable: {0}")] - Encode(#[from] litellm_python_compat::Error), -} - -impl Error { - pub fn kind(&self) -> ErrorKind { - match self { - Self::Redis(error) => error.into(), - Self::Decode(_) => ErrorKind::Decode, - Self::Encode(_) => ErrorKind::Encode, - Self::Memory(_) => ErrorKind::Memory, - } - } -} - -/// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every -/// other object under a prefixed key. -pub fn is_key_object(key: &str) -> bool { - key.len() == 64 - && key - .bytes() - .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) -} - -/// `BaseCache.get_ttl` sends `int(ttl)`, so Redis sees whole seconds. -fn whole_seconds(ttl: Duration) -> Duration { - Duration::from_secs(ttl.as_secs()) -} - -fn decode(raw: Option>) -> Result, DecodeError> { - raw.as_deref().map(python_value).transpose() -} - -/// The proxy's auth object cache with `enable_redis_auth_cache` on: two memory partitions, so churn -/// in management objects cannot evict user-key objects, in front of one Redis shared by every pod. -/// Keys are the logical cache keys `UserApiKeyCache` callers already use; values are what -/// `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`. -pub struct AuthCache { - key_objects: InMemoryCache, - objects: InMemoryCache, - remote: Option, - settings: Settings, -} - -impl AuthCache { - pub fn new(client: Option, settings: Settings) -> Self { - Self::with_clock(client, settings, system_clock()) - } - - pub fn with_clock(client: Option, settings: Settings, now: Clock) -> Self { - let memory = || { - let now = now.clone(); - InMemoryCache::with_clock(Some(settings.memory_entries), Some(MEMORY_TTL), move || { - now() - }) - }; - Self { - key_objects: memory(), - objects: memory(), - remote: client, - settings, - } - } - - pub fn has_redis(&self) -> bool { - self.remote.is_some() - } - - pub fn settings(&self) -> Settings { - self.settings - } - - fn memory(&self, key: &str) -> &InMemoryCache { - if is_key_object(key) { - &self.key_objects - } else { - &self.objects - } - } - - /// `DualCache`'s rule: a write without a TTL takes `default_in_memory_ttl` in memory and Redis. - fn memory_ttl(&self, ttl: Option) -> Option { - ttl.or(self.settings.memory_ttl) - } - - fn redis_ttl(&self, ttl: Option) -> Duration { - whole_seconds(self.memory_ttl(ttl).unwrap_or(self.settings.redis_ttl)) - } - - pub fn get_local(&self, key: &str) -> Result, MemoryError> { - self.memory(key).get_cache(key) - } - - /// A memory write with `ttl`, or `default_in_memory_ttl` when it names none. - pub fn set_local( - &self, - key: &str, - value: Value, - ttl: Option, - ) -> Result<(), MemoryError> { - self.memory(key) - .set_cache(key, value, self.memory_ttl(ttl)) - .map(drop) - } - - pub fn delete_local(&self, key: &str) -> Result<(), MemoryError> { - self.memory(key).delete_cache(key) - } - - pub fn flush_local(&self) -> Result<(), MemoryError> { - self.key_objects.flush_cache()?; - self.objects.flush_cache() - } - - /// Entries held by the management-object and the user-key partitions. - pub fn len_local(&self) -> Result<(usize, usize), MemoryError> { - Ok((self.objects.len()?, self.key_objects.len()?)) - } - - /// `DualCache.async_get_cache`: memory, then Redis, backfilling memory. A failed Redis read is a - /// miss, as `DualCache` logs and swallows it. - pub async fn get(&self, key: &str) -> Result, MemoryError> { - if let Some(value) = self.get_local(key)? { - return Ok(Some(value)); - } - let Some(remote) = &self.remote else { - return Ok(None); - }; - let read = remote - .get::>(key) - .await - .map_err(Error::from) - .and_then(|raw| Ok(decode(raw)?)); - match read { - Ok(Some(value)) => { - self.set_local(key, value.clone(), None)?; - Ok(Some(value)) - } - Ok(None) => Ok(None), - Err(error) => { - tracing::warn!(%error, "auth cache Redis read failed, treating it as a miss"); - Ok(None) - } - } - } - - /// `DualCache.async_set_cache`: memory, then `SET` of `json.dumps(value)`. Redis failures and - /// values `json.dumps` refuses only log, as there. - pub async fn set( - &self, - key: &str, - value: Value, - ttl: Option, - ) -> Result<(), MemoryError> { - let encoded = json::dumps(&value); - self.set_local(key, value, ttl)?; - let Some(remote) = &self.remote else { - return Ok(()); - }; - let written = match encoded { - Ok(bytes) => remote - .set(key, bytes, self.redis_ttl(ttl)) - .await - .map_err(Error::from), - Err(error) => Err(error.into()), - }; - if let Err(error) = written { - tracing::warn!(%error, "writing an auth cache entry to Redis failed"); - } - Ok(()) - } - - /// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`. - pub async fn set_str( - &self, - key: &str, - value: Value, - ttl: Option, - ) -> Result<(), MemoryError> { - let encoded = to_str(&value); - self.set_local(key, value, ttl)?; - let Some(remote) = &self.remote else { - return Ok(()); - }; - if let Err(error) = remote.set(key, encoded, self.redis_ttl(ttl)).await { - tracing::warn!(%error, "writing an auth cache entry to Redis failed"); - } - Ok(()) - } - - /// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the request's pre-call batch, - /// sent with whatever read goes out next. `None` without Redis. - pub fn set_pre_call( - &self, - request: &RequestRedis, - key: &str, - value: Value, - ttl: Option, - ) -> Result, MemoryError> { - let encoded = json::dumps(&value); - let Some(remote) = &self.remote else { - return Ok(None); - }; - self.set_local(key, value, ttl)?; - Ok(Some(match encoded { - Ok(bytes) => { - DeclaredWrite::Set(request.pre(remote).set(key, bytes, self.redis_ttl(ttl))) - } - Err(error) => DeclaredWrite::Unencodable(error), - })) - } - - /// `DualCache.async_delete_cache_pre_call`. `None` without Redis. - pub fn delete_pre_call( - &self, - request: &RequestRedis, - key: &str, - ) -> Result, MemoryError> { - let Some(remote) = &self.remote else { - return Ok(None); - }; - self.delete_local(key)?; - Ok(Some(DeclaredWrite::Delete(request.pre(remote).del(key)))) - } - - /// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then - /// the rest in another, each with one TTL. - pub async fn set_many( - &self, - entries: Vec<(String, Value)>, - ttl: Option, - ) -> Result<(), MemoryError> { - let (key_objects, others): (Vec<_>, Vec<_>) = - entries.into_iter().partition(|(key, _)| is_key_object(key)); - for partition in [key_objects, others] { - if !partition.is_empty() { - self.set_partition(partition, ttl).await?; - } - } - Ok(()) - } - - async fn set_partition( - &self, - entries: Vec<(String, Value)>, - ttl: Option, - ) -> Result<(), MemoryError> { - let mut pipeline = self.remote.as_ref().map(RedisClient::pipeline); - for (key, value) in entries { - if let Some(writes) = pipeline.take() { - pipeline = Some(match json::dumps(&value) { - Ok(bytes) => writes.set(&key, bytes, self.redis_ttl(ttl)), - Err(error) => { - tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only"); - writes - } - }); - } - self.set_local(&key, value, ttl)?; - } - if let (Some(remote), Some(pipeline)) = (&self.remote, pipeline) { - for reply in remote.run(&pipeline).await { - if let Err(error) = reply { - tracing::warn!(%error, "writing an auth cache entry to Redis failed"); - } - } - } - Ok(()) - } - - /// The auth prefetch write-back, `async_set_cache_pipeline_with_ttls`: an entry without a TTL - /// takes `default_in_memory_ttl` in memory but the connection's default in Redis. On the - /// request's pre-call batch when there is one, else in a round trip of its own. - pub async fn write_back( - &self, - request: Option<&RequestRedis>, - entries: Vec<(String, Value, Option)>, - ) -> Result<(), Error> { - let Some(remote) = &self.remote else { - for (key, value, ttl) in entries { - self.set_local(&key, value, ttl)?; - } - return Ok(()); - }; - if let Some(request) = request { - let batch = request.pre(remote); - for (key, value, ttl) in entries { - let redis_ttl = whole_seconds(ttl.unwrap_or(self.settings.redis_ttl)); - drop(batch.set(&key, json::dumps(&value)?, redis_ttl)); - self.set_local(&key, value, ttl)?; - } - return Ok(()); - } - let mut pipeline = remote.pipeline(); - for (key, value, ttl) in entries { - let redis_ttl = whole_seconds(ttl.unwrap_or(self.settings.redis_ttl)); - pipeline = pipeline.set(&key, json::dumps(&value)?, redis_ttl); - self.set_local(&key, value, ttl)?; - } - for reply in remote.run(&pipeline).await { - reply?; - } - Ok(()) - } - - /// `UserApiKeyCache.async_delete_cache_keys`: memory, then the key-object partition and the - /// rest as two concurrent `DualCache.async_delete_cache_keys`, each a multi-key `DEL` per - /// chunk of 1000. Both run even when one fails. - pub async fn delete(&self, keys: &[String]) -> Result<(), Error> { - for key in keys { - self.delete_local(key)?; - } - let Some(remote) = &self.remote else { - return Ok(()); - }; - let (key_objects, others): (Vec<&str>, Vec<&str>) = keys - .iter() - .map(String::as_str) - .partition(|key| is_key_object(key)); - let (key_objects, others) = tokio::join!( - delete_chunks(remote, &key_objects), - delete_chunks(remote, &others) - ); - key_objects.and(others) - } - - /// Reads for keys the caller found missing in memory, on the request's pre-call batch. `None` - /// without Redis. - pub fn declare_get_many( - &self, - request: &RequestRedis, - keys: &[String], - ) -> Option { - Some(self.declare_gets(&request.pre(self.remote.as_ref()?), keys)) - } - - /// `RedisCache.async_batch_get_cache` for keys missing in memory, in one round trip. `None` - /// without Redis. - pub async fn get_many(&self, keys: &[String]) -> Option>, Error>> { - Some( - self.declare_gets(&self.remote.as_ref()?.batch(), keys) - .resolve() - .await, - ) - } - - fn declare_gets(&self, batch: &RedisBatch, keys: &[String]) -> DeclaredGets { - let keys: Vec<&str> = keys.iter().map(String::as_str).collect(); - DeclaredGets(batch.mget::>(&keys)) - } - - /// `RedisCache.async_increment_if_present`: `INCRBYFLOAT` only an entry Redis still holds. - /// `None` without Redis or for an absent entry. - pub async fn increment_if_present( - &self, - key: &str, - amount: Number, - ) -> Result, Error> { - let Some(remote) = &self.remote else { - return Ok(None); - }; - let total: Option> = remote - .eval_source( - INCREMENT_IF_PRESENT, - &[key], - litellm_redis_client::Amount::from(amount), - ) - .await?; - total - .as_deref() - .map(|raw| { - std::str::from_utf8(raw) - .ok() - .and_then(|text| text.trim().parse::().ok()) - .ok_or(Error::Decode(DecodeError::Shape( - "INCRBYFLOAT returned a non-number", - ))) - }) - .transpose() - } -} - -/// A write declared on a request's pre-call batch, settled once the batch goes out. -pub enum DeclaredWrite { - Set(Declared<()>), - Delete(Declared), - /// Kept in memory only, as `json.dumps` refused the value. - Unencodable(litellm_python_compat::Error), -} - -impl DeclaredWrite { - /// Flushes the batch if nobody has yet. - pub async fn resolve(self) -> Result<(), Error> { - match self { - Self::Set(written) => written.resolve().await?, - Self::Delete(deleted) => drop(deleted.resolve().await?), - Self::Unencodable(error) => { - tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only"); - } - } - Ok(()) - } -} - -/// Reads declared on a batch, answered once it goes out. -pub struct DeclaredGets(Declared>>>); - -impl DeclaredGets { - /// Flushes the batch if nobody has yet. Any failed read or undecodable value fails the whole - /// read, as a failed `MGET` does. - pub async fn resolve(self) -> Result>, Error> { - self.0 - .resolve() - .await? - .into_iter() - .map(|reply| Ok(decode(reply)?)) - .collect() - } -} - -async fn delete_chunks(remote: &RedisClient, keys: &[&str]) -> Result<(), Error> { - for chunk in keys.chunks(DELETE_CHUNK) { - remote.del(chunk).await?; - } - Ok(()) -} diff --git a/litellm-rust/crates/redis-store/src/auth_cache/keys.rs b/litellm-rust/crates/redis-store/src/auth_cache/keys.rs new file mode 100644 index 00000000000..feb87a40b00 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/keys.rs @@ -0,0 +1,8 @@ +/// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every +/// other object under a prefixed key. +pub fn is_key_object(key: &str) -> bool { + key.len() == 64 + && key + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} diff --git a/litellm-rust/crates/redis-store/src/auth_cache/mod.rs b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs new file mode 100644 index 00000000000..260bb78a6f2 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs @@ -0,0 +1,292 @@ +mod keys; +mod types; + +use std::time::Duration; + +use litellm_python_compat::{Value, json, repr::to_str}; +use litellm_redis_client::{Declared, RedisBatch, commands}; + +pub use self::keys::is_key_object; +pub use self::types::Settings; +use crate::{ + Error, Number, Pending, Send, + codec::python_value, + memory::{Clock, Memory, system_clock}, + remote::{self, Remote, RemoteSlot}, + request::{pending, ready}, +}; + +/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it. +const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"; + +/// The proxy's auth object cache with `enable_redis_auth_cache` on: two memory partitions, so churn +/// in management objects cannot evict user-key objects, in front of one Redis shared by every pod. +/// Keys are the logical cache keys `UserApiKeyCache` callers already use; values are what +/// `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`. +pub struct AuthCache { + key_objects: Memory, + objects: Memory, + remote: RemoteSlot, + settings: Settings, +} + +impl AuthCache { + pub fn new(remote: Option, settings: Settings) -> Self { + Self::with_clock(remote, settings, system_clock()) + } + + pub fn with_clock(remote: Option, settings: Settings, now: Clock) -> Self { + Self { + key_objects: Memory::new(Some(settings.memory_entries), &now), + objects: Memory::new(Some(settings.memory_entries), &now), + remote: RemoteSlot::new(remote), + settings, + } + } + + pub fn has_redis(&self) -> bool { + self.remote.get().is_some() + } + + pub fn settings(&self) -> Settings { + self.settings + } + + pub fn get_local(&self, key: &str) -> Result, Error> { + Ok(self.memory(key).get(key)?) + } + + /// A memory write with `ttl`, or `default_in_memory_ttl` when it names none. + pub fn set_local(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { + Ok(self.memory(key).set(key, value, self.memory_ttl(ttl))?) + } + + pub fn delete_local(&self, key: &str) -> Result<(), Error> { + Ok(self.memory(key).delete(key)?) + } + + pub fn flush_local(&self) -> Result<(), Error> { + self.key_objects.flush()?; + Ok(self.objects.flush()?) + } + + /// Entries held by the management-object and the user-key partitions. + pub fn len_local(&self) -> Result<(usize, usize), Error> { + Ok((self.objects.len()?, self.key_objects.len()?)) + } + + /// `DualCache.async_get_cache`, backfilling memory for `default_in_memory_ttl`. + pub async fn get(&self, key: &str) -> Result, Error> { + self.remote + .read_through(self.memory(key), key, python_value, self.settings.memory_ttl) + .await + } + + /// `DualCache.async_set_cache`: memory, then `SET` of `json.dumps(value)`. Redis failures and + /// values `json.dumps` refuses only log, as there. + pub async fn set(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { + let encoded = json::dumps(&value); + self.set_local(key, value, ttl)?; + let Some(remote) = self.remote.get() else { + return Ok(()); + }; + let written = match encoded { + Ok(bytes) => remote + .client + .send(commands::set(key, bytes, self.redis_ttl(remote, ttl))) + .await + .map_err(Error::from), + Err(error) => Err(error.into()), + }; + if let Err(error) = written { + tracing::warn!(%error, "writing an auth cache entry to Redis failed"); + } + Ok(()) + } + + /// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`. + pub async fn set_str(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { + let encoded = to_str(&value); + self.set_local(key, value, ttl)?; + let Some(remote) = self.remote.get() else { + return Ok(()); + }; + let write = commands::set(key, encoded, self.redis_ttl(remote, ttl)); + if let Err(error) = remote.client.send(write).await { + tracing::warn!(%error, "writing an auth cache entry to Redis failed"); + } + Ok(()) + } + + /// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the batch `send` picks. + pub fn declare_set( + &self, + send: Send<'_>, + key: &str, + value: Value, + ttl: Option, + ) -> Result, Error> { + let encoded = json::dumps(&value); + self.set_local(key, value, ttl)?; + let Some(remote) = self.remote.get() else { + return Ok(ready(())); + }; + match encoded { + Ok(bytes) => Ok(pending( + send.batch(&remote.client) + .add(commands::set(key, bytes, self.redis_ttl(remote, ttl))), + )), + Err(error) => { + tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only"); + Ok(ready(())) + } + } + } + + /// `DualCache.async_delete_cache_pre_call`, memory now and the `DEL` on the batch `send` picks. + pub fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { + self.delete_local(key)?; + let Some(batch) = self.remote.batch(send) else { + return Ok(ready(())); + }; + let deleted = batch.add(commands::del(&[key])); + Ok(Box::pin(async move { Ok(deleted.resolve().await.map(drop)?) })) + } + + /// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then + /// the rest in another, each with one TTL. + pub async fn set_many(&self, entries: Vec<(String, Value)>, ttl: Option) -> Result<(), Error> { + let (key_objects, others): (Vec<_>, Vec<_>) = + entries.into_iter().partition(|(key, _)| is_key_object(key)); + for partition in [key_objects, others] { + self.set_partition(partition, ttl).await?; + } + Ok(()) + } + + /// The auth prefetch write-back, `async_set_cache_pipeline_with_ttls`: memory now, and an entry + /// without a TTL takes `default_in_memory_ttl` in memory but the connection's default in Redis. + /// On a request batch the `SET`s are declared and nothing is left to wait for. [`Send::Now`] is + /// one pipeline of its own, as `async_set_cache_pipeline_with_ttls` sends it, rather than one + /// round trip per entry. + pub fn write_back( + &self, + send: Send<'_>, + entries: Vec<(String, Value, Option)>, + ) -> Result, Error> { + let Some(remote) = self.remote.get() else { + for (key, value, ttl) in entries { + self.set_local(&key, value, ttl)?; + } + return Ok(ready(())); + }; + let sets = entries + .into_iter() + .map(|(key, value, ttl)| { + let redis_ttl = ttl.unwrap_or(remote.default_ttl); + let set = commands::set(&key, json::dumps(&value)?, redis_ttl); + self.set_local(&key, value, ttl)?; + Ok(set) + }) + .collect::, Error>>()?; + let Send::Now = send else { + let batch = send.batch(&remote.client); + sets.into_iter().for_each(|set| drop(batch.add(set))); + return Ok(ready(())); + }; + let pipeline = RedisBatch::pipelined(remote.client.clone()); + let written: Vec<_> = sets.into_iter().map(|set| pipeline.add(set)).collect(); + Ok(Box::pin(async move { + Declared::all(written).await.into_iter().try_for_each(|reply| Ok(reply?)) + })) + } + + /// `UserApiKeyCache.async_delete_cache_keys`: memory, then the key-object partition and the + /// rest as two concurrent `DualCache.async_delete_cache_keys`, each a multi-key `DEL` per + /// chunk of 1000. Both run even when one fails. + pub async fn delete(&self, keys: &[String]) -> Result<(), Error> { + keys.iter().try_for_each(|key| self.delete_local(key))?; + let Some(client) = self.remote.client() else { + return Ok(()); + }; + let (key_objects, others): (Vec, Vec) = + keys.iter().cloned().partition(|key| is_key_object(key)); + let (key_objects, others) = tokio::join!( + remote::delete(client, &key_objects), + remote::delete(client, &others) + ); + key_objects.and(others) + } + + /// `RedisCache.async_batch_get_cache` for keys the caller found missing in memory, one `MGET` on + /// the batch `send` picks. `None` without Redis. Any failed read or undecodable value fails the + /// whole read, as a failed `MGET` does. + pub fn get_many(&self, send: Send<'_>, keys: &[String]) -> Option>>> { + let keys: Vec<&str> = keys.iter().map(String::as_str).collect(); + let read = self.remote.batch(send)?.add(commands::mget::>(&keys)); + Some(Box::pin(async move { + read.resolve() + .await? + .into_iter() + .map(|raw| Ok(types::decode(raw)?)) + .collect() + })) + } + + /// `RedisCache.async_increment_if_present`: `INCRBYFLOAT` only an entry Redis still holds. + /// `None` without Redis or for an absent entry. + pub async fn increment_if_present(&self, key: &str, amount: Number) -> Result, Error> { + let Some(client) = self.remote.client() else { + return Ok(None); + }; + let amount = litellm_redis_client::Amount::from(amount); + Ok(client.send(commands::eval_source(INCREMENT_IF_PRESENT, &[key], amount)).await?) + } + + fn memory(&self, key: &str) -> &Memory { + if is_key_object(key) { + &self.key_objects + } else { + &self.objects + } + } + + /// `DualCache`'s rule: a write without a TTL takes `default_in_memory_ttl` in memory and Redis. + fn memory_ttl(&self, ttl: Option) -> Option { + ttl.or(self.settings.memory_ttl) + } + + fn redis_ttl(&self, remote: &Remote, ttl: Option) -> Duration { + self.memory_ttl(ttl).unwrap_or(remote.default_ttl) + } + + async fn set_partition(&self, entries: Vec<(String, Value)>, ttl: Option) -> Result<(), Error> { + let target = self + .remote + .get() + .map(|remote| (RedisBatch::pipelined(remote.client.clone()), self.redis_ttl(remote, ttl))); + let written = entries + .into_iter() + .map(|(key, value)| { + let set = match (&target, json::dumps(&value)) { + (Some((pipeline, redis_ttl)), Ok(bytes)) => { + Some(pipeline.add(commands::set(&key, bytes, *redis_ttl))) + } + (Some(_), Err(error)) => { + tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only"); + None + } + (None, _) => None, + }; + self.set_local(&key, value, ttl)?; + Ok(set) + }) + .collect::, Error>>()?; + for reply in Declared::all(written.into_iter().flatten()).await { + if let Err(error) = reply { + tracing::warn!(%error, "writing an auth cache entry to Redis failed"); + } + } + Ok(()) + } +} diff --git a/litellm-rust/crates/redis-store/src/auth_cache/types.rs b/litellm-rust/crates/redis-store/src/auth_cache/types.rs new file mode 100644 index 00000000000..d66ae858e44 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/types.rs @@ -0,0 +1,26 @@ +use std::time::Duration; + +use litellm_python_compat::Value; + +use crate::{DecodeError, codec::python_value}; + +/// `UserApiKeyCache`'s memory configuration: each partition's size and `default_in_memory_ttl`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Settings { + pub memory_entries: usize, + pub memory_ttl: Option, +} + +impl Default for Settings { + /// `InMemoryCache`'s 200 entries. + fn default() -> Self { + Self { + memory_entries: 200, + memory_ttl: None, + } + } +} + +pub(super) fn decode(raw: Option>) -> Result, DecodeError> { + raw.as_deref().map(python_value).transpose() +} diff --git a/litellm-rust/crates/redis-store/src/codec.rs b/litellm-rust/crates/redis-store/src/codec.rs index a8662fcbe0c..43029b8f337 100644 --- a/litellm-rust/crates/redis-store/src/codec.rs +++ b/litellm-rust/crates/redis-store/src/codec.rs @@ -1,5 +1,7 @@ use litellm_python_compat::{Value, json, literal}; +use crate::Number; + #[derive(Debug, thiserror::Error)] pub enum DecodeError { #[error("Redis value is not UTF-8")] @@ -18,3 +20,8 @@ pub(crate) fn python_value(raw: &[u8]) -> Result { Err(_) => literal::literal_eval(text).map_err(DecodeError::Unparsed), } } + +/// A number written by Python, `shape` naming what it was meant to be when it is not one. +pub(crate) fn number(raw: &[u8], shape: &'static str) -> Result { + Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape(shape)) +} diff --git a/litellm-rust/crates/redis-store/src/error.rs b/litellm-rust/crates/redis-store/src/error.rs index 89dd61433ad..b22e73db1ac 100644 --- a/litellm-rust/crates/redis-store/src/error.rs +++ b/litellm-rust/crates/redis-store/src/error.rs @@ -1,3 +1,28 @@ +use crate::DecodeError; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Redis(#[from] litellm_redis_client::Error), + #[error(transparent)] + Memory(#[from] litellm_cache::Error), + #[error(transparent)] + Decode(#[from] DecodeError), + #[error("the value is not JSON serializable: {0}")] + Encode(#[from] litellm_python_compat::Error), +} + +impl Error { + pub fn kind(&self) -> ErrorKind { + match self { + Self::Redis(error) => error.into(), + Self::Memory(_) => ErrorKind::Memory, + Self::Decode(_) => ErrorKind::Decode, + Self::Encode(_) => ErrorKind::Encode, + } + } +} + /// The failure class a host reports with a store error, `timeout` and `breaker_open` being the ones a /// caller may drop. #[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr)] diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs index 6d58aeb043f..16627a3c548 100644 --- a/litellm-rust/crates/redis-store/src/lib.rs +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -1,16 +1,22 @@ pub mod auth_cache; mod codec; mod error; +mod memory; mod number; pub mod rate_limits; +mod remote; mod request; pub mod router_cooldowns; pub mod router_usage; pub mod spend_counters; mod tier; +mod value; pub use codec::DecodeError; -pub use error::ErrorKind; +pub use error::{Error, ErrorKind}; +pub use memory::{Clock, system_clock}; pub use number::Number; -pub use request::RequestRedis; -pub use tier::{Clock, TierPolicy}; +pub use remote::Remote; +pub use request::{Pending, RequestRedis, Send}; +pub use tier::TierPolicy; +pub use value::{FromValue, IntoValue}; diff --git a/litellm-rust/crates/redis-store/src/memory.rs b/litellm-rust/crates/redis-store/src/memory.rs new file mode 100644 index 00000000000..c616a6fc6e8 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/memory.rs @@ -0,0 +1,84 @@ +use std::{ + sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::Error; +use litellm_cache_memory::InMemoryCache; + +pub type Clock = Arc Duration + Send + Sync>; + +pub fn system_clock() -> Clock { + Arc::new(|| { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + }) +} + +/// A store's memory tier: Python's `InMemoryCache`, with its 200 entries and 600 s TTL unless the +/// store names others, plus one lock for every read-modify-write on it. +pub struct Memory { + cache: InMemoryCache, + updates: Mutex<()>, +} + +impl Memory { + pub fn new(entries: Option, now: &Clock) -> Self { + let now = now.clone(); + Self { + cache: InMemoryCache::with_clock(entries, None, move || now()), + updates: Mutex::new(()), + } + } + + pub fn get(&self, key: &str) -> Result, Error> { + self.cache.get_cache(key) + } + + pub fn set(&self, key: &str, value: V, ttl: Option) -> Result<(), Error> { + self.cache.set_cache(key, value, ttl).map(drop) + } + + pub fn delete(&self, key: &str) -> Result<(), Error> { + self.cache.delete_cache(key) + } + + pub fn flush(&self) -> Result<(), Error> { + self.cache.flush_cache() + } + + pub fn len(&self) -> Result { + self.cache.len() + } + + pub fn expires_at(&self, key: &str) -> Result, Error> { + self.cache.expires_at(key) + } + + pub fn entries(&self) -> Result, Error> { + self.cache.entries() + } + + /// Read-modify-write of one entry, atomic against every other update of this tier. + pub fn update( + &self, + key: &str, + ttl: Option, + update: impl FnOnce(Option) -> Result, + ) -> Result { + let _guard = self.guard(); + let value = update(self.get(key)?)?; + self.set(key, value.clone(), ttl)?; + Ok(value) + } + + /// Held across a read and a conditional write that `update` cannot express. + pub fn guard(&self) -> MutexGuard<'_, ()> { + lock(&self.updates) + } +} + +pub(crate) fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} diff --git a/litellm-rust/crates/redis-store/src/number.rs b/litellm-rust/crates/redis-store/src/number.rs index 9fe67930860..7f3e9f4d10b 100644 --- a/litellm-rust/crates/redis-store/src/number.rs +++ b/litellm-rust/crates/redis-store/src/number.rs @@ -1,6 +1,6 @@ use std::{cmp::Ordering, ops::Add}; -use litellm_python_compat::{BigInt, Value}; +use litellm_python_compat::{BigInt, Value, json, repr}; /// A Python `int` or `float`, kept apart because `repr` and `json.dumps` write them differently. #[derive(Clone, Copy, Debug, PartialEq)] @@ -41,6 +41,12 @@ impl Number { } } + /// `json.dumps`, which `RedisCache.async_set_cache` writes. + pub fn to_json(self) -> String { + let value = self.to_value(); + json::dumps(&value).unwrap_or_else(|_| repr::to_str(&value)) + } + pub fn from_value(value: &Value) -> Option { match value { Value::Int(value) => i64::try_from(value).ok().map(Self::Int), diff --git a/litellm-rust/crates/redis-store/src/rate_limits.rs b/litellm-rust/crates/redis-store/src/rate_limits.rs deleted file mode 100644 index 198c0bd157e..00000000000 --- a/litellm-rust/crates/redis-store/src/rate_limits.rs +++ /dev/null @@ -1,508 +0,0 @@ -mod lua; -mod scripts; - -use std::{ - sync::{Arc, Mutex, PoisonError}, - time::Duration, -}; - -use futures_util::future::join_all; -use litellm_cache::Error as MemoryError; -use litellm_cache_memory::InMemoryCache; -use litellm_python_compat::Value; -use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient, Slot}; - -use self::scripts::Route; -pub use self::scripts::{ - AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply, - Pending, WindowCount, -}; -use crate::{Number, RequestRedis, codec::python_value}; - -/// A rate-limit descriptor, `{descriptor:value}`, the hash tag every key of its window shares. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct Scope { - pub descriptor: String, - pub value: String, -} - -impl Scope { - fn key(&self, suffix: &str) -> String { - format!("{{{}:{}}}:{suffix}", self.descriptor, self.value) - } - - fn window_key(&self) -> String { - self.key(WINDOW) - } - - fn gauge_key(&self) -> String { - self.key(GAUGE) - } -} - -const WINDOW: &str = "window"; -const GAUGE: &str = "max_parallel_requests"; - -/// A counter's key suffix, which is also how the Python limiter names it. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)] -pub enum Kind { - #[strum(serialize = "requests")] - Requests, - #[strum(serialize = "tokens")] - Tokens, -} - -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct Counter { - pub scope: Scope, - pub kind: Kind, -} - -impl Counter { - fn key(&self) -> String { - self.scope.key(self.kind.into()) - } -} - -/// One value of the limiter's keyspace, as both tiers address it. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub enum Entry { - Window(Scope), - Counter(Counter), - Gauge(Scope), -} - -impl Entry { - /// The entry the Python limiter names by its key suffix: `window`, `max_parallel_requests`, or a - /// counter's `Kind`. - pub fn new(suffix: &str, scope: Scope) -> Result { - match suffix { - WINDOW => Ok(Self::Window(scope)), - GAUGE => Ok(Self::Gauge(scope)), - kind => kind - .parse() - .map(|kind| Self::Counter(Counter { scope, kind })), - } - } - - fn key(&self) -> String { - match self { - Self::Window(scope) => scope.window_key(), - Self::Counter(counter) => counter.key(), - Self::Gauge(scope) => scope.gauge_key(), - } - } -} - -/// What the memory tier holds: counters and mirrored gauge counts, window starts, and the per-slot -/// registry a gauge keeps while Redis is unreachable. -#[derive(Clone, Debug, PartialEq)] -pub enum Local { - Number(Number), - Text(String), - Slots(Vec<(String, f64)>), -} - -/// A `RedisPipelineIncrementOperation`. -#[derive(Clone, Debug, PartialEq)] -pub struct Increment { - pub counter: Counter, - pub amount: Number, - pub ttl: Option, -} - -#[derive(Debug, thiserror::Error)] -pub enum Error { - #[error(transparent)] - Redis(#[from] litellm_redis_client::Error), - #[error("rate limiter script reply {0}")] - Reply(&'static str), - #[error("the rate limiter has no Redis to run its scripts on")] - NoRedis, -} - -/// The v3 rate limiter's keyspace: its memory tier, its Lua scripts and the plain counter writes -/// `DualCache` made for it on `main`. The admission logic stays with the caller. -pub struct RateLimits { - redis: Option, - memory: InMemoryCache, - increments: Mutex<()>, - backfill_ttl: Option, - redis_ttl: Duration, -} - -impl RateLimits { - /// `backfill_ttl` is the `default_in_memory_ttl` of the `DualCache` this tier replaces, and - /// `redis_ttl` its `RedisCache.default_ttl`, which an increment without a TTL takes. - pub fn new( - redis: Option, - backfill_ttl: Option, - redis_ttl: Duration, - ) -> Self { - Self::with_memory( - redis, - backfill_ttl, - redis_ttl, - InMemoryCache::new(None, None), - ) - } - - pub fn with_memory( - redis: Option, - backfill_ttl: Option, - redis_ttl: Duration, - memory: InMemoryCache, - ) -> Self { - Self { - redis, - memory, - increments: Mutex::new(()), - backfill_ttl, - redis_ttl, - } - } - - pub fn has_redis(&self) -> bool { - self.redis.is_some() - } - - /// `_group_keys_by_hash_tag`: indices of `scopes` sharing a cluster slot, in first-seen order, - /// or one group when the client is not a cluster and grouping is not forced. - pub fn slot_groups(&self, scopes: &[Scope], force: bool) -> Vec> { - let by_slot = force || self.redis.as_ref().is_some_and(RedisClient::is_cluster); - if !by_slot { - return vec![(0..scopes.len()).collect()]; - } - scopes - .iter() - .enumerate() - .fold( - Vec::<(Slot, Vec)>::new(), - |mut groups, (index, scope)| { - let slot = Slot::for_key(scope.window_key()); - match groups.iter_mut().find(|(seen, _)| *seen == slot) { - Some((_, members)) => members.push(index), - None => groups.push((slot, vec![index])), - } - groups - }, - ) - .into_iter() - .map(|(_, members)| members) - .collect() - } - - pub fn local(&self, entry: &Entry) -> Result, MemoryError> { - self.memory.get_cache(&entry.key()) - } - - pub fn set_local( - &self, - entry: &Entry, - value: Local, - ttl: Option, - ) -> Result<(), MemoryError> { - self.memory.set_cache(entry.key(), value, ttl).map(drop) - } - - /// `DualCache.async_get_cache`: memory, else Redis written back to memory for `backfill_ttl`, - /// or the store's default. A Redis failure reads as a miss, as `DualCache` logs and swallows it. - pub async fn get( - &self, - entry: &Entry, - backfill_ttl: Option, - ) -> Result, MemoryError> { - let key = entry.key(); - if let Some(local) = self.memory.get_cache(&key)? { - return Ok(Some(local)); - } - let Some(redis) = &self.redis else { - return Ok(None); - }; - let raw = match redis.get::>(&key).await { - Ok(raw) => raw, - Err(error) => { - tracing::error!(%error, "LiteLLM Cache: exception in async_get_cache"); - return Ok(None); - } - }; - let Some(value) = raw - .map(|raw| decode(&raw)) - .transpose() - .unwrap_or_else(|error| { - tracing::error!(%error, "LiteLLM Cache: exception in async_get_cache"); - None - }) - else { - return Ok(None); - }; - self.memory - .set_cache(key, value.clone(), backfill_ttl.or(self.backfill_ttl))?; - Ok(Some(value)) - } - - pub fn sliding_window( - &self, - request: Option<&RequestRedis>, - counters: &[Counter], - now: i64, - window_size: i64, - ) -> Result>, Error> { - Ok(scripts::sliding_window( - &self.pre(request)?, - counters, - now, - window_size, - )) - } - - pub fn read_counters( - &self, - request: Option<&RequestRedis>, - counters: &[Counter], - ) -> Result>>, Error> { - Ok(scripts::read_counters(&self.pre(request)?, counters)) - } - - pub fn check_and_increment( - &self, - request: Option<&RequestRedis>, - counters: &[AtomicCounter], - ) -> Result, Error> { - Ok(scripts::check_and_increment(&self.pre(request)?, counters)) - } - - pub fn guarded_increment( - &self, - request: Option<&RequestRedis>, - increment: &GuardedIncrement, - ) -> Result, Error> { - Ok(scripts::guarded_increment(&self.pre(request)?, increment)) - } - - pub fn parallel_acquire( - &self, - request: Option<&RequestRedis>, - gauges: &[Gauge], - slot_ttl: Duration, - slot_id: &str, - ) -> Result, Error> { - Ok(scripts::parallel_acquire( - &self.pre(request)?, - gauges, - slot_ttl, - slot_id, - )) - } - - pub fn parallel_release( - &self, - request: Option<&RequestRedis>, - gauges: &[Scope], - slot_id: &str, - ) -> Result>, Error> { - Ok(scripts::parallel_release( - &self.pre(request)?, - gauges, - slot_id, - )) - } - - pub fn parallel_count( - &self, - request: Option<&RequestRedis>, - gauges: &[Scope], - slot_ttl: Duration, - ) -> Result>, Error> { - Ok(scripts::parallel_count( - &self.pre(request)?, - gauges, - slot_ttl, - )) - } - - pub fn token_increment( - &self, - request: Option<&RequestRedis>, - increments: &[Increment], - ) -> Result>, Error> { - Ok(scripts::token_increment(&self.pre(request)?, increments)) - } - - /// `TOKEN_INCREMENT_SCRIPT` on the request's post-call batch. A failed script falls back to - /// `increment_many` for these increments once the batch has gone out. - pub fn token_increment_post_call( - self: &Arc, - request: &RequestRedis, - increments: Vec, - ) -> Result<(), Error> { - let pending = scripts::token_increment(&self.post(request)?, &increments); - let store = Arc::clone(self); - request.spawn(async move { - let Err(error) = pending.settled().await else { - return; - }; - tracing::warn!(%error, "TTL preservation failed, falling back to regular pipeline"); - if let Err(error) = store.increment_many(&increments).await { - tracing::warn!(%error, "in-memory increment failed"); - } - }); - Ok(()) - } - - /// `PARALLEL_RELEASE_SCRIPT` on the request's post-call batch; the caller has already freed - /// the slot in memory. - pub fn parallel_release_post_call( - &self, - request: &RequestRedis, - gauges: &[Scope], - slot_id: &str, - ) -> Result<(), Error> { - let pending = scripts::parallel_release(&self.post(request)?, gauges, slot_id); - request.spawn(async move { - if let Err(error) = pending.settled().await { - tracing::warn!( - %error, - "parallel_release_script failed, the slot stays released in memory only" - ); - } - }); - Ok(()) - } - - /// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL. A - /// Redis failure answers with the memory count. - pub async fn increment(&self, increment: &Increment) -> Result { - let local = self.add_local(increment)?; - let Some(redis) = &self.redis else { - return Ok(local); - }; - let key = increment.counter.key(); - let ttl = increment.ttl.unwrap_or(self.redis_ttl); - let total = redis - .incr_by_float(&key, increment.amount.into(), Expiry::KeepExisting(ttl)) - .await - .map_err(Error::from); - Ok(total.map(Number::Float).unwrap_or_else(|error| { - tracing::warn!( - %error, - "Redis async_increment_cache failed, falling back to in-memory result" - ); - local - })) - } - - /// `DualCache.async_increment_cache_pipeline`: memory, then one pipeline of `INCRBYFLOAT` - /// with `EXPIRE` per increment. - pub async fn increment_many(&self, increments: &[Increment]) -> Result<(), MemoryError> { - increments - .iter() - .try_for_each(|increment| self.add_local(increment).map(drop))?; - let Some(redis) = &self.redis else { - return Ok(()); - }; - let pipeline = increments - .iter() - .fold(redis.pipeline(), |pipeline, increment| { - pipeline.incr_by_float( - &increment.counter.key(), - increment.amount.into(), - increment.ttl, - ) - }); - if let Some(error) = redis.run(&pipeline).await.into_iter().find_map(Result::err) { - tracing::warn!( - %error, - "Redis async_increment_cache_pipeline failed, falling back to in-memory result" - ); - } - Ok(()) - } - - /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, the Redis increments on - /// the request's post-call batch. - pub fn increment_many_post_call( - &self, - request: &RequestRedis, - increments: &[Increment], - ) -> Result<(), MemoryError> { - let Some(declared) = self.declare_increments(increments, |client| request.post(client))? - else { - return Ok(()); - }; - request.spawn(async move { - if let Some(error) = settle(declared).await { - tracing::warn!(%error, "deferred Redis increment failed"); - } - }); - Ok(()) - } - - fn declare_increments( - &self, - increments: &[Increment], - batch_for: impl FnOnce(&RedisClient) -> RedisBatch, - ) -> Result>>, MemoryError> { - increments - .iter() - .try_for_each(|increment| self.add_local(increment).map(drop))?; - let Some(redis) = &self.redis else { - return Ok(None); - }; - let batch = batch_for(redis); - Ok(Some( - increments - .iter() - .map(|increment| scripts::counter_increment(&batch, increment)) - .collect(), - )) - } - - /// `InMemoryCache.increment_cache`: `(current or 0) + amount` under one lock. - fn add_local(&self, increment: &Increment) -> Result { - let _guard = self - .increments - .lock() - .unwrap_or_else(PoisonError::into_inner); - let key = increment.counter.key(); - let current = match self.memory.get_cache(&key)? { - None => Number::Int(0), - Some(Local::Number(current)) => current, - Some(_) => return Err(MemoryError::InvalidEntry(None)), - }; - let total = current + increment.amount; - self.memory - .set_cache(key, Local::Number(total), increment.ttl)?; - Ok(total) - } - - fn pre(&self, request: Option<&RequestRedis>) -> Result { - let redis = self.redis.as_ref().ok_or(Error::NoRedis)?; - Ok(request.map_or_else( - || Route::Direct(redis.clone()), - |request| Route::Batch(request.pre(redis)), - )) - } - - fn post(&self, request: &RequestRedis) -> Result { - let redis = self.redis.as_ref().ok_or(Error::NoRedis)?; - Ok(Route::Batch(request.post(redis))) - } -} - -/// Every increment's outcome; the first failure when there is one. `flush` sends the batch first. -async fn settle(declared: Vec>) -> Option { - let outcomes = join_all(declared.into_iter().map(Declared::settled)).await; - outcomes.into_iter().find_map(Result::err).map(Error::from) -} - -fn decode(raw: &[u8]) -> Result { - match python_value(raw)? { - Value::Str(text) => Ok(Local::Text(text)), - value => Number::from_value(&value) - .map(Local::Number) - .ok_or(crate::DecodeError::Shape( - "a rate-limit value is a number or a string", - )), - } -} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/keys.rs b/litellm-rust/crates/redis-store/src/rate_limits/keys.rs new file mode 100644 index 00000000000..7f5714a19ad --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/keys.rs @@ -0,0 +1,74 @@ +/// A rate-limit descriptor, `{descriptor:value}`, the hash tag every key of its window shares. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Scope { + pub descriptor: String, + pub value: String, +} + +impl Scope { + pub(super) fn key(&self, suffix: &str) -> String { + format!("{{{}:{}}}:{suffix}", self.descriptor, self.value) + } + + pub(super) fn window_key(&self) -> String { + self.key(WINDOW) + } + + pub(super) fn gauge_key(&self) -> String { + self.key(GAUGE) + } +} + +const WINDOW: &str = "window"; +const GAUGE: &str = "max_parallel_requests"; + +/// A counter's key suffix, which is also how the Python limiter names it. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)] +pub enum Kind { + #[strum(serialize = "requests")] + Requests, + #[strum(serialize = "tokens")] + Tokens, +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Counter { + pub scope: Scope, + pub kind: Kind, +} + +impl Counter { + pub(super) fn key(&self) -> String { + self.scope.key(self.kind.into()) + } +} + +/// One value of the limiter's keyspace, as both tiers address it. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum Entry { + Window(Scope), + Counter(Counter), + Gauge(Scope), +} + +impl Entry { + /// The entry the Python limiter names by its key suffix: `window`, `max_parallel_requests`, or a + /// counter's `Kind`. + pub fn new(suffix: &str, scope: Scope) -> Result { + match suffix { + WINDOW => Ok(Self::Window(scope)), + GAUGE => Ok(Self::Gauge(scope)), + kind => kind + .parse() + .map(|kind| Self::Counter(Counter { scope, kind })), + } + } + + pub(super) fn key(&self) -> String { + match self { + Self::Window(scope) => scope.window_key(), + Self::Counter(counter) => counter.key(), + Self::Gauge(scope) => scope.gauge_key(), + } + } +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/mod.rs b/litellm-rust/crates/redis-store/src/rate_limits/mod.rs new file mode 100644 index 00000000000..abc978d6bcf --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/mod.rs @@ -0,0 +1,259 @@ +mod keys; +mod lua; +mod scripts; +mod types; + +use std::{sync::Arc, time::Duration}; + +use futures_util::future::join_all; +use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient, Slot, commands}; + +pub use self::keys::{Counter, Entry, Kind, Scope}; +pub use self::types::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply, + Increment, Local, WindowCount, +}; +use crate::{ + Error, Number, Pending, RequestRedis, Send, + memory::{Clock, Memory, system_clock}, + remote::{Remote, RemoteSlot}, + request::pending, +}; + +/// The v3 rate limiter's keyspace: its memory tier, its Lua scripts and the plain counter writes +/// `DualCache` made for it on `main`. The admission logic stays with the caller. A script needs +/// Redis, so without one it declares nothing and answers `None`. +pub struct RateLimits { + memory: Memory, + remote: RemoteSlot, + backfill_ttl: Option, +} + +impl RateLimits { + /// `backfill_ttl` is the `default_in_memory_ttl` of the `DualCache` this tier replaces. An + /// increment without a TTL takes the remote's default. + pub fn new(remote: Option, backfill_ttl: Option) -> Self { + Self::with_clock(remote, backfill_ttl, system_clock()) + } + + pub fn with_clock(remote: Option, backfill_ttl: Option, now: Clock) -> Self { + Self { + memory: Memory::new(None, &now), + remote: RemoteSlot::new(remote), + backfill_ttl, + } + } + + pub fn has_redis(&self) -> bool { + self.remote.get().is_some() + } + + /// `_group_keys_by_hash_tag`: indices of `scopes` sharing a cluster slot, in first-seen order, + /// or one group when the client is not a cluster and grouping is not forced. + pub fn slot_groups(&self, scopes: &[Scope], force: bool) -> Vec> { + let by_slot = force || self.remote.client().is_some_and(RedisClient::is_cluster); + if !by_slot { + return vec![(0..scopes.len()).collect()]; + } + scopes + .iter() + .enumerate() + .fold(Vec::<(Slot, Vec)>::new(), |mut groups, (index, scope)| { + let slot = Slot::for_key(scope.window_key()); + match groups.iter_mut().find(|(seen, _)| *seen == slot) { + Some((_, members)) => members.push(index), + None => groups.push((slot, vec![index])), + } + groups + }) + .into_iter() + .map(|(_, members)| members) + .collect() + } + + pub fn get_local(&self, entry: &Entry) -> Result, Error> { + Ok(self.memory.get(&entry.key())?) + } + + pub fn set_local(&self, entry: &Entry, value: Local, ttl: Option) -> Result<(), Error> { + Ok(self.memory.set(&entry.key(), value, ttl)?) + } + + /// `DualCache.async_get_cache`, backfilling memory for `backfill_ttl` or the store's default. + pub async fn get(&self, entry: &Entry, backfill_ttl: Option) -> Result, Error> { + let backfill_ttl = backfill_ttl.or(self.backfill_ttl); + self.remote + .read_through(&self.memory, &entry.key(), types::decode, backfill_ttl) + .await + } + + pub fn sliding_window( + &self, + send: Send<'_>, + counters: &[Counter], + now: i64, + window_size: i64, + ) -> Option>> { + Some(pending(scripts::sliding_window(&self.batch(send)?, counters, now, window_size))) + } + + pub fn read_counters(&self, send: Send<'_>, counters: &[Counter]) -> Option>>> { + Some(pending(scripts::read_counters(&self.batch(send)?, counters))) + } + + pub fn check_and_increment(&self, send: Send<'_>, counters: &[AtomicCounter]) -> Option> { + Some(pending(scripts::check_and_increment(&self.batch(send)?, counters))) + } + + pub fn guarded_increment(&self, send: Send<'_>, increment: &GuardedIncrement) -> Option> { + Some(pending(scripts::guarded_increment(&self.batch(send)?, increment))) + } + + pub fn parallel_acquire( + &self, + send: Send<'_>, + gauges: &[Gauge], + slot_ttl: Duration, + slot_id: &str, + ) -> Option> { + Some(pending(scripts::parallel_acquire(&self.batch(send)?, gauges, slot_ttl, slot_id))) + } + + pub fn parallel_release(&self, send: Send<'_>, gauges: &[Scope], slot_id: &str) -> Option>> { + Some(pending(self.declare_release(send, gauges, slot_id)?)) + } + + pub fn parallel_count(&self, send: Send<'_>, gauges: &[Scope], slot_ttl: Duration) -> Option>> { + Some(pending(scripts::parallel_count(&self.batch(send)?, gauges, slot_ttl))) + } + + pub fn token_increment(&self, send: Send<'_>, increments: &[Increment]) -> Option>> { + Some(pending(self.declare_token_increment(send, increments)?)) + } + + /// `TOKEN_INCREMENT_SCRIPT` on the request's post-call batch. A failed script falls back to + /// `increment_many` for these increments once the batch has gone out. `false` without Redis. + pub fn token_increment_post_call(self: &Arc, request: &RequestRedis, increments: Vec) -> bool { + let Some(declared) = self.declare_token_increment(Send::PostCall(request), &increments) else { + return false; + }; + let store = Arc::clone(self); + request.spawn(async move { + let Err(error) = declared.settled().await else { + return; + }; + tracing::warn!(%error, "TTL preservation failed, falling back to regular pipeline"); + if let Err(error) = store.increment_many(&increments).await { + tracing::warn!(%error, "in-memory increment failed"); + } + }); + true + } + + /// `PARALLEL_RELEASE_SCRIPT` on the request's post-call batch; the caller has already freed + /// the slot in memory. `false` without Redis. + pub fn parallel_release_post_call(&self, request: &RequestRedis, gauges: &[Scope], slot_id: &str) -> bool { + let Some(declared) = self.declare_release(Send::PostCall(request), gauges, slot_id) else { + return false; + }; + request.spawn(async move { + if let Err(error) = declared.settled().await { + tracing::warn!(%error, "parallel_release_script failed, the slot stays released in memory only"); + } + }); + true + } + + /// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL. A + /// Redis failure answers with the memory count. + pub async fn increment(&self, increment: &Increment) -> Result { + let local = self.add_local(increment)?; + let Some(remote) = self.remote.get() else { + return Ok(local); + }; + let ttl = increment.ttl.unwrap_or(remote.default_ttl); + let total = remote + .client + .send(commands::incr_by_float( + &increment.counter.key(), + increment.amount.into(), + Some(Expiry::KeepExisting(ttl)), + )) + .await; + Ok(total.map(Number::Float).unwrap_or_else(|error| { + tracing::warn!(%error, "Redis async_increment_cache failed, falling back to in-memory result"); + local + })) + } + + /// `DualCache.async_increment_cache_pipeline`: memory, then one pipeline of `INCRBYFLOAT` + /// with `EXPIRE` per increment. + pub async fn increment_many(&self, increments: &[Increment]) -> Result<(), Error> { + increments + .iter() + .try_for_each(|increment| self.add_local(increment).map(drop))?; + let Some(client) = self.remote.client() else { + return Ok(()); + }; + let pipeline = RedisBatch::pipelined(client.clone()); + let totals = increments.iter().map(|increment| { + pipeline.add(commands::incr_by_float( + &increment.counter.key(), + increment.amount.into(), + increment.ttl.map(Expiry::Refresh), + )) + }); + let totals = Declared::all(totals.collect::>()).await; + if let Some(error) = totals.into_iter().find_map(Result::err) { + tracing::warn!(%error, "Redis async_increment_cache_pipeline failed, falling back to in-memory result"); + } + Ok(()) + } + + /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, the Redis increments on + /// the request's post-call batch. + pub fn increment_many_post_call(&self, request: &RequestRedis, increments: &[Increment]) -> Result<(), Error> { + increments + .iter() + .try_for_each(|increment| self.add_local(increment).map(drop))?; + let Some(batch) = self.remote.batch(Send::PostCall(request)) else { + return Ok(()); + }; + let declared = increments + .iter() + .map(|increment| scripts::counter_increment(&batch, increment)) + .collect::>(); + request.spawn(async move { + let outcomes = join_all(declared.into_iter().map(Declared::settled)).await; + if let Some(error) = outcomes.into_iter().find_map(Result::err) { + tracing::warn!(%error, "deferred Redis increment failed"); + } + }); + Ok(()) + } + + fn declare_release(&self, send: Send<'_>, gauges: &[Scope], slot_id: &str) -> Option>> { + Some(scripts::parallel_release(&self.batch(send)?, gauges, slot_id)) + } + + fn declare_token_increment(&self, send: Send<'_>, increments: &[Increment]) -> Option>> { + Some(scripts::token_increment(&self.batch(send)?, increments)) + } + + /// `InMemoryCache.increment_cache`: `(current or 0) + amount` under one lock. + fn add_local(&self, increment: &Increment) -> Result { + let key = increment.counter.key(); + let _guard = self.memory.guard(); + let total = match self.memory.get(&key)? { + None => increment.amount, + Some(Local::Number(current)) => current + increment.amount, + Some(_) => return Err(litellm_cache::Error::InvalidEntry(None).into()), + }; + self.memory.set(&key, Local::Number(total), increment.ttl)?; + Ok(total) + } + + fn batch(&self, send: Send<'_>) -> Option { + self.remote.batch(send) + } +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs index 9b5c723b214..93c7605c932 100644 --- a/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs +++ b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs @@ -1,10 +1,16 @@ use std::{sync::LazyLock, time::Duration}; -use futures_util::future::BoxFuture; use litellm_python_compat::repr::float_repr; -use litellm_redis_client::{Declared, RedisBatch, RedisClient, Script, Value}; +use litellm_redis_client::{ + Command, Declared, Error, Expiry, RedisBatch, Script, Value, + commands::{self, unexpected_reply}, +}; -use super::{Counter, Error, Increment, Scope, lua}; +use super::{ + keys::{Counter, Scope}, + lua, + types::{AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply, Increment, WindowCount}, +}; use crate::Number; static BATCH_RATE_LIMITER: LazyLock