diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index f9e7a148706..9863c46783f 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -6,9 +6,9 @@ `litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations -`litellm-cache-response` owns response keys, controls, entries, and the Python-compatible response codec. It has no runtime dependency on a specific cache backend or Python +`litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python -The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host +The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum, which only dispatches. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host ## Native Rust use diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs new file mode 100644 index 00000000000..68af5278c8c --- /dev/null +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -0,0 +1,46 @@ +use std::{sync::Mutex, time::Duration}; + +use litellm_cache::{BaseCache, Error}; +use serde_json::Value; + +use crate::{CacheEntry, ResponseCache, ResponseCacheRequest}; + +/// Defers async writes until `flush_size` entries are pending, then stores them as one batch. +pub struct WriteBuffer { + flush_size: usize, + entries: Mutex>, +} + +impl WriteBuffer { + pub fn new(flush_size: usize) -> Self { + Self { + flush_size: flush_size.max(1), + entries: Mutex::new(Vec::new()), + } + } + + pub async fn async_store>( + &self, + cache: &ResponseCache, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + let pending = { + let mut entries = self.entries.lock().map_err(|_| Error::Unavailable)?; + entries.push((request.clone(), response, now)); + (entries.len() >= self.flush_size).then(|| std::mem::take(&mut *entries)) + }; + // A failed flush drops its batch, as Python does. Requeueing would grow the + // buffer and re-send an ever larger pipeline on every write during an outage. + match pending { + Some(pending) => cache.async_store_entries(pending).await, + None => Ok(()), + } + } + + pub fn clear(&self) -> Result<(), Error> { + self.entries.lock().map_err(|_| Error::Unavailable)?.clear(); + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index 72a507f8ee9..91b36ebe24b 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -1,8 +1,10 @@ +mod buffer; mod caching; mod codec; mod embedding; mod response; +pub use buffer::WriteBuffer; pub use caching::{ CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key, get_cache_key, should_use_cache, diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 7c69e5a1d55..94b25626f86 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -11,7 +11,7 @@ use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec, - ResponseCacheRequest, + ResponseCacheRequest, WriteBuffer, }; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; @@ -414,3 +414,71 @@ async fn deferred_entries_keep_the_time_they_were_produced() { None ); } + +#[tokio::test] +async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let buffer = WriteBuffer::new(2); + let mut first = request(); + first.max_age = Some(Duration::from_secs(10)); + let mut second = request(); + second.key.preset = Some("tenant:other".into()); + + buffer + .async_store( + &cache, + &first, + json!({"answer": 7}), + Duration::from_secs(100), + ) + .await + .unwrap(); + assert_eq!( + cache.lookup(&first, Duration::from_secs(100)).unwrap(), + None + ); + + buffer + .async_store( + &cache, + &second, + json!({"answer": 8}), + Duration::from_secs(200), + ) + .await + .unwrap(); + assert_eq!( + cache.lookup(&first, Duration::from_secs(110)).unwrap(), + Some(json!({"answer": 7})) + ); + assert_eq!( + cache.lookup(&first, Duration::from_secs(111)).unwrap(), + None + ); + assert_eq!( + cache.lookup(&second, Duration::from_secs(200)).unwrap(), + Some(json!({"answer": 8})) + ); +} + +#[tokio::test] +async fn write_buffer_clear_drops_pending_entries() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let buffer = WriteBuffer::new(2); + let mut other = request(); + other.key.preset = Some("tenant:other".into()); + let now = Duration::from_secs(100); + + buffer + .async_store(&cache, &request(), json!({"answer": 7}), now) + .await + .unwrap(); + buffer.clear().unwrap(); + buffer + .async_store(&cache, &other, json!({"answer": 8}), now) + .await + .unwrap(); + + assert_eq!(cache.lookup(&request(), now).unwrap(), None); + assert_eq!(cache.lookup(&other, now).unwrap(), None); +} diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs new file mode 100644 index 00000000000..44c133d4611 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -0,0 +1,294 @@ +use litellm_cache_response::PartialHits; +use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + types::PyDict, +}; +use serde_json::Value; + +use super::{ + cache_error, + callback::PythonCallback, + future::{ready_none, ready_value}, + native::NativeResponseCache, + request::{now, request, requests}, +}; + +pub(super) enum CacheBinding { + Disabled, + Native(NativeResponseCache), + PythonCallback(PythonCallback), +} + +#[pyclass(frozen, name = "_CacheTestBinding")] +pub(crate) struct ResolvedCache { + binding: CacheBinding, + pid: u32, +} + +impl ResolvedCache { + pub(super) fn new(binding: CacheBinding) -> Self { + Self { + binding, + pid: std::process::id(), + } + } + + fn check_process(&self) -> PyResult<()> { + if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache bindings must be resolved again after fork", + )); + } + Ok(()) + } + + pub(crate) fn lookup_step( + &self, + py: Python<'_>, + input: &Bound<'_, PyAny>, + kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + self.check_process()?; + let awaitable = match &self.binding { + CacheBinding::Disabled => ready_none(py)?, + CacheBinding::Native(service) => { + let request = request(input)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup(&request, now()).await }, + cache_error, + )? + } + CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?, + }; + Ok(ExecutionStep::Await(awaitable.unbind())) + } +} + +#[pymethods] +impl ResolvedCache { + #[getter] + fn kind(&self) -> &'static str { + match self.binding { + CacheBinding::Disabled => "disabled", + CacheBinding::Native(_) => "native", + CacheBinding::PythonCallback(_) => "python_callback", + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn lookup( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(py.None()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup(&request, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(callback) => { + callback.lookup(py, callback_kwargs).map(Bound::unbind) + } + } + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn store( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + response: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult<()> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + release_gil(py, move || service.store(&request, response, now())) + .map_err(cache_error) + } + CacheBinding::PythonCallback(callback) => callback.store(py, response, callback_kwargs), + } + } + + /// Native bindings return `{values, missing_indices}`, while a Python callback returns the + /// list of its per-request results. + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn lookup_batch( + &self, + py: Python<'_>, + requests: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyAny>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + to_py(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup_batch(&requests, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(callback) => callback + .lookup_batch(py, requests, callback_kwargs) + .map(Bound::unbind), + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn async_lookup<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)? + else { + unreachable!() + }; + Ok(awaitable.into_bound(py)) + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn async_store<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + response: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + run_async( + py, + async move { service.async_store(&request, response, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => { + callback.async_store(py, response, callback_kwargs) + } + } + } + + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn async_lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + ready_value(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup_batch(&requests, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => { + callback.async_lookup_batch(py, requests, callback_kwargs) + } + } + } + + /// A Python callback receives the caller's original result through `callback_result`. + #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] + fn async_store_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + responses: &Bound<'py, PyAny>, + callback_result: Option<&Bound<'py, PyAny>>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let responses: Vec = from_py(responses)?; + if requests.len() != responses.len() { + return Err(PyValueError::new_err( + "batch cache requests and responses must have equal lengths", + )); + } + let entries = requests.into_iter().zip(responses).collect(); + let service = service.clone(); + run_async( + py, + async move { service.async_store_batch(entries, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => { + callback.async_store_batch(py, callback_result, callback_kwargs) + } + } + } + + fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async(py, async move { service.async_flush().await }, cache_error) + } + CacheBinding::PythonCallback(callback) => callback.async_flush(py), + } + } + + fn ping<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async( + py, + async move { service.test_connection().await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => callback.ping(py), + } + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let CacheBinding::PythonCallback(callback) = &self.binding { + callback.traverse(&visit)?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/callback.rs b/litellm-rust/crates/python-bridge/src/cache/callback.rs new file mode 100644 index 00000000000..318f9d02080 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/callback.rs @@ -0,0 +1,169 @@ +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyTypeError, PyValueError}, + prelude::*, + types::{PyDict, PyList, PyTuple}, +}; + +use super::future::ready_none; + +/// A custom Python cache object, driven through the built-in `Cache` API so a `Cache` subclass +/// works unchanged. +pub(super) struct PythonCallback(Py); + +impl PythonCallback { + pub(super) fn new(object: Py) -> Self { + Self(object) + } + + pub(super) fn lookup<'py>( + &self, + py: Python<'py>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.0 + .bind(py) + .call_method("get_cache", (), Some(callback_kwargs(kwargs)?)) + } + + pub(super) fn async_lookup<'py>( + &self, + py: Python<'py>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.0 + .bind(py) + .call_method("async_get_cache", (), Some(callback_kwargs(kwargs)?)) + } + + pub(super) fn store( + &self, + py: Python<'_>, + response: &Bound<'_, PyAny>, + kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult<()> { + self.0 + .bind(py) + .call_method("add_cache", (response,), Some(callback_kwargs(kwargs)?)) + .map(|_| ()) + } + + pub(super) fn async_store<'py>( + &self, + py: Python<'py>, + response: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.0.bind(py).call_method( + "async_add_cache", + (response,), + Some(callback_kwargs(kwargs)?), + ) + } + + /// The built-in `Cache` API has no batch read, so the callback receives one + /// `get_cache(**kwargs)` call per request, in order, and the results come back as a list. + pub(super) fn lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let results = PyList::empty(py); + for kwargs in batch_callback_kwargs(requests, kwargs)? { + results.append( + self.0 + .bind(py) + .call_method("get_cache", (), Some(&kwargs))?, + )?; + } + Ok(results.into_any()) + } + + pub(super) fn async_lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let awaitables = batch_callback_kwargs(requests, kwargs)? + .iter() + .map(|kwargs| { + self.0 + .bind(py) + .call_method("async_get_cache", (), Some(kwargs)) + }) + .collect::>>()?; + py.import("asyncio")? + .call_method1("gather", PyTuple::new(py, awaitables)?) + } + + /// Receives the caller's original result, because the built-in + /// `Cache.async_add_cache_pipeline` splits the batch itself. + pub(super) fn async_store_batch<'py>( + &self, + py: Python<'py>, + result: Option<&Bound<'py, PyAny>>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + let result = result.ok_or_else(|| { + PyTypeError::new_err("Python cache callbacks require their original callback_result") + })?; + self.0.bind(py).call_method( + "async_add_cache_pipeline", + (result,), + Some(callback_kwargs(kwargs)?), + ) + } + + /// The built-in `Cache` facade has no flush of its own; its backend does. + pub(super) fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { + let object = self.0.bind(py); + let backend = match object.getattr_opt("cache")? { + Some(backend) if !backend.is_none() => backend, + _ => object.clone(), + }; + if backend.hasattr("async_flush_cache")? { + return backend.call_method0("async_flush_cache"); + } + backend.call_method0("flush_cache")?; + ready_none(py) + } + + pub(super) fn ping<'py>(&self, py: Python<'py>) -> PyResult> { + self.0.bind(py).call_method0("ping") + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } +} + +fn callback_kwargs<'a, 'py>( + kwargs: Option<&'a Bound<'py, PyDict>>, +) -> PyResult<&'a Bound<'py, PyDict>> { + kwargs.ok_or_else(|| { + PyTypeError::new_err("Python cache callbacks require their original callback_kwargs") + }) +} + +fn batch_callback_kwargs<'py>( + requests: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyAny>>, +) -> PyResult>> { + let kwargs = kwargs + .ok_or_else(|| { + PyTypeError::new_err( + "Python cache callbacks require one original callback_kwargs mapping per request", + ) + })? + .try_iter()? + .map(|item| Ok(item?.cast_into::()?)) + .collect::>>()?; + if kwargs.len() != requests.len()? { + return Err(PyValueError::new_err( + "batch cache requests and callback_kwargs must have equal lengths", + )); + } + Ok(kwargs) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 2ad13ce200f..19220bf868f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -9,7 +9,7 @@ use pyo3::{ }; use serde_json::Value; -use super::{CacheTestHandle, native::NativeResponseCache}; +use super::{handle::CacheTestHandle, native::NativeResponseCache}; struct ClassGuard { class: Py, diff --git a/litellm-rust/crates/python-bridge/src/cache/future.rs b/litellm-rust/crates/python-bridge/src/cache/future.rs new file mode 100644 index 00000000000..42593eee1f4 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/future.rs @@ -0,0 +1,18 @@ +use litellm_host_python::to_py; +use pyo3::prelude::*; + +pub(super) fn ready_none(py: Python<'_>) -> PyResult> { + ready_value(py, &()) +} + +pub(super) fn ready_value<'py, T: serde::Serialize>( + py: Python<'py>, + value: &T, +) -> PyResult> { + let future = py + .import("asyncio")? + .call_method0("get_running_loop")? + .call_method0("create_future")?; + future.call_method1("set_result", (to_py(py, value)?,))?; + Ok(future) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs new file mode 100644 index 00000000000..42d7f2c2f3d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -0,0 +1,106 @@ +use std::time::Duration; + +use litellm_host_python::release_gil; +use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; + +use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; +use crate::python_settings::PythonSettings; + +const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); + +#[derive(FromPyObject)] +struct PythonCacheSettings { + default_redis_ttl: Option, +} + +fn redis_default_ttl(py: Python<'_>) -> PyResult { + let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; + settings + .default_redis_ttl + .map(duration) + .transpose() + .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) +} + +#[pyclass(frozen, name = "_CacheTestHandle")] +pub(crate) struct CacheTestHandle { + service: NativeResponseCache, + pub(super) guard: Option, + pid: u32, +} + +impl CacheTestHandle { + pub(super) fn service(&self) -> PyResult { + if self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache handles must be recreated after fork", + )); + } + Ok(self.service.clone()) + } +} + +#[pymethods] +impl CacheTestHandle { + #[staticmethod] + #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] + fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { + Ok(Self { + service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes), + guard: None, + pid: std::process::id(), + }) + } + + #[staticmethod] + #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] + fn redis( + py: Python<'_>, + url: String, + ttl_seconds: Option, + namespace: Option, + ) -> PyResult { + let ttl = Some(match ttl_seconds { + Some(seconds) => duration(seconds)?, + None => redis_default_ttl(py)?, + }); + let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + + #[getter] + fn backend(&self) -> &'static str { + self.service.kind() + } + + fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { + let service = self.service()?; + let guard = FacadeGuard::capture(py, facade, &service)?; + let service = service.with_redis_flush_size( + facade + .getattr("redis_flush_size")? + .extract::>()?, + ); + let handle = Py::new( + py, + Self { + service, + guard: Some(guard), + pid: self.pid, + }, + )?; + facade.setattr("_native_cache_handle", handle) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some(guard) = &self.guard { + guard.traverse(visit)?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index f83788f6036..7955ed934b2 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,81 +1,21 @@ +mod binding; +mod callback; mod facade; +mod future; +mod handle; mod native; +mod request; +mod resolver; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -use facade::FacadeGuard; use litellm_cache::Error; -use litellm_cache_response::{CacheControls, CacheKeyInput, PartialHits, ResponseCacheRequest}; -use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; -use native::NativeResponseCache; use pyo3::{ - PyTraverseError, PyVisit, - exceptions::{PyRuntimeError, PyTypeError, PyValueError}, + exceptions::{PyRuntimeError, PyValueError}, prelude::*, - types::{PyDict, PyList, PyTuple}, }; -use serde::Deserialize; -use serde_json::Value; -use crate::python_settings::PythonSettings; - -const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); - -#[derive(FromPyObject)] -struct PythonCacheSettings { - default_redis_ttl: Option, -} - -fn redis_default_ttl(py: Python<'_>) -> PyResult { - let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; - settings - .default_redis_ttl - .map(duration) - .transpose() - .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) -} - -#[derive(Deserialize)] -#[serde(deny_unknown_fields)] -struct RequestInput { - key: CacheKeyInput, - controls: Option, - ttl_seconds: Option, - max_age_seconds: Option, -} - -fn request(value: &Bound<'_, PyAny>) -> PyResult { - let input: RequestInput = from_py(value)?; - request_input(input) -} - -fn request_input(input: RequestInput) -> PyResult { - let mut request = ResponseCacheRequest::new(input.key); - if let Some(controls) = input.controls { - request.controls = controls; - } - request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?; - request.max_age = input.max_age_seconds.map(duration).transpose()?; - Ok(request) -} - -fn requests(value: &Bound<'_, PyAny>) -> PyResult> { - from_py::>(value)? - .into_iter() - .map(request_input) - .collect() -} - -fn duration(seconds: f64) -> PyResult { - Duration::try_from_secs_f64(seconds) - .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) -} - -fn now() -> Duration { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() -} +pub(crate) use self::{ + binding::ResolvedCache, handle::CacheTestHandle, resolver::CacheTestResolver, +}; fn cache_error(error: Error) -> PyErr { match error { @@ -83,493 +23,3 @@ fn cache_error(error: Error) -> PyErr { _ => PyRuntimeError::new_err(error.to_string()), } } - -#[pyclass(frozen, name = "_CacheTestHandle")] -pub(crate) struct CacheTestHandle { - service: NativeResponseCache, - guard: Option, - pid: u32, -} - -impl CacheTestHandle { - fn service(&self) -> PyResult { - if self.pid != std::process::id() { - return Err(PyRuntimeError::new_err( - "native cache handles must be recreated after fork", - )); - } - Ok(self.service.clone()) - } -} - -#[pymethods] -impl CacheTestHandle { - #[staticmethod] - #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] - fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { - Ok(Self { - service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes), - guard: None, - pid: std::process::id(), - }) - } - - #[staticmethod] - #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] - fn redis( - py: Python<'_>, - url: String, - ttl_seconds: Option, - namespace: Option, - ) -> PyResult { - let ttl = Some(match ttl_seconds { - Some(seconds) => duration(seconds)?, - None => redis_default_ttl(py)?, - }); - let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) - .map_err(cache_error)?; - Ok(Self { - service, - guard: None, - pid: std::process::id(), - }) - } - - #[getter] - fn backend(&self) -> &'static str { - self.service.kind() - } - - fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { - let service = self.service()?; - let guard = FacadeGuard::capture(py, facade, &service)?; - let service = service.with_redis_flush_size( - facade - .getattr("redis_flush_size")? - .extract::>()?, - ); - let handle = Py::new( - py, - Self { - service, - guard: Some(guard), - pid: self.pid, - }, - )?; - facade.setattr("_native_cache_handle", handle) - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - if let Some(guard) = &self.guard { - guard.traverse(visit)?; - } - Ok(()) - } -} - -enum CacheBinding { - Disabled, - Native(NativeResponseCache), - PythonCallback(Py), -} - -#[pyclass(frozen, name = "_CacheTestBinding")] -pub(crate) struct ResolvedCache { - binding: CacheBinding, - pid: u32, -} - -impl ResolvedCache { - fn check_process(&self) -> PyResult<()> { - if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() { - return Err(PyRuntimeError::new_err( - "native cache bindings must be resolved again after fork", - )); - } - Ok(()) - } - - pub(crate) fn lookup_step( - &self, - py: Python<'_>, - input: &Bound<'_, PyAny>, - kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult { - self.check_process()?; - let awaitable = match &self.binding { - CacheBinding::Disabled => ready_none(py)?, - CacheBinding::Native(service) => { - let request = request(input)?; - let service = service.clone(); - run_async( - py, - async move { service.async_lookup(&request, now()).await }, - cache_error, - )? - } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_get_cache", - (), - Some(callback_kwargs(kwargs)?), - )?, - }; - Ok(ExecutionStep::Await(awaitable.unbind())) - } -} - -#[pymethods] -impl ResolvedCache { - #[getter] - fn kind(&self) -> &'static str { - match self.binding { - CacheBinding::Disabled => "disabled", - CacheBinding::Native(_) => "native", - CacheBinding::PythonCallback(_) => "python_callback", - } - } - - #[pyo3(signature = (request, *, callback_kwargs=None))] - fn lookup( - &self, - py: Python<'_>, - request: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => Ok(py.None()), - CacheBinding::Native(service) => { - let request = self::request(request)?; - let service = service.clone(); - let response = release_gil(py, move || service.lookup(&request, now())) - .map_err(cache_error)?; - to_py(py, &response) - } - CacheBinding::PythonCallback(object) => object - .bind(py) - .call_method( - "get_cache", - (), - Some(self::callback_kwargs(callback_kwargs)?), - ) - .map(Bound::unbind), - } - } - - #[pyo3(signature = (request, response, *, callback_kwargs=None))] - fn store( - &self, - py: Python<'_>, - request: &Bound<'_, PyAny>, - response: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult<()> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => Ok(()), - CacheBinding::Native(service) => { - let request = self::request(request)?; - let response: Value = from_py(response)?; - let service = service.clone(); - release_gil(py, move || service.store(&request, response, now())) - .map_err(cache_error) - } - CacheBinding::PythonCallback(object) => object - .bind(py) - .call_method( - "add_cache", - (response,), - Some(self::callback_kwargs(callback_kwargs)?), - ) - .map(|_| ()), - } - } - - /// Native bindings return `{values, missing_indices}`. The built-in `Cache` API has no batch - /// read, so a Python callback receives one `get_cache(**kwargs)` call per request, in order, - /// and the results come back as a list. - #[pyo3(signature = (requests, *, callback_kwargs=None))] - fn lookup_batch( - &self, - py: Python<'_>, - requests: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyAny>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => { - let requests = self::requests(requests)?; - to_py(py, &PartialHits::new(vec![None; requests.len()])) - } - CacheBinding::Native(service) => { - let requests = self::requests(requests)?; - let service = service.clone(); - let response = release_gil(py, move || service.lookup_batch(&requests, now())) - .map_err(cache_error)?; - to_py(py, &response) - } - CacheBinding::PythonCallback(object) => { - let results = PyList::empty(py); - for kwargs in batch_callback_kwargs(requests, callback_kwargs)? { - results.append(object.bind(py).call_method( - "get_cache", - (), - Some(&kwargs), - )?)?; - } - Ok(results.into_any().unbind()) - } - } - } - - #[pyo3(signature = (request, *, callback_kwargs=None))] - fn async_lookup<'py>( - &self, - py: Python<'py>, - request: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyDict>>, - ) -> PyResult> { - let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)? - else { - unreachable!() - }; - Ok(awaitable.into_bound(py)) - } - - #[pyo3(signature = (request, response, *, callback_kwargs=None))] - fn async_store<'py>( - &self, - py: Python<'py>, - request: &Bound<'py, PyAny>, - response: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyDict>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let request = self::request(request)?; - let response: Value = from_py(response)?; - let service = service.clone(); - run_async( - py, - async move { service.async_store(&request, response, now()).await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_add_cache", - (response,), - Some(self::callback_kwargs(callback_kwargs)?), - ), - } - } - - #[pyo3(signature = (requests, *, callback_kwargs=None))] - fn async_lookup_batch<'py>( - &self, - py: Python<'py>, - requests: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyAny>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => { - let requests = self::requests(requests)?; - ready_value(py, &PartialHits::new(vec![None; requests.len()])) - } - CacheBinding::Native(service) => { - let requests = self::requests(requests)?; - let service = service.clone(); - run_async( - py, - async move { service.async_lookup_batch(&requests, now()).await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => { - let awaitables = batch_callback_kwargs(requests, callback_kwargs)? - .iter() - .map(|kwargs| { - object - .bind(py) - .call_method("async_get_cache", (), Some(kwargs)) - }) - .collect::>>()?; - py.import("asyncio")? - .call_method1("gather", PyTuple::new(py, awaitables)?) - } - } - } - - /// A Python callback receives the caller's original result through `callback_result`, because - /// the built-in `Cache.async_add_cache_pipeline` splits the batch itself. - #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] - fn async_store_batch<'py>( - &self, - py: Python<'py>, - requests: &Bound<'py, PyAny>, - responses: &Bound<'py, PyAny>, - callback_result: Option<&Bound<'py, PyAny>>, - callback_kwargs: Option<&Bound<'py, PyDict>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let requests = self::requests(requests)?; - let responses: Vec = from_py(responses)?; - if requests.len() != responses.len() { - return Err(PyValueError::new_err( - "batch cache requests and responses must have equal lengths", - )); - } - let entries = requests.into_iter().zip(responses).collect(); - let service = service.clone(); - run_async( - py, - async move { service.async_store_batch(entries, now()).await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => { - let result = callback_result.ok_or_else(|| { - PyTypeError::new_err( - "Python cache callbacks require their original callback_result", - ) - })?; - object.bind(py).call_method( - "async_add_cache_pipeline", - (result,), - Some(self::callback_kwargs(callback_kwargs)?), - ) - } - } - } - - fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let service = service.clone(); - run_async(py, async move { service.async_flush().await }, cache_error) - } - // The built-in `Cache` facade has no flush of its own; its backend does. - CacheBinding::PythonCallback(object) => { - let object = object.bind(py); - let backend = match object.getattr_opt("cache")? { - Some(backend) if !backend.is_none() => backend, - _ => object.clone(), - }; - if backend.hasattr("async_flush_cache")? { - return backend.call_method0("async_flush_cache"); - } - backend.call_method0("flush_cache")?; - ready_none(py) - } - } - } - - fn ping<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let service = service.clone(); - run_async( - py, - async move { service.test_connection().await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => object.bind(py).call_method0("ping"), - } - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - if let CacheBinding::PythonCallback(object) = &self.binding { - visit.call(object)?; - } - Ok(()) - } -} - -fn callback_kwargs<'a, 'py>( - kwargs: Option<&'a Bound<'py, PyDict>>, -) -> PyResult<&'a Bound<'py, PyDict>> { - kwargs.ok_or_else(|| { - PyTypeError::new_err("Python cache callbacks require their original callback_kwargs") - }) -} - -fn batch_callback_kwargs<'py>( - requests: &Bound<'py, PyAny>, - kwargs: Option<&Bound<'py, PyAny>>, -) -> PyResult>> { - let kwargs = kwargs - .ok_or_else(|| { - PyTypeError::new_err( - "Python cache callbacks require one original callback_kwargs mapping per request", - ) - })? - .try_iter()? - .map(|item| Ok(item?.cast_into::()?)) - .collect::>>()?; - if kwargs.len() != requests.len()? { - return Err(PyValueError::new_err( - "batch cache requests and callback_kwargs must have equal lengths", - )); - } - Ok(kwargs) -} - -fn ready_none(py: Python<'_>) -> PyResult> { - ready_value(py, &()) -} - -fn ready_value<'py, T: serde::Serialize>( - py: Python<'py>, - value: &T, -) -> PyResult> { - let future = py - .import("asyncio")? - .call_method0("get_running_loop")? - .call_method0("create_future")?; - future.call_method1("set_result", (to_py(py, value)?,))?; - Ok(future) -} - -#[pyclass(frozen, name = "_CacheTestResolver")] -pub(crate) struct CacheTestResolver { - namespace: Py, -} - -#[pymethods] -impl CacheTestResolver { - #[new] - fn new(namespace: Py) -> Self { - Self { namespace } - } - - pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult { - let object = self.namespace.bind(py).getattr("cache")?; - let binding = if object.is_none() { - CacheBinding::Disabled - } else if let Ok(handle) = object.extract::>() { - CacheBinding::Native(handle.service()?) - } else if let Some(service) = facade::resolve(py, &object)? { - CacheBinding::Native(service) - } else { - CacheBinding::PythonCallback(object.unbind()) - }; - Ok(ResolvedCache { - binding, - pid: std::process::id(), - }) - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.namespace) - } -} diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 3fc8f61dff6..a718d07b286 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -4,25 +4,19 @@ use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ - CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, + CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; use serde_json::Value; -use tokio::sync::Mutex; #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), Redis { cache: Arc>>, - buffer: Option>, + buffer: Option>, }, } -pub(super) struct RedisWriteBuffer { - flush_size: usize, - entries: Mutex>, -} - impl NativeResponseCache { pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { Self::Memory(Arc::new(ResponseCache::new(Arc::new( @@ -33,7 +27,7 @@ impl NativeResponseCache { Some(Arc::new(|entry| { ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) })), - super::now, + super::request::now, ), )))) } @@ -84,12 +78,7 @@ impl NativeResponseCache { match self { Self::Redis { cache, .. } => Self::Redis { cache, - buffer: flush_size.map(|flush_size| { - Arc::new(RedisWriteBuffer { - flush_size: flush_size.max(1), - entries: Mutex::new(Vec::new()), - }) - }), + buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), }, memory => memory, } @@ -155,19 +144,7 @@ impl NativeResponseCache { Self::Redis { cache, buffer: Some(buffer), - } => { - let pending = { - let mut entries = buffer.entries.lock().await; - entries.push((request.clone(), response, now)); - (entries.len() >= buffer.flush_size).then(|| std::mem::take(&mut *entries)) - }; - // A failed flush drops its batch, as Python does. Requeueing would grow the - // buffer and re-send an ever larger pipeline on every write during an outage. - match pending { - Some(pending) => cache.async_store_entries(pending).await, - None => Ok(()), - } - } + } => buffer.async_store(cache, request, response, now).await, } } @@ -198,7 +175,7 @@ impl NativeResponseCache { Self::Memory(cache) => cache.async_flush().await, Self::Redis { cache, buffer } => { if let Some(buffer) = buffer { - buffer.entries.lock().await.clear(); + buffer.clear()?; } cache.async_flush().await } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs new file mode 100644 index 00000000000..d8793abd115 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -0,0 +1,48 @@ +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; +use litellm_host_python::from_py; +use pyo3::{exceptions::PyValueError, prelude::*}; +use serde::Deserialize; + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RequestInput { + key: CacheKeyInput, + controls: Option, + ttl_seconds: Option, + max_age_seconds: Option, +} + +pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { + let input: RequestInput = from_py(value)?; + request_input(input) +} + +fn request_input(input: RequestInput) -> PyResult { + let mut request = ResponseCacheRequest::new(input.key); + if let Some(controls) = input.controls { + request.controls = controls; + } + request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?; + request.max_age = input.max_age_seconds.map(duration).transpose()?; + Ok(request) +} + +pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { + from_py::>(value)? + .into_iter() + .map(request_input) + .collect() +} + +pub(super) fn duration(seconds: f64) -> PyResult { + Duration::try_from_secs_f64(seconds) + .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) +} + +pub(super) fn now() -> Duration { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() +} diff --git a/litellm-rust/crates/python-bridge/src/cache/resolver.rs b/litellm-rust/crates/python-bridge/src/cache/resolver.rs new file mode 100644 index 00000000000..ef6f142e0a1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/resolver.rs @@ -0,0 +1,39 @@ +use pyo3::{PyTraverseError, PyVisit, prelude::*}; + +use super::{ + binding::{CacheBinding, ResolvedCache}, + callback::PythonCallback, + facade, + handle::CacheTestHandle, +}; + +#[pyclass(frozen, name = "_CacheTestResolver")] +pub(crate) struct CacheTestResolver { + namespace: Py, +} + +#[pymethods] +impl CacheTestResolver { + #[new] + fn new(namespace: Py) -> Self { + Self { namespace } + } + + pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult { + let object = self.namespace.bind(py).getattr("cache")?; + let binding = if object.is_none() { + CacheBinding::Disabled + } else if let Ok(handle) = object.extract::>() { + CacheBinding::Native(handle.service()?) + } else if let Some(service) = facade::resolve(py, &object)? { + CacheBinding::Native(service) + } else { + CacheBinding::PythonCallback(PythonCallback::new(object.unbind())) + }; + Ok(ResolvedCache::new(binding)) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.namespace) + } +}