mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
feat(redis): authenticate host credential providers through Rust
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
LiteLLM Rust DB / db-queries (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
LiteLLM Rust DB / db-queries (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
A credential_provider object reaches the resolver as a ProviderHandle, so it resolves to Auth::Provider and two RedisCaches share a client only when they share the provider. The process registry of native clients now holds each client weakly, so a discarded Router drops its client and its provider instead of keeping them for the life of the process.
This commit is contained in:
parent
ddaf8c6c2b
commit
f7a162a519
15 changed files with 307 additions and 52 deletions
|
|
@ -87,6 +87,7 @@ pub(crate) async fn connect(
|
|||
events: Some(events.clone()),
|
||||
http: resources.http.clone(),
|
||||
environment: resources.lookup.clone(),
|
||||
credential_provider: None,
|
||||
},
|
||||
breaker,
|
||||
)?)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,47 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_host_python::{PythonContext, attach_blocking};
|
||||
use litellm_redis_connect::{CredentialSource, CredentialStage, Error, Expiring, SourceFuture};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
const MODULE: &str = "litellm.rust_bridge.host.redis_credentials";
|
||||
|
||||
/// A redis-py `CredentialProvider` the SDK caller passed in `cache_kwargs`, asked through the host
|
||||
/// module the way redis-py's async connection asks it.
|
||||
pub(crate) struct PythonCredentialSource {
|
||||
provider: Arc<Py<PyAny>>,
|
||||
context: PythonContext,
|
||||
}
|
||||
|
||||
impl PythonCredentialSource {
|
||||
pub(crate) fn new(provider: Py<PyAny>, context: PythonContext) -> Self {
|
||||
Self {
|
||||
provider: Arc::new(provider),
|
||||
context,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CredentialSource for PythonCredentialSource {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
let provider = Arc::clone(&self.provider);
|
||||
let context = self.context.clone();
|
||||
Box::pin(async move {
|
||||
let (username, password) = attach_blocking(context, move |py| {
|
||||
py.import(MODULE)?
|
||||
.getattr("credentials")?
|
||||
.call1((provider.bind(py),))?
|
||||
.extract::<(Option<String>, String)>()
|
||||
})
|
||||
.await
|
||||
.and_then(|credentials| credentials)
|
||||
.map_err(Error::credential(CredentialStage::Provider))?;
|
||||
Ok(Expiring {
|
||||
username,
|
||||
password: SecretValue::new(password),
|
||||
expires_at: None,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,5 @@
|
|||
mod credentials;
|
||||
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
sync::{Arc, LazyLock, Mutex, PoisonError},
|
||||
|
|
@ -6,22 +8,28 @@ use std::{
|
|||
|
||||
use litellm_circuit_breaker::Config as BreakerConfig;
|
||||
use litellm_core_utils::settings::ProcessEnvironment;
|
||||
use litellm_redis_client::{RedisClient, RedisEvents};
|
||||
use litellm_redis_connect::{HostServices, SharedEnvironment};
|
||||
use litellm_host_python::PythonContext;
|
||||
use litellm_redis_client::{RedisClient, RedisEvents, WeakRedisClient};
|
||||
use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment};
|
||||
use litellm_redis_resolve::Kwargs;
|
||||
use litellm_redis_types::Resolution;
|
||||
use litellm_redis_types::{ProviderHandle, Resolution};
|
||||
use pyo3::{
|
||||
prelude::*,
|
||||
sync::MutexExt,
|
||||
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple},
|
||||
};
|
||||
|
||||
use self::credentials::PythonCredentialSource;
|
||||
use crate::redis::events::{ServiceLoggingEvents, state_label};
|
||||
|
||||
/// The kwargs `_merge_redis_kwargs` returns for a `RedisCache`, which the resolver reads the way
|
||||
/// redis-py's constructors do. A key whose value is a Python object, such as `retry`, reaches the
|
||||
/// resolver by name only, so it can name the setting it cannot honor.
|
||||
pub(crate) struct RedisSpec(Kwargs);
|
||||
/// resolver by name only, so it can name the setting it cannot honor. A `credential_provider` object
|
||||
/// is also kept, so the connector can ask it for credentials, and reaches the resolver as a handle.
|
||||
pub(crate) struct RedisSpec {
|
||||
kwargs: Kwargs,
|
||||
credential_provider: Option<Py<PyAny>>,
|
||||
}
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
|
||||
type Error = PyErr;
|
||||
|
|
@ -36,12 +44,20 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
|
|||
false => host_objects.push(name.extract::<String>()?),
|
||||
}
|
||||
}
|
||||
let credential_provider = kwargs
|
||||
.cast::<PyDict>()?
|
||||
.get_item("credential_provider")?
|
||||
.filter(|provider| !provider.is_none() && !is_plain(provider))
|
||||
.map(Bound::unbind);
|
||||
let fields: BTreeMap<String, serde_json::Value> =
|
||||
litellm_host_python::from_py_argument(plain.as_any())?;
|
||||
let objects = host_objects
|
||||
.into_iter()
|
||||
.map(|name| (name, serde_json::Value::Bool(true)));
|
||||
Ok(Self(litellm_redis_resolve::kwargs(fields.into_iter().chain(objects))))
|
||||
Ok(Self {
|
||||
kwargs: litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)),
|
||||
credential_provider,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -76,14 +92,23 @@ fn is_plain(value: &Bound<'_, PyAny>) -> bool {
|
|||
|| value.is_instance_of::<PyDict>()
|
||||
}
|
||||
|
||||
/// A client the process built, held weakly so a discarded `Router` drops its client and provider.
|
||||
struct Shared {
|
||||
resolved: Resolution,
|
||||
client: RedisClient,
|
||||
client: WeakRedisClient,
|
||||
}
|
||||
|
||||
static CLIENTS: LazyLock<Mutex<Vec<Shared>>> = LazyLock::new(|| Mutex::new(Vec::new()));
|
||||
|
||||
/// The one client for a resolved Redis config, so every family in this process shares its pool
|
||||
fn live_clients(py: Python<'_>) -> std::sync::MutexGuard<'static, Vec<Shared>> {
|
||||
let mut clients = CLIENTS
|
||||
.lock_py_attached(py)
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
clients.retain(|shared| shared.client.upgrade().is_some());
|
||||
clients
|
||||
}
|
||||
|
||||
/// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool
|
||||
/// and its breaker. Like redis-py, which never connects at construction, a config the resolver or
|
||||
/// the connector cannot honor leaves the family on memory with a warning rather than failing the
|
||||
/// `Router`.
|
||||
|
|
@ -94,8 +119,16 @@ fn shared_client(
|
|||
breaker: Option<RedisBreaker>,
|
||||
) -> PyResult<Option<RedisClient>> {
|
||||
let environment: SharedEnvironment = Arc::new(ProcessEnvironment);
|
||||
let RedisSpec {
|
||||
kwargs,
|
||||
credential_provider,
|
||||
} = spec;
|
||||
// The registry entry's client keeps the provider alive, so no live entry shares its address
|
||||
let handle = credential_provider
|
||||
.as_ref()
|
||||
.map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64));
|
||||
let logger = crate::logger::capture(py);
|
||||
let resolved = match logger.scope(|| litellm_redis_resolve::interpret(&spec.0, None)) {
|
||||
let resolved = match logger.scope(|| litellm_redis_resolve::interpret(&kwargs, handle)) {
|
||||
Ok(resolution) => {
|
||||
logger.scope(|| {
|
||||
for ignored in &resolution.ignored {
|
||||
|
|
@ -119,12 +152,23 @@ fn shared_client(
|
|||
events: Some(events),
|
||||
http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?,
|
||||
environment,
|
||||
credential_provider: credential_provider
|
||||
.as_ref()
|
||||
.map(|provider| -> PyResult<Arc<dyn CredentialSource>> {
|
||||
Ok(Arc::new(PythonCredentialSource::new(
|
||||
provider.clone_ref(py),
|
||||
PythonContext::capture(py)?,
|
||||
)))
|
||||
})
|
||||
.transpose()?,
|
||||
};
|
||||
let mut clients = CLIENTS
|
||||
.lock_py_attached(py)
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
if let Some(shared) = clients.iter().find(|shared| shared.resolved == resolved) {
|
||||
return Ok(Some(shared.client.clone()));
|
||||
let mut clients = live_clients(py);
|
||||
if let Some(client) = clients
|
||||
.iter()
|
||||
.filter(|shared| shared.resolved == resolved)
|
||||
.find_map(|shared| shared.client.upgrade())
|
||||
{
|
||||
return Ok(Some(client));
|
||||
}
|
||||
let connected = {
|
||||
let _runtime = runtime.enter();
|
||||
|
|
@ -143,7 +187,7 @@ fn shared_client(
|
|||
};
|
||||
clients.push(Shared {
|
||||
resolved,
|
||||
client: client.clone(),
|
||||
client: client.downgrade(),
|
||||
});
|
||||
Ok(Some(client))
|
||||
}
|
||||
|
|
@ -178,12 +222,10 @@ impl NativeRedisClient {
|
|||
/// gauge counts them next to the Python `RedisCache` breakers.
|
||||
#[pyfunction]
|
||||
pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> {
|
||||
let clients = CLIENTS
|
||||
.lock_py_attached(py)
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
clients
|
||||
live_clients(py)
|
||||
.iter()
|
||||
.filter_map(|shared| shared.client.breaker_state())
|
||||
.filter_map(|shared| shared.client.upgrade())
|
||||
.filter_map(|client| client.breaker_state())
|
||||
.map(|state| state_label(&state))
|
||||
.collect()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use std::{
|
||||
future::Future,
|
||||
sync::{Arc, Mutex, PoisonError},
|
||||
sync::{Arc, Mutex, PoisonError, Weak},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
|
|
@ -14,15 +14,28 @@ use crate::{
|
|||
pubsub::Subscriber,
|
||||
};
|
||||
|
||||
/// One Redis client. Clones share the connection, the breaker and the subscriber.
|
||||
#[derive(Clone)]
|
||||
pub struct RedisClient {
|
||||
pub(crate) connection: Connection,
|
||||
pub struct RedisClient(Arc<Shared>);
|
||||
|
||||
/// A handle that does not keep the client alive, for a registry that must not outlive its users.
|
||||
#[derive(Clone)]
|
||||
pub struct WeakRedisClient(Weak<Shared>);
|
||||
|
||||
impl WeakRedisClient {
|
||||
pub fn upgrade(&self) -> Option<RedisClient> {
|
||||
self.0.upgrade().map(RedisClient)
|
||||
}
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
connection: Connection,
|
||||
subscriber: Subscriber,
|
||||
namespace: Namespace,
|
||||
breaker: Option<Breaker>,
|
||||
events: Option<Arc<dyn RedisEvents>>,
|
||||
health_check: Option<Duration>,
|
||||
last_used: Arc<Mutex<Instant>>,
|
||||
last_used: Mutex<Instant>,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
|
|
@ -35,7 +48,7 @@ pub enum Ttl {
|
|||
impl RedisClient {
|
||||
pub fn new(config: RedisConfig) -> Result<Self, ConnectError> {
|
||||
let connection = Connection::new(&config.target, &config.credentials, config.timeouts)?;
|
||||
Ok(Self {
|
||||
Ok(Self(Arc::new(Shared {
|
||||
subscriber: Subscriber::new(
|
||||
&config.target,
|
||||
&config.credentials,
|
||||
|
|
@ -47,13 +60,17 @@ impl RedisClient {
|
|||
breaker: config.breaker,
|
||||
events: config.events,
|
||||
health_check: config.health_check,
|
||||
last_used: Arc::new(Mutex::new(Instant::now())),
|
||||
})
|
||||
last_used: Mutex::new(Instant::now()),
|
||||
})))
|
||||
}
|
||||
|
||||
pub fn downgrade(&self) -> WeakRedisClient {
|
||||
WeakRedisClient(Arc::downgrade(&self.0))
|
||||
}
|
||||
|
||||
/// The breaker's current state, or none when the breaker is disabled.
|
||||
pub fn breaker_state(&self) -> Option<State> {
|
||||
self.breaker.as_ref().map(Breaker::state)
|
||||
self.0.breaker.as_ref().map(Breaker::state)
|
||||
}
|
||||
|
||||
pub async fn ping(&self) -> Result<(), Error> {
|
||||
|
|
@ -119,7 +136,7 @@ impl RedisClient {
|
|||
pub async fn incr_by(&self, key: &str, delta: i64, ttl: Duration) -> Result<i64, Error> {
|
||||
let key = self.key(key);
|
||||
self.guarded(Operation::Command("INCRBY"), async {
|
||||
let mut connection = self.connection.clone();
|
||||
let mut connection = self.0.connection.clone();
|
||||
let value: i64 = commands::incr_by(key.clone(), delta)
|
||||
.query_async(&mut connection)
|
||||
.await?;
|
||||
|
|
@ -159,7 +176,7 @@ impl RedisClient {
|
|||
) -> Result<f64, Error> {
|
||||
let key = self.key(key);
|
||||
self.guarded(Operation::Command("INCRBYFLOAT"), async {
|
||||
let mut connection = self.connection.clone();
|
||||
let mut connection = self.0.connection.clone();
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline.add_command(commands::incr_by_float(key.clone(), amount));
|
||||
match expiry {
|
||||
|
|
@ -198,7 +215,7 @@ impl RedisClient {
|
|||
self.guarded(Operation::Script, async {
|
||||
script
|
||||
.invocation(&keys, &args)
|
||||
.invoke_async(&mut self.connection.clone())
|
||||
.invoke_async(&mut self.0.connection.clone())
|
||||
.await
|
||||
})
|
||||
.await
|
||||
|
|
@ -213,19 +230,19 @@ impl RedisClient {
|
|||
/// 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<u8>> {
|
||||
self.subscriber.clone().messages(channel.to_owned())
|
||||
self.0.subscriber.clone().messages(channel.to_owned())
|
||||
}
|
||||
|
||||
pub fn pipeline(&self) -> Pipeline {
|
||||
Pipeline::new(self.namespace.clone())
|
||||
Pipeline::new(self.0.namespace.clone())
|
||||
}
|
||||
|
||||
pub fn same_client(&self, other: &Self) -> bool {
|
||||
Arc::ptr_eq(&self.last_used, &other.last_used)
|
||||
Arc::ptr_eq(&self.0, &other.0)
|
||||
}
|
||||
|
||||
pub fn is_cluster(&self) -> bool {
|
||||
matches!(self.connection, Connection::Cluster(_))
|
||||
matches!(self.0.connection, Connection::Cluster(_))
|
||||
}
|
||||
|
||||
pub fn batch(&self) -> RedisBatch {
|
||||
|
|
@ -251,7 +268,7 @@ impl RedisClient {
|
|||
}
|
||||
};
|
||||
self.check_health().await;
|
||||
let results = pipeline.run(&self.connection).await;
|
||||
let results = pipeline.run(&self.0.connection).await;
|
||||
let transport = results.iter().find_map(|result| {
|
||||
result
|
||||
.as_ref()
|
||||
|
|
@ -268,7 +285,7 @@ impl RedisClient {
|
|||
}
|
||||
|
||||
fn key(&self, key: &str) -> String {
|
||||
self.namespace.apply(key)
|
||||
self.0.namespace.apply(key)
|
||||
}
|
||||
|
||||
fn keys(&self, keys: &[&str]) -> Vec<String> {
|
||||
|
|
@ -281,7 +298,7 @@ impl RedisClient {
|
|||
command: &Cmd,
|
||||
) -> Result<T, Error> {
|
||||
self.guarded(Operation::Command(name), async {
|
||||
command.query_async(&mut self.connection.clone()).await
|
||||
command.query_async(&mut self.0.connection.clone()).await
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
|
@ -297,7 +314,7 @@ impl RedisClient {
|
|||
self.check_health().await;
|
||||
let result = call
|
||||
.await
|
||||
.map_err(|error| Error::classify(error, self.connection.deployment()));
|
||||
.map_err(|error| Error::classify(error, self.0.connection.deployment()));
|
||||
settle(admission, result.as_ref().err());
|
||||
result
|
||||
}
|
||||
|
|
@ -309,10 +326,11 @@ impl RedisClient {
|
|||
|
||||
fn observe(&self, operation: Operation, started: Instant, outcome: Outcome) {
|
||||
*self
|
||||
.0
|
||||
.last_used
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner) = Instant::now();
|
||||
if let Some(events) = &self.events {
|
||||
if let Some(events) = &self.0.events {
|
||||
events.call(&CallEvent {
|
||||
operation,
|
||||
duration: started.elapsed(),
|
||||
|
|
@ -325,10 +343,11 @@ impl RedisClient {
|
|||
/// sending the next command on it, and reconnects when the ping fails. The ping is not a
|
||||
/// call of the caller's and never reaches the breaker or the events.
|
||||
async fn check_health(&self) {
|
||||
let Some(interval) = self.health_check else {
|
||||
let Some(interval) = self.0.health_check else {
|
||||
return;
|
||||
};
|
||||
let idle = self
|
||||
.0
|
||||
.last_used
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
|
|
@ -336,13 +355,14 @@ impl RedisClient {
|
|||
if idle < interval {
|
||||
return;
|
||||
}
|
||||
if let Err(error) = self.connection.ping_every_node().await {
|
||||
if let Err(error) = self.0.connection.ping_every_node().await {
|
||||
tracing::debug!(%error, "idle Redis connection failed its health check, the next command reconnects");
|
||||
}
|
||||
}
|
||||
|
||||
fn admit(&self) -> Result<Option<Admission>, Error> {
|
||||
self.breaker
|
||||
self.0
|
||||
.breaker
|
||||
.as_ref()
|
||||
.map(|breaker| {
|
||||
breaker
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ mod pubsub;
|
|||
mod script;
|
||||
|
||||
pub use batch::{Declared, RedisBatch};
|
||||
pub use client::{RedisClient, Ttl};
|
||||
pub use client::{RedisClient, Ttl, WeakRedisClient};
|
||||
pub use commands::{Amount, Expiry};
|
||||
pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts};
|
||||
pub use error::{ConnectError, Error, KeyError};
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ use litellm_redis_client::{
|
|||
use litellm_redis_types::{Auth, HealthCheck, Resolution, StaticAuth};
|
||||
|
||||
use crate::{
|
||||
Error, RefreshingCredentials,
|
||||
CredentialSource, Error, RefreshingCredentials,
|
||||
sources::{AzureAd, ElastiCacheIam, GcpIam},
|
||||
};
|
||||
|
||||
|
|
@ -19,6 +19,7 @@ pub struct HostServices {
|
|||
pub events: Option<Arc<dyn RedisEvents>>,
|
||||
pub http: litellm_http::Client,
|
||||
pub environment: SharedEnvironment,
|
||||
pub credential_provider: Option<Arc<dyn CredentialSource>>,
|
||||
}
|
||||
|
||||
// litellm/constants.py REDIS_CLUSTER_HEALTH_CHECK_INTERVAL, the default _redis.py sets for cluster clients only
|
||||
|
|
@ -78,7 +79,10 @@ fn credentials(auth: Auth, host: &HostServices) -> Result<Credentials, Error> {
|
|||
Auth::Static(StaticAuth { username, password }) => {
|
||||
Credentials::Static { username, password }
|
||||
}
|
||||
Auth::Provider(_) => return Err(Error::NoHostCredentialProvider),
|
||||
Auth::Provider(_) => match &host.credential_provider {
|
||||
Some(source) => refreshing(RefreshingCredentials::new(Arc::clone(source))),
|
||||
None => return Err(Error::NoHostCredentialProvider),
|
||||
},
|
||||
Auth::AzureAd(auth) => refreshing(RefreshingCredentials::new(AzureAd::new(
|
||||
auth,
|
||||
host.environment.clone(),
|
||||
|
|
|
|||
|
|
@ -40,10 +40,12 @@ pub enum CredentialStage {
|
|||
Exchange,
|
||||
#[strum(serialize = "Redis token response")]
|
||||
Response,
|
||||
#[strum(serialize = "Redis credential_provider")]
|
||||
Provider,
|
||||
}
|
||||
|
||||
impl Error {
|
||||
pub(crate) fn credential<E: std::error::Error + Send + Sync + 'static>(
|
||||
pub fn credential<E: std::error::Error + Send + Sync + 'static>(
|
||||
stage: CredentialStage,
|
||||
) -> impl FnOnce(E) -> Self {
|
||||
move |cause| Self::Credential {
|
||||
|
|
@ -86,6 +88,7 @@ impl Error {
|
|||
CredentialStage::Signing => "token-signing",
|
||||
CredentialStage::Exchange => "token-exchange",
|
||||
CredentialStage::Response => "token-response",
|
||||
CredentialStage::Provider => "credential-provider",
|
||||
},
|
||||
Self::CredentialTimeout { .. } => "acquisition-timeout",
|
||||
Self::ExpiredCredential => "expired-credential",
|
||||
|
|
|
|||
|
|
@ -7,4 +7,5 @@ mod sources;
|
|||
pub use connect::{HostServices, SharedEnvironment, connect};
|
||||
pub use error::{CredentialStage, Error};
|
||||
pub use events::TracingEvents;
|
||||
pub(crate) use refresh::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture};
|
||||
pub(crate) use refresh::RefreshingCredentials;
|
||||
pub use refresh::{CredentialSource, Expiring, SourceFuture};
|
||||
|
|
|
|||
|
|
@ -32,6 +32,12 @@ pub trait CredentialSource: Send + Sync + 'static {
|
|||
fn fetch(&self) -> SourceFuture<'_>;
|
||||
}
|
||||
|
||||
impl CredentialSource for Arc<dyn CredentialSource> {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
(**self).fetch()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
enum Latest {
|
||||
Pending,
|
||||
|
|
|
|||
|
|
@ -1,10 +1,13 @@
|
|||
use std::{collections::BTreeMap, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_redis_connect::{Error, HostServices, connect};
|
||||
use litellm_redis_connect::{
|
||||
CredentialSource, Error, Expiring, HostServices, SourceFuture, connect,
|
||||
};
|
||||
use litellm_redis_testing::{
|
||||
RedisDeployment, RedisSentinel, RedisTlsServer, SERVICE_NAME, Topology, eventually,
|
||||
RedisDeployment, RedisSentinel, RedisServer, RedisTlsServer, SERVICE_NAME, Topology, eventually,
|
||||
};
|
||||
use litellm_redis_types::{ProviderHandle, Resolution};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use rstest::rstest;
|
||||
|
||||
// a host hands its credential_provider object over as a handle, so these tests do the same
|
||||
|
|
@ -21,6 +24,21 @@ fn services() -> HostServices {
|
|||
events: None,
|
||||
http: litellm_http::Client::for_test(),
|
||||
environment: Arc::new(|_: &str| None),
|
||||
credential_provider: None,
|
||||
}
|
||||
}
|
||||
|
||||
struct HostPassword(&'static str);
|
||||
|
||||
impl CredentialSource for HostPassword {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
Box::pin(async move {
|
||||
Ok(Expiring {
|
||||
username: None,
|
||||
password: SecretValue::new(self.0),
|
||||
expires_at: None,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -126,7 +144,7 @@ async fn python_sentinel_kwargs_follow_the_master_through_a_failover() {
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_python_credential_provider_fails_at_boot() {
|
||||
fn a_python_credential_provider_without_a_host_source_fails_at_boot() {
|
||||
let error = connect(
|
||||
resolved(
|
||||
serde_json::json!({"host": "redis", "credential_provider": true})),
|
||||
|
|
@ -170,3 +188,39 @@ fn config_only_aws_identities_fail_before_connecting(#[case] setting: &'static s
|
|||
Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::right_password("right", true)]
|
||||
#[case::wrong_password("wrong", false)]
|
||||
#[tokio::test]
|
||||
async fn a_python_credential_provider_authenticates_through_the_host_source(
|
||||
#[case] password: &'static str,
|
||||
#[case] connects: bool,
|
||||
) {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
server
|
||||
.cli(&["ACL", "SETUSER", "default", "on", ">right"])
|
||||
.await
|
||||
.unwrap();
|
||||
let kwargs = serde_json::json!({
|
||||
"host": "127.0.0.1",
|
||||
"port": port(server.url()),
|
||||
"password": "ignored",
|
||||
"credential_provider": true,
|
||||
"socket_timeout": 2,
|
||||
});
|
||||
|
||||
let pinged = connect(
|
||||
resolved(kwargs),
|
||||
HostServices {
|
||||
credential_provider: Some(Arc::new(HostPassword(password))),
|
||||
..services()
|
||||
},
|
||||
None,
|
||||
)
|
||||
.unwrap()
|
||||
.ping()
|
||||
.await;
|
||||
|
||||
assert_eq!(pinged.is_ok(), connects, "{pinged:?}");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ async fn roundtrip(config: LiveConfig) {
|
|||
events: None,
|
||||
http: http.clone(),
|
||||
environment: Arc::new(|name: &str| std::env::var(name).ok()),
|
||||
credential_provider: None,
|
||||
};
|
||||
let client = connect(resolved.clone(), services(), None).expect("build live Redis client");
|
||||
let marker = std::time::SystemTime::now()
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name
|
||||
|
||||
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`)
|
||||
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`), `redis_credentials` (what a caller's redis-py `CredentialProvider` authenticates with)
|
||||
- A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here
|
||||
- Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers
|
||||
- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green
|
||||
|
|
|
|||
19
litellm/rust_bridge/host/redis_credentials.py
Normal file
19
litellm/rust_bridge/host/redis_credentials.py
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
"""Credentials from a redis-py ``CredentialProvider`` an SDK caller passed in ``cache_kwargs``, asked the way
|
||||
redis-py's async connection asks it."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Protocol
|
||||
|
||||
|
||||
class CredentialProvider(Protocol):
|
||||
async def get_credentials_async(self) -> tuple[str] | tuple[str, str]: ...
|
||||
|
||||
|
||||
def credentials(provider: CredentialProvider) -> tuple[str | None, str]:
|
||||
match asyncio.run(provider.get_credentials_async()):
|
||||
case (password,):
|
||||
return None, password
|
||||
case (username, password):
|
||||
return username, password
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
import gc
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Final
|
||||
|
|
@ -26,3 +27,21 @@ def test_concurrent_first_connects_share_one_client_without_deadlocking() -> Non
|
|||
|
||||
assert all(client is not None for client in clients)
|
||||
assert len(native_redis_breaker_states()) == before + 1
|
||||
|
||||
|
||||
def _breaker_states_while_connected(namespace: str) -> int:
|
||||
client: Final = NativeRedisClient.connect({"host": "127.0.0.1", "port": 1}, namespace=namespace)
|
||||
assert client is not None
|
||||
return len(native_redis_breaker_states())
|
||||
|
||||
|
||||
def test_a_discarded_client_leaves_the_process_registry(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true")
|
||||
gc.collect()
|
||||
before: Final = len(native_redis_breaker_states())
|
||||
|
||||
while_connected: Final = _breaker_states_while_connected(f"discarded-{uuid.uuid4()}")
|
||||
gc.collect()
|
||||
|
||||
assert while_connected == before + 1
|
||||
assert len(native_redis_breaker_states()) == before
|
||||
|
|
|
|||
38
tests/unit/rust_bridge/host/test_redis_credentials.py
Normal file
38
tests/unit/rust_bridge/host/test_redis_credentials.py
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
import asyncio
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from redis.credentials import CredentialProvider
|
||||
|
||||
from litellm.rust_bridge.host.redis_credentials import credentials
|
||||
|
||||
|
||||
class _SyncOnly(CredentialProvider):
|
||||
def __init__(self, answer: tuple[str] | tuple[str, str]) -> None:
|
||||
self.answer: Final = answer
|
||||
|
||||
def get_credentials(self) -> tuple[str] | tuple[str, str]:
|
||||
return self.answer
|
||||
|
||||
|
||||
class _AsyncOnly(CredentialProvider):
|
||||
def get_credentials(self) -> tuple[str, str]:
|
||||
raise AssertionError("redis-py's async connection asks get_credentials_async")
|
||||
|
||||
async def get_credentials_async(self) -> tuple[str, str]:
|
||||
await asyncio.sleep(0)
|
||||
return "async-user", "async-password"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("provider", "expected"),
|
||||
[
|
||||
(_SyncOnly(("token",)), (None, "token")),
|
||||
(_SyncOnly(("user", "password")), ("user", "password")),
|
||||
(_AsyncOnly(), ("async-user", "async-password")),
|
||||
],
|
||||
)
|
||||
def test_credentials_are_what_redis_py_would_authenticate_with(
|
||||
provider: CredentialProvider, expected: tuple[str | None, str]
|
||||
) -> None:
|
||||
assert credentials(provider) == expected
|
||||
Loading…
Add table
Reference in a new issue