diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 636a5031e5f..50a4e2151f3 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -104,9 +104,21 @@ jobs: rust-test: runs-on: ubuntu-latest timeout-minutes: 30 + # The gateway's shared-state tests run against this Redis. + services: + redis: + image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f + ports: + - 6379:6379 + options: >- + --health-cmd "redis-cli ping" + --health-interval 10s + --health-timeout 5s + --health-retries 5 # Trim debuginfo to keep the job's target dir within the runner's disk env: CARGO_PROFILE_DEV_DEBUG: line-tables-only + LITELLM_TEST_REDIS_URL: redis://127.0.0.1:6379 defaults: run: working-directory: litellm-rust diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5c8811d749d..c47e5b78af0 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -876,6 +876,15 @@ dependencies = [ "time", ] +[[package]] +name = "backon" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" +dependencies = [ + "fastrand", +] + [[package]] name = "base16ct" version = "0.2.0" @@ -1248,7 +1257,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" dependencies = [ "bytes", + "futures-core", "memchr", + "pin-project-lite", + "tokio", + "tokio-util", ] [[package]] @@ -3750,8 +3763,10 @@ dependencies = [ "litellm-http", "litellm-inference", "litellm-llms", + "litellm-router", "litellm-secrets", "litellm-tracing", + "redis", "rstest", "serde", "serde_json", @@ -4299,6 +4314,7 @@ dependencies = [ "litellm-router-types", "rand 0.8.7", "rstest", + "sha2 0.10.9", "strum", "thiserror 2.0.19", "tokio", @@ -5773,18 +5789,29 @@ version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed" dependencies = [ + "arc-swap", "arcstr", + "async-lock", + "backon", + "bytes", + "cfg-if", "combine", "crc16", + "futures-channel", + "futures-util", "itoa", "num-bigint 0.5.1", "percent-encoding", + "pin-project-lite", "rand 0.10.2", "rustls 0.23.42", "rustls-native-certs", "ryu", "sha1_smol", "socket2 0.6.5", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util", "url", "xxhash-rust", ] diff --git a/litellm-rust/crates/gateway-auth/src/authentication.rs b/litellm-rust/crates/gateway-auth/src/authentication.rs index c63b877942e..d3ee4f6a39b 100644 --- a/litellm-rust/crates/gateway-auth/src/authentication.rs +++ b/litellm-rust/crates/gateway-auth/src/authentication.rs @@ -121,7 +121,7 @@ impl Auth { } /// Also accepts virtual keys from `keys`; the configured authenticator is tried first. - pub fn with_virtual_keys(self, keys: crate::CachedKeyLookup) -> Self { + pub fn with_virtual_keys(self, keys: Arc) -> Self { Self { authenticator: Arc::new(crate::VirtualKeyAuthenticator::new( self.authenticator.clone(), diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 89d2828563f..99042aa6d81 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -12,6 +12,15 @@ impl KeyHash { Self(hash_token(token.expose())) } + /// A hash as the proxy stores and broadcasts it: 64 lowercase hex digits. + pub fn from_hashed(hashed: &str) -> Option { + (hashed.len() == 64 + && hashed + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))) + .then(|| Self(hashed.to_owned())) + } + pub fn as_str(&self) -> &str { &self.0 } @@ -321,8 +330,43 @@ pub struct VerifiedKey { pub type LookupFuture<'a> = Pin, KeyError>> + Send + 'a>>; +pub type EntryFuture<'a> = Pin> + Send + 'a>>; + +/// A key's status and the rows it was read from besides its own, whose changes make it stale. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct KeyEntry { + pub status: Option, + pub team_id: Option, + pub user_id: Option, +} + +impl From> for KeyEntry { + fn from(status: Option) -> Self { + let team_id = match &status { + Some(KeyStatus::Active { team_id, .. }) => team_id.clone(), + Some(KeyStatus::Refused { + refusal: Refusal::TeamMissing { team_id } | Refusal::TeamBlocked { team_id }, + .. + }) => Some(team_id.clone()), + _ => None, + }; + Self { + status, + team_id, + user_id: None, + } + } +} + pub trait KeyLookup: Send + Sync { fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a>; + + /// [`KeyLookup::lookup`] with the team and user the answer depends on. The default reads + /// the team from the status alone; a store that knows both should say so, so that a cache + /// can drop the answer when either changes. + fn lookup_entry<'a>(&'a self, hash: &'a KeyHash) -> EntryFuture<'a> { + Box::pin(async move { self.lookup(hash).await.map(KeyEntry::from) }) + } } pub async fn verify_key( diff --git a/litellm-rust/crates/gateway-auth/src/virtual_keys.rs b/litellm-rust/crates/gateway-auth/src/virtual_keys.rs index 48bbf3946d5..6f7c45fb571 100644 --- a/litellm-rust/crates/gateway-auth/src/virtual_keys.rs +++ b/litellm-rust/crates/gateway-auth/src/virtual_keys.rs @@ -12,7 +12,7 @@ use moka::Expiry; use crate::{ AuthFuture, Authentication, AuthenticationMethod, Authenticator, Clock, Credential, Error, KeyError, KeyGrant, Permissions, Principal, PrincipalKind, VerifiedIdentity, - keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key}, + keys::{KeyEntry, KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key}, }; pub(crate) const AUTHORITY: &str = "litellm:virtual-key"; @@ -44,7 +44,7 @@ impl Default for KeyCacheSettings { #[derive(Clone)] struct Cached { - status: Option, + entry: KeyEntry, ttl: Duration, } @@ -62,14 +62,20 @@ impl Expiry for PerEntryTtl { /// for one key share a single store read. The last answer per key is also kept for `max_stale`: /// when the store fails, that answer is served and further reads pause for `negative_ttl`, so keys /// seen recently keep working through an outage and requests do not queue on a dead store. +/// +/// [`CachedKeyLookup::invalidate`] drops answers the proxy reports changed, so a key revoked or +/// edited there stops being served from here before its TTL runs out. pub struct CachedKeyLookup { store: Arc, fresh: moka::future::Cache, - stale: moka::sync::Cache>, + stale: moka::sync::Cache, settings: KeyCacheSettings, started: Instant, /// Milliseconds after `started` until which the store is not read; 0 when healthy. paused_until_ms: AtomicU64, + /// Bumped by every invalidation, so a store read that overlapped one is not cached: it may + /// have read the row before the change it reports. + generation: AtomicU64, } impl CachedKeyLookup { @@ -79,17 +85,58 @@ impl CachedKeyLookup { fresh: moka::future::Cache::builder() .max_capacity(settings.capacity) .expire_after(PerEntryTtl) + .support_invalidation_closures() .build(), stale: moka::sync::Cache::builder() .max_capacity(settings.capacity) .time_to_live(settings.max_stale) + .support_invalidation_closures() .build(), settings, started: Instant::now(), paused_until_ms: AtomicU64::new(0), + generation: AtomicU64::new(0), } } + /// Drops the cached answers that a proxy auth cache invalidation for `cache_key` makes + /// stale: a key's hash drops that key, `team_id:` every key of the team, and + /// `team_membership::` or a bare user id every key of the user. The last known + /// answer is dropped too, so an outage cannot bring a revoked key back. + pub async fn invalidate(&self, cache_key: &str) { + self.generation.fetch_add(1, Ordering::Relaxed); + if let Some(hash) = KeyHash::from_hashed(cache_key) { + self.fresh.invalidate(&hash).await; + self.stale.invalidate(&hash); + return; + } + let owner = match cache_key.strip_prefix("team_id:") { + Some(team_id) => Owner::Team(team_id.to_owned()), + None => Owner::User(match cache_key.strip_prefix("team_membership:") { + Some(membership) => membership + .rsplit_once(':') + .map_or(membership, |(user_id, _)| user_id) + .to_owned(), + None => cache_key.to_owned(), + }), + }; + let stale_owner = owner.clone(); + // Fails only without invalidation closures, which `new` enables. + let _ = self + .fresh + .invalidate_entries_if(move |_, cached| owner.owns(&cached.entry)); + let _ = self + .stale + .invalidate_entries_if(move |_, entry| stale_owner.owns(entry)); + } + + /// Drops every fresh answer, for when invalidations may have been missed. The last known + /// answers stay to serve through a store outage. + pub fn invalidate_all(&self) { + self.generation.fetch_add(1, Ordering::Relaxed); + self.fresh.invalidate_all(); + } + fn elapsed_ms(&self) -> u64 { u64::try_from(self.started.elapsed().as_millis()).unwrap_or(u64::MAX) } @@ -113,29 +160,34 @@ impl CachedKeyLookup { } async fn read_store(&self, hash: &KeyHash) -> Result { - let result = tokio::time::timeout(self.settings.lookup_timeout, self.store.lookup(hash)) - .await - .unwrap_or_else(|elapsed| Err(KeyError::Lookup(Box::new(elapsed)))); + let generation = self.generation.load(Ordering::Relaxed); + let result = + tokio::time::timeout(self.settings.lookup_timeout, self.store.lookup_entry(hash)) + .await + .unwrap_or_else(|elapsed| Err(KeyError::Lookup(Box::new(elapsed)))); match result { - Ok(status) => { + Ok(entry) => { self.paused_until_ms.store(0, Ordering::Relaxed); - self.stale.insert(hash.clone(), status.clone()); - Ok(Cached { - ttl: self.ttl(&status), - status, - }) + let ttl = match self.generation.load(Ordering::Relaxed) == generation { + true => self.ttl(&entry.status), + false => Duration::ZERO, + }; + if !ttl.is_zero() { + self.stale.insert(hash.clone(), entry.clone()); + } + Ok(Cached { entry, ttl }) } Err(error) => { self.pause(); match self.stale.get(hash) { - Some(status) => { + Some(entry) => { tracing::warn!( error = %error, "virtual key store unavailable, serving the last known answer" ); // Cache it briefly so the next read retries the store soon after it returns. Ok(Cached { - status, + entry, ttl: self.settings.negative_ttl, }) } @@ -146,22 +198,38 @@ impl CachedKeyLookup { } } +/// Whose change an invalidation reports. +#[derive(Clone)] +enum Owner { + Team(String), + User(String), +} + +impl Owner { + fn owns(&self, entry: &KeyEntry) -> bool { + match self { + Self::Team(team_id) => entry.team_id.as_ref() == Some(team_id), + Self::User(user_id) => entry.user_id.as_ref() == Some(user_id), + } + } +} + impl KeyLookup for CachedKeyLookup { fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> { Box::pin(async move { if let Some(cached) = self.fresh.get(hash).await { - return Ok(cached.status); + return Ok(cached.entry.status); } if self.paused() { return match self.stale.get(hash) { - Some(status) => Ok(status), + Some(entry) => Ok(entry.status), None => Err(KeyError::Lookup("virtual key store unavailable".into())), }; } self.fresh .try_get_with_by_ref(hash, self.read_store(hash)) .await - .map(|cached| cached.status) + .map(|cached| cached.entry.status) .map_err(|error| KeyError::Lookup(Box::new(error))) }) } @@ -170,14 +238,14 @@ impl KeyLookup for CachedKeyLookup { /// Accepts the configured credential (normally the master key) and, failing that, virtual keys. pub struct VirtualKeyAuthenticator { primary: Arc, - keys: CachedKeyLookup, + keys: Arc, clock: Arc, } impl VirtualKeyAuthenticator { pub fn new( primary: Arc, - keys: CachedKeyLookup, + keys: Arc, clock: Arc, ) -> Self { Self { @@ -206,7 +274,7 @@ impl Authenticator for VirtualKeyAuthenticator { let Credential::Token(token) = credential else { return Err(Error::InvalidToken); }; - let key = verify_key(&self.keys, token, || self.clock.now()).await?; + let key = verify_key(self.keys.as_ref(), token, || self.clock.now()).await?; let id = key.hash.as_str().to_owned(); Ok(VerifiedIdentity { principal: Principal::new(AUTHORITY.into(), id.clone(), PrincipalKind::Service), @@ -417,7 +485,7 @@ mod tests { Arc::new(NoAdditionalPolicy), Arc::new(SystemClock), ) - .with_virtual_keys(CachedKeyLookup::new(store.clone(), settings())) + .with_virtual_keys(Arc::new(CachedKeyLookup::new(store.clone(), settings()))) } fn model(name: &str) -> AccessRequest { @@ -669,6 +737,124 @@ mod tests { assert!(call(&auth, "sk-all", "gpt-a").await.is_ok()); } + /// A store whose keys belong to `user-1` on `team-1`, counting reads per key. + #[derive(Default)] + struct Owned { + reads: Mutex>, + blocked: AtomicBool, + } + + impl KeyLookup for Owned { + fn lookup<'a>(&'a self, _: &'a KeyHash) -> LookupFuture<'a> { + unreachable!("the cache reads entries") + } + + fn lookup_entry<'a>(&'a self, hash: &'a KeyHash) -> crate::keys::EntryFuture<'a> { + Box::pin(async move { + *self.reads.lock().unwrap().entry(hash.clone()).or_default() += 1; + tokio::task::yield_now().await; + Ok(KeyEntry { + status: Some(match self.blocked.load(Ordering::SeqCst) { + true => KeyStatus::Blocked, + false => active(&[], None), + }), + team_id: Some("team-1".into()), + user_id: Some("user-1".into()), + }) + }) + } + } + + fn hash(token: &str) -> KeyHash { + KeyHash::from_token(&SecretValue::new(token)) + } + + #[rstest] + #[case::key_hash(None, &[true, false])] + #[case::team(Some("team_id:team-1"), &[true, true])] + #[case::user(Some("user-1"), &[true, true])] + #[case::membership(Some("team_membership:user-1:team-1"), &[true, true])] + #[case::other_team(Some("team_id:team-2"), &[false, false])] + #[case::other_user(Some("user-2"), &[false, false])] + #[tokio::test] + async fn invalidation_drops_the_keys_it_names( + #[case] cache_key: Option<&str>, + #[case] reread: &[bool], + ) { + let store = Arc::new(Owned::default()); + let keys = CachedKeyLookup::new(store.clone(), KeyCacheSettings::default()); + let tokens = ["sk-a", "sk-b"]; + for token in tokens { + keys.lookup(&hash(token)).await.unwrap(); + } + + let first = hash("sk-a"); + keys.invalidate(cache_key.unwrap_or(first.as_str())).await; + for token in tokens { + keys.lookup(&hash(token)).await.unwrap(); + } + + let reads = store.reads.lock().unwrap(); + let expected: Vec = reread.iter().map(|again| 1 + usize::from(*again)).collect(); + assert_eq!( + tokens.map(|token| reads[&hash(token)]).to_vec(), + expected, + "{cache_key:?}" + ); + } + + #[rstest] + #[tokio::test] + async fn a_revoked_key_is_refused_as_soon_as_it_is_invalidated() { + let store = Arc::new(Owned::default()); + let keys = CachedKeyLookup::new(store.clone(), KeyCacheSettings::default()); + let key = hash("sk-a"); + assert!(matches!( + keys.lookup(&key).await, + Ok(Some(KeyStatus::Active { .. })) + )); + + store.blocked.store(true, Ordering::SeqCst); + keys.invalidate(key.as_str()).await; + + assert_eq!(keys.lookup(&key).await.unwrap(), Some(KeyStatus::Blocked)); + } + + #[rstest] + #[tokio::test] + async fn invalidate_all_rereads_every_key() { + let store = Arc::new(Owned::default()); + let keys = CachedKeyLookup::new(store.clone(), KeyCacheSettings::default()); + keys.lookup(&hash("sk-a")).await.unwrap(); + + keys.invalidate_all(); + keys.lookup(&hash("sk-a")).await.unwrap(); + + assert_eq!(store.reads.lock().unwrap()[&hash("sk-a")], 2); + } + + #[rstest] + #[tokio::test] + async fn a_read_overlapping_an_invalidation_is_not_cached() { + let store = Arc::new(Owned::default()); + let keys = Arc::new(CachedKeyLookup::new( + store.clone(), + KeyCacheSettings::default(), + )); + let key = hash("sk-a"); + let invalidating = { + let keys = keys.clone(); + let key = key.clone(); + async move { keys.invalidate(key.as_str()).await } + }; + // The invalidation lands while the read waits on the store. + let (read, ()) = tokio::join!(keys.lookup(&key), invalidating); + read.unwrap(); + keys.lookup(&key).await.unwrap(); + + assert_eq!(store.reads.lock().unwrap()[&key], 2); + } + #[rstest] #[tokio::test] async fn slow_store_reads_time_out() { @@ -694,7 +880,7 @@ mod tests { async fn rejects_transport_credentials(store: Arc) { let authenticator = VirtualKeyAuthenticator::new( Arc::new(Master), - CachedKeyLookup::new(store, settings()), + Arc::new(CachedKeyLookup::new(store, settings())), Arc::new(SystemClock), ); let parts = Request::new(()).into_parts().0; diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 4f85034280b..d74866b623e 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -22,6 +22,8 @@ litellm-config.workspace = true litellm-db.workspace = true litellm-http = { workspace = true, features = ["mcp"] } litellm-llms.workspace = true +litellm-router.workspace = true +redis = { version = "1.7.0", default-features = false, features = ["connection-manager", "tokio-comp", "tokio-rustls-comp"] } litellm-secrets.workspace = true litellm-tracing.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/gateway/src/coordination.rs b/litellm-rust/crates/gateway/src/coordination.rs new file mode 100644 index 00000000000..eeb01158e59 --- /dev/null +++ b/litellm-rust/crates/gateway/src/coordination.rs @@ -0,0 +1,663 @@ +//! State shared with the gateway's other replicas, and with Python proxies, through the +//! coordination Redis the Python proxy uses. +//! +//! Nothing here sits on the path of a call that succeeds: cooldowns are written when a +//! deployment cools down and read by a background poll, `allowed_fails` counts are only +//! touched by a failed attempt, and key cache invalidations arrive by subscription. While +//! Redis is unreachable each replica keeps its own state, as it does without Redis. + +use std::{ + sync::{ + Arc, OnceLock, + atomic::{AtomicBool, Ordering}, + }, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use futures_util::StreamExt; +use litellm_config::{Config, Object, Value}; +use litellm_gateway_auth::CachedKeyLookup; +use litellm_gateway_inference::ModelRouter; +use litellm_router::{CountFuture, Failure, SharedHealth}; +use litellm_secrets::source::SecretSource; +use redis::aio::{ConnectionManager, ConnectionManagerConfig}; +use tokio_util::sync::CancellationToken; + +use crate::Error; + +/// The Python proxy's channel for auth cache invalidations (`auth_cache_invalidation_pubsub`). +const AUTH_CACHE_INVALIDATION_CHANNEL: &str = "litellm_proxy.auth_cache_invalidation"; +/// How often cooldowns written by other replicas are read, as Python's +/// `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS`. +const DEFAULT_COOLDOWN_READ_INTERVAL: Duration = Duration::from_secs(1); +/// Upper bound on the `allowed_fails` round trip a failed attempt waits for before it counts +/// this replica's failures alone. +const COUNT_TIMEOUT: Duration = Duration::from_millis(250); +const COMMAND_TIMEOUT: Duration = Duration::from_secs(2); +const RECONNECT_MIN: Duration = Duration::from_secs(1); +const RECONNECT_MAX: Duration = Duration::from_secs(30); +const MGET_CHUNK: usize = 500; + +/// Where the coordination Redis is, resolved as the Python proxy resolves it. +#[derive(Clone, PartialEq, Eq)] +pub struct RedisSettings { + url: String, + namespace: Option, + cooldown_read_interval: Duration, +} + +impl std::fmt::Debug for RedisSettings { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RedisSettings") + .field("namespace", &self.namespace) + .finish_non_exhaustive() + } +} + +impl RedisSettings { + /// A Redis at `url`, with keys and channels under `namespace` as the proxy's `namespace` + /// cache param puts them. + pub fn new(url: impl Into, namespace: Option) -> Self { + Self { + url: url.into(), + namespace, + cooldown_read_interval: read_interval(), + } + } + + /// How often cooldowns other replicas wrote are read; one second by default, or + /// `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS` as for the Python router. + pub fn with_cooldown_read_interval(self, interval: Duration) -> Self { + Self { + cooldown_read_interval: interval, + ..self + } + } + + /// `general_settings.coordination_redis`, else the response cache's Redis + /// (`litellm_settings.cache_params` with `type: redis` when `cache` is on), else the + /// `REDIS_*` environment variables. `None` when none names a Redis. + pub async fn resolve( + config: &Config, + secrets: &dyn SecretSource, + ) -> Result, Error> { + if let Some(block) = &config.general_settings.coordination_redis { + let settings = Self::from_params(block, "general_settings.coordination_redis", secrets) + .await? + .ok_or(Error::Redis( + "general_settings.coordination_redis needs a host or url".into(), + ))?; + return Ok(Some(settings)); + } + let litellm = &config.litellm_settings; + if litellm.cache == Some(true) + && let Some(params) = &litellm.cache_params + && params.get("type").and_then(Value::as_str) == Some("redis") + && let Some(settings) = + Self::from_params(params, "litellm_settings.cache_params", secrets).await? + { + return Ok(Some(settings)); + } + Self::from_environment(secrets).await + } + + async fn from_params( + params: &Object, + setting: &str, + secrets: &dyn SecretSource, + ) -> Result, Error> { + for unsupported in [ + "startup_nodes", + "redis_startup_nodes", + "sentinel_nodes", + "service_name", + ] { + if params + .get(unsupported) + .is_some_and(|value| !value.is_null()) + { + return Err(Error::Redis(format!( + "{setting}.{unsupported}: Redis Cluster and Sentinel are not supported by this gateway yet" + ))); + } + } + let mut fields = Fields::default(); + for (name, slot) in [ + ("url", &mut fields.url), + ("host", &mut fields.host), + ("port", &mut fields.port), + ("username", &mut fields.username), + ("password", &mut fields.password), + ("ssl", &mut fields.ssl), + ("db", &mut fields.db), + ] { + *slot = match params.get(name) { + Some(value) => scalar(value, secrets).await?, + None => None, + }; + } + let namespace = match params.get("namespace") { + Some(value) => scalar(value, secrets).await?, + None => None, + }; + Ok(fields.url().map(|url| Self::new(url, namespace))) + } + + async fn from_environment(secrets: &dyn SecretSource) -> Result, Error> { + for unsupported in ["REDIS_CLUSTER_NODES", "REDIS_SENTINEL_NODES"] { + if secrets.get_secret_str(unsupported).await?.is_some() { + return Err(Error::Redis(format!( + "{unsupported}: Redis Cluster and Sentinel are not supported by this gateway yet" + ))); + } + } + let read = async |name: &str| -> Result, Error> { + Ok(secrets + .get_secret_str(name) + .await? + .map(|value| value.expose().to_owned()) + .filter(|value| !value.trim().is_empty())) + }; + let fields = Fields { + url: read("REDIS_URL").await?, + host: read("REDIS_HOST").await?, + port: read("REDIS_PORT").await?, + username: read("REDIS_USERNAME").await?, + password: read("REDIS_PASSWORD").await?, + ssl: read("REDIS_SSL").await?, + db: read("REDIS_DB").await?, + }; + Ok(fields.url().map(|url| Self::new(url, None))) + } + + fn key(&self, key: &str) -> String { + match &self.namespace { + Some(namespace) => format!("{namespace}:{key}"), + None => key.to_owned(), + } + } +} + +#[derive(Default)] +struct Fields { + url: Option, + host: Option, + port: Option, + username: Option, + password: Option, + ssl: Option, + db: Option, +} + +impl Fields { + /// `url` when given, else one built from `host`; `None` when neither is set. + fn url(self) -> Option { + if let Some(url) = self.url { + return Some(url); + } + let host = self.host?; + let scheme = match self.ssl.as_deref().map(str::to_ascii_lowercase).as_deref() { + Some("true" | "1" | "yes") => "rediss", + _ => "redis", + }; + let credentials = match (self.username, self.password) { + (None, None) => String::new(), + (username, password) => format!( + "{}:{}@", + encode(username.as_deref().unwrap_or_default()), + encode(password.as_deref().unwrap_or_default()) + ), + }; + let port = self.port.unwrap_or_else(|| "6379".into()); + let db = self.db.map(|db| format!("/{db}")).unwrap_or_default(); + Some(format!("{scheme}://{credentials}{host}:{port}{db}")) + } +} + +/// Percent-encodes a URL userinfo part. +fn encode(part: &str) -> String { + part.bytes() + .map(|byte| match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => { + char::from(byte).to_string() + } + byte => format!("%{byte:02X}"), + }) + .collect() +} + +/// A config scalar as text, resolving an `os.environ/` reference. +async fn scalar(value: &Value, secrets: &dyn SecretSource) -> Result, Error> { + let text = match value { + Value::Null => return Ok(None), + Value::String(text) => text.clone(), + Value::Bool(flag) => flag.to_string(), + Value::Number(number) => number.to_string(), + _ => return Err(Error::Redis("Redis settings must be scalars".into())), + }; + let text = match text.strip_prefix("os.environ/") { + Some(name) => match secrets.get_secret_str(name).await? { + Some(value) => value.expose().to_owned(), + None => return Err(Error::Redis(format!("{name} is not set"))), + }, + None => text, + }; + Ok(Some(text).filter(|text| !text.trim().is_empty())) +} + +/// The coordination Redis: a client and, once it has connected, a multiplexed connection that +/// reconnects on its own. +pub struct Coordination { + settings: RedisSettings, + client: redis::Client, + connection: OnceLock, + subscribed: AtomicBool, +} + +impl Coordination { + pub fn new(settings: RedisSettings) -> Result, Error> { + let client = redis::Client::open(settings.url.as_str()) + .map_err(|error| Error::Redis(format!("invalid coordination Redis: {error}")))?; + Ok(Arc::new(Self { + settings, + client, + connection: OnceLock::new(), + subscribed: AtomicBool::new(false), + })) + } + + fn connection(&self) -> Option { + self.connection.get().cloned() + } + + /// Whether the connection shared state is written and read through is up yet. + pub fn connected(&self) -> bool { + self.connection.get().is_some() + } + + /// Whether key cache invalidations are being followed right now. + pub fn subscribed(&self) -> bool { + self.subscribed.load(Ordering::Relaxed) + } + + /// Connects in the background, retrying until Redis answers, then starts the cooldown poll + /// and, with `keys`, the key cache invalidation subscriber. Returns at once. + pub fn start( + self: &Arc, + router: Arc, + keys: Option>, + shutdown: CancellationToken, + ) { + let coordination = self.clone(); + let cooldowns = shutdown.clone(); + tokio::spawn(async move { + let connected = tokio::select! { + connected = coordination.connect() => connected, + () = cooldowns.cancelled() => return, + }; + let _ = coordination.connection.set(connected); + tracing::info!("coordination Redis connected: sharing deployment cooldowns"); + tokio::select! { + () = coordination.poll_cooldowns(router) => (), + () = cooldowns.cancelled() => (), + } + }); + if let Some(keys) = keys { + let coordination = self.clone(); + tokio::spawn(async move { + tokio::select! { + () = coordination.subscribe(&keys) => (), + () = shutdown.cancelled() => (), + } + }); + } + } + + async fn connect(&self) -> ConnectionManager { + let config = ConnectionManagerConfig::new() + .set_connection_timeout(Some(COMMAND_TIMEOUT)) + .set_response_timeout(Some(COMMAND_TIMEOUT)); + let mut backoff = RECONNECT_MIN; + loop { + match self + .client + .get_connection_manager_with_config(config.clone()) + .await + { + Ok(connection) => return connection, + Err(error) => { + tracing::warn!(%error, retry_secs = backoff.as_secs(), "coordination Redis unreachable; replicas keep their own router health until it answers"); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(RECONNECT_MAX); + } + } + } + } + + /// Reads every deployment's cooldown key each interval and applies the active ones. + async fn poll_cooldowns(&self, router: Arc) { + let router = router.router(); + let keys: Vec = router + .deployment_ids() + .map(|id| self.settings.key(&cooldown_key(id))) + .collect(); + let mut interval = tokio::time::interval(self.settings.cooldown_read_interval); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + interval.tick().await; + let Some(mut connection) = self.connection() else { + continue; + }; + for (chunk_index, chunk) in keys.chunks(MGET_CHUNK).enumerate() { + let values: Vec> = match redis::cmd("MGET") + .arg(chunk) + .query_async(&mut connection) + .await + { + Ok(values) => values, + Err(error) => { + tracing::debug!(%error, "reading shared cooldowns failed"); + break; + } + }; + let now = epoch_seconds(); + for (offset, value) in values.into_iter().enumerate() { + if let Some(remaining) = + value.as_deref().and_then(|value| remaining(value, now)) + { + router.cool_down(chunk_index * MGET_CHUNK + offset, remaining); + } + } + } + } + } + + /// Follows the proxy's auth cache invalidations, dropping what each names from `keys`. + /// Every (re)subscription first drops all fresh answers, since invalidations published + /// while unsubscribed were missed. + async fn subscribe(&self, keys: &CachedKeyLookup) { + let channel = self.settings.key(AUTH_CACHE_INVALIDATION_CHANNEL); + let mut backoff = RECONNECT_MIN; + loop { + match self.client.get_async_pubsub().await { + Ok(mut pubsub) => match pubsub.subscribe(&channel).await { + Ok(()) => { + backoff = RECONNECT_MIN; + keys.invalidate_all(); + self.subscribed.store(true, Ordering::Relaxed); + tracing::info!(%channel, "following proxy virtual key invalidations"); + let mut messages = pubsub.into_on_message(); + while let Some(message) = messages.next().await { + if let Some(cache_key) = invalidated(message.get_payload_bytes()) { + keys.invalidate(&cache_key).await; + } + } + self.subscribed.store(false, Ordering::Relaxed); + tracing::warn!( + "virtual key invalidation subscription dropped; resubscribing" + ); + } + Err(error) => { + tracing::warn!(%error, "subscribing to virtual key invalidations failed") + } + }, + Err(error) => { + tracing::warn!(%error, retry_secs = backoff.as_secs(), "coordination Redis unreachable; virtual keys refresh on their TTL until it answers") + } + } + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(RECONNECT_MAX); + } + } +} + +/// What the cooldown poll needs from the gateway: its router. +pub trait RouterHandle: Send + Sync + 'static { + fn router(&self) -> &ModelRouter; +} + +impl RouterHandle for litellm_gateway_inference::Gateway { + fn router(&self) -> &ModelRouter { + &self.models + } +} + +/// Shares cooldowns and `allowed_fails` counts through the coordination Redis, under the keys +/// and in the format the Python router uses, so replicas and Python routers that name +/// deployments by the same `model_info.id` see each other's. +pub struct RedisHealth(pub Arc); + +impl SharedHealth for RedisHealth { + fn cooled(&self, deployment: &str, cooldown: Duration, failure: Failure) { + let Some(mut connection) = self.0.connection() else { + return; + }; + let key = self.0.settings.key(&cooldown_key(deployment)); + let value = serde_json::json!({ + "exception_received": format!("litellm-gateway: {failure:?}"), + "status_code": match failure { + Failure::RateLimited => "429", + Failure::Rejected => "401", + Failure::Unavailable => "503", + }, + "timestamp": epoch_seconds(), + "cooldown_time": cooldown.as_secs_f64(), + }) + .to_string(); + let milliseconds = u64::try_from(cooldown.as_millis()) + .unwrap_or(u64::MAX) + .max(1); + tokio::spawn(async move { + let result: redis::RedisResult<()> = redis::cmd("SET") + .arg(&key) + .arg(value) + .arg("PX") + .arg(milliseconds) + .query_async(&mut connection) + .await; + if let Err(error) = result { + tracing::warn!(%error, "sharing a deployment cooldown failed"); + } + }); + } + + fn count_failure<'a>(&'a self, deployment: &'a str, window: Duration) -> CountFuture<'a> { + Box::pin(async move { + let mut connection = self.0.connection()?; + let key = self + .0 + .settings + .key(&format!("deployment:{deployment}:allowed_fails")); + let milliseconds = u64::try_from(window.as_millis()).unwrap_or(u64::MAX).max(1); + let mut pipeline = redis::pipe(); + pipeline + .cmd("SET") + .arg(&key) + .arg(0) + .arg("PX") + .arg(milliseconds) + .arg("NX") + .ignore() + .cmd("INCR") + .arg(&key); + let count = pipeline.query_async::<(u64,)>(&mut connection); + match tokio::time::timeout(COUNT_TIMEOUT, count).await { + Ok(Ok((count,))) => Some(u32::try_from(count).unwrap_or(u32::MAX)), + Ok(Err(error)) => { + tracing::debug!(%error, "shared allowed_fails count failed"); + None + } + Err(_) => None, + } + }) + } +} + +/// Python's `CooldownCache.get_cooldown_cache_key`. +fn cooldown_key(deployment: &str) -> String { + format!("deployment:{deployment}:cooldown") +} + +fn read_interval() -> Duration { + std::env::var("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS") + .ok() + .and_then(|seconds| seconds.trim().parse::().ok()) + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .filter(|interval| !interval.is_zero()) + .unwrap_or(DEFAULT_COOLDOWN_READ_INTERVAL) +} + +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64() +} + +/// How much of a stored cooldown is left at `now`, or `None` when it is over or unreadable. +/// Reads both this gateway's JSON and the Python dict repr the Python router stores. +fn remaining(value: &str, now: f64) -> Option { + let timestamp = number_after(value, "timestamp")?; + let cooldown_time = number_after(value, "cooldown_time")?; + Duration::try_from_secs_f64(timestamp + cooldown_time - now) + .ok() + .filter(|remaining| !remaining.is_zero()) +} + +/// The number following the quoted `field` and its colon in a flat JSON or Python dict. +fn number_after(text: &str, field: &str) -> Option { + ["\"", "'"].into_iter().find_map(|quote| { + let start = text.find(&format!("{quote}{field}{quote}"))? + field.len() + 2; + let rest = text[start..].trim_start().strip_prefix(':')?.trim_start(); + let end = rest + .find(|character: char| !(character.is_ascii_digit() || "+-.eE".contains(character))) + .unwrap_or(rest.len()); + rest[..end].parse().ok() + }) +} + +/// The cache key an invalidation message drops; `None` for a malformed message or one that +/// sets a value (spend counters), which keys do not cache. +fn invalidated(payload: &[u8]) -> Option { + let message: serde_json::Value = serde_json::from_slice(payload).ok()?; + if message + .get("new_value") + .is_some_and(|value| !value.is_null()) + { + return None; + } + message.get("cache_key")?.as_str().map(str::to_owned) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use futures_util::future::BoxFuture; + use litellm_auth_types::SecretValue; + use rstest::rstest; + + use super::*; + + struct Secrets(HashMap<&'static str, &'static str>); + + impl SecretSource for Secrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { Ok(self.0.get(name).map(|value| SecretValue::new(*value))) }) + } + } + + async fn resolve( + yaml: &str, + environment: &[(&'static str, &'static str)], + ) -> Result)>, String> { + let config = Config::from_yaml(&format!("model_list: []\n{yaml}")).unwrap(); + RedisSettings::resolve(&config, &Secrets(environment.iter().copied().collect())) + .await + .map(|settings| settings.map(|settings| (settings.url, settings.namespace))) + .map_err(|error| error.to_string()) + } + + #[rstest] + #[case::none("", &[], Ok(None))] + #[case::environment_host( + "", + &[("REDIS_HOST", "redis.internal"), ("REDIS_PASSWORD", "p@ss")], + Ok(Some(("redis://:p%40ss@redis.internal:6379".into(), None))) + )] + #[case::environment_url("", &[("REDIS_URL", "rediss://r:6380")], Ok(Some(("rediss://r:6380".into(), None))))] + #[case::coordination_block_wins( + "general_settings:\n coordination_redis: {host: coord, port: 6390, ssl: true, namespace: prod}\n", + &[("REDIS_HOST", "env")], + Ok(Some(("rediss://coord:6390".into(), Some("prod".into())))) + )] + #[case::coordination_reference( + "general_settings:\n coordination_redis: {host: os.environ/COORD_HOST, password: os.environ/COORD_PASSWORD}\n", + &[("COORD_HOST", "coord"), ("COORD_PASSWORD", "secret")], + Ok(Some(("redis://:secret@coord:6379".into(), None))) + )] + #[case::response_cache_redis( + "litellm_settings:\n cache: true\n cache_params: {type: redis, host: cache, port: '6381', namespace: ns}\n", + &[], + Ok(Some(("redis://cache:6381".into(), Some("ns".into())))) + )] + #[case::response_cache_off( + "litellm_settings:\n cache_params: {type: redis, host: cache}\n", + &[("REDIS_HOST", "env")], + Ok(Some(("redis://env:6379".into(), None))) + )] + #[case::cluster_refused( + "", + &[("REDIS_CLUSTER_NODES", "[]")], + Err("coordination Redis: REDIS_CLUSTER_NODES: Redis Cluster and Sentinel are not supported by this gateway yet".into()) + )] + #[case::block_without_target( + "general_settings:\n coordination_redis: {port: 6379}\n", + &[], + Err("coordination Redis: general_settings.coordination_redis needs a host or url".into()) + )] + #[tokio::test] + async fn resolves_the_coordination_redis_like_the_proxy( + #[case] yaml: &str, + #[case] environment: &[(&'static str, &'static str)], + #[case] expected: Result)>, String>, + ) { + assert_eq!(resolve(yaml, environment).await, expected); + } + + #[rstest] + #[case::gateway_json(r#"{"exception_received": "x", "status_code": "429", "timestamp": 100.0, "cooldown_time": 5.0}"#, 102.0, Some(3_000))] + #[case::python_repr( + "{'exception_received': 'RateLimitError', 'status_code': '429', 'timestamp': 1.0e2, 'cooldown_time': 5}", + 104.5, + Some(500) + )] + #[case::over(r#"{"timestamp": 100.0, "cooldown_time": 5.0}"#, 105.0, None)] + #[case::unreadable("garbage", 0.0, None)] + fn reads_cooldowns_written_by_either_gateway( + #[case] value: &str, + #[case] now: f64, + #[case] expected_ms: Option, + ) { + assert_eq!( + remaining(value, now).map(|remaining| remaining.as_millis()), + expected_ms + ); + } + + #[rstest] + #[case::key(br#"{"cache_key": "abc"}"#, Some("abc"))] + #[case::team(br#"{"cache_key": "team_id:t"}"#, Some("team_id:t"))] + #[case::spend_counter( + br#"{"cache_key": "spend:key:abc", "new_value": 1.5, "ttl": 60}"#, + None + )] + #[case::malformed(b"not json", None)] + fn reads_proxy_invalidation_messages(#[case] payload: &[u8], #[case] expected: Option<&str>) { + assert_eq!(invalidated(payload).as_deref(), expected); + } +} diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index 9e8e5c9daef..a2097417ce5 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -20,6 +20,8 @@ pub enum Error { DatabaseUrl(String), #[error("reading a gateway setting from the secret source failed")] Secret(#[from] litellm_secrets::Error), + #[error("coordination Redis: {0}")] + Redis(String), #[error("MCP requires a configured master key")] Auth(#[from] litellm_gateway_auth::Error), } diff --git a/litellm-rust/crates/gateway/src/keys.rs b/litellm-rust/crates/gateway/src/keys.rs index ca1135929fd..e4a92953e84 100644 --- a/litellm-rust/crates/gateway/src/keys.rs +++ b/litellm-rust/crates/gateway/src/keys.rs @@ -5,7 +5,7 @@ use litellm_config::Config; use litellm_db::{PgPool, PoolSettings}; use litellm_gateway_auth::{ KeyError, - keys::{KeyHash, KeyLookup, LookupFuture}, + keys::{EntryFuture, KeyEntry, KeyHash, KeyLookup, LookupFuture}, }; use litellm_secrets::source::SecretSource; @@ -19,11 +19,22 @@ pub struct PostgresKeys { impl KeyLookup for PostgresKeys { fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> { + Box::pin(async move { self.lookup_entry(hash).await.map(|entry| entry.status) }) + } + + fn lookup_entry<'a>(&'a self, hash: &'a KeyHash) -> EntryFuture<'a> { Box::pin(async move { - litellm_db::find_verification_token(&self.pool, hash.as_str()) + let row = litellm_db::find_verification_token(&self.pool, hash.as_str()) .await - .map(|row| row.map(|row| key_status::status(row, self.personal_keys_use_teams))) - .map_err(|error| KeyError::Lookup(Box::new(error))) + .map_err(|error| KeyError::Lookup(Box::new(error)))?; + Ok(match row { + Some(row) => KeyEntry { + team_id: row.team_id.clone(), + user_id: row.user_id.clone(), + status: Some(key_status::status(row, self.personal_keys_use_teams)), + }, + None => KeyEntry::default(), + }) }) } } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 1b80d354e2b..952f7cf9e2c 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -1,4 +1,5 @@ mod auth; +pub mod coordination; pub mod drain; mod error; pub mod health; @@ -152,12 +153,47 @@ pub fn router_with_keys( ui: Option, mcp: Option, keys: Option>, +) -> Router { + let keys = keys.map(|keys| key_cache(keys, config)); + router_with_key_cache(inference, config, ui, mcp, keys) +} + +/// The cache virtual keys are read through: `general_settings.user_api_key_cache_ttl` seconds +/// for a found key, as the proxy's in-memory key cache, 60 by default. +pub fn key_cache(keys: Arc, config: &Config) -> Arc { + let defaults = KeyCacheSettings::default(); + let positive_ttl = match config + .general_settings + .additional_fields + .get("user_api_key_cache_ttl") + { + Some(litellm_config::Value::Number(seconds)) => seconds.as_f64(), + Some(litellm_config::Value::String(seconds)) => seconds.trim().parse().ok(), + _ => None, + } + .and_then(|seconds| std::time::Duration::try_from_secs_f64(seconds).ok()) + .unwrap_or(defaults.positive_ttl); + Arc::new(CachedKeyLookup::new( + keys, + KeyCacheSettings { + positive_ttl, + ..defaults + }, + )) +} + +/// [`router_with_keys`] over an already built key cache, which the caller can keep to +/// invalidate. +pub fn router_with_key_cache( + inference: Arc, + config: &Config, + ui: Option, + mcp: Option, + keys: Option>, ) -> Router { let auth = Auth::from_config(config, inference.secrets.clone()); let auth = match keys { - Some(keys) => { - auth.with_virtual_keys(CachedKeyLookup::new(keys, KeyCacheSettings::default())) - } + Some(keys) => auth.with_virtual_keys(keys), None => auth, }; let inference = litellm_gateway_inference::router(inference) diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 41a322506f6..31c4a741a84 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -1,5 +1,6 @@ use std::{ error::Error, + sync::Arc, time::{SystemTime, UNIX_EPOCH}, }; @@ -90,8 +91,23 @@ async fn main() -> Result<(), Box> { }, ) }); - let keys = litellm_gateway::build_keys(&config, inference.secrets.as_ref()).await?; + let keys = litellm_gateway::build_keys(&config, inference.secrets.as_ref()) + .await? + .map(|keys| litellm_gateway::key_cache(keys, &config)); let virtual_keys = keys.is_some(); + let coordination = + litellm_gateway::coordination::RedisSettings::resolve(&config, inference.secrets.as_ref()) + .await? + .map(litellm_gateway::coordination::Coordination::new) + .transpose()?; + if let Some(coordination) = &coordination { + inference + .models + .share_health(Arc::new(litellm_gateway::coordination::RedisHealth( + coordination.clone(), + ))); + coordination.start(inference.clone(), keys.clone(), shutdown.clone()); + } let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; let address = listener.local_addr()?; // Without TCP_NODELAY, Nagle holds small SSE chunks until the client's delayed ACK (~40 ms per stream). @@ -101,13 +117,14 @@ async fn main() -> Result<(), Box> { } }); - tracing::info!(%address, models = config.model_list.len(), virtual_keys, log_level = %level, "gateway listening"); + let shared_state = coordination.is_some(); + tracing::info!(%address, models = config.model_list.len(), virtual_keys, shared_state, log_level = %level, "gateway listening"); let draining = litellm_gateway::health::Draining::default(); let metrics = litellm_gateway::metrics::Metrics::new(&config, draining.flag()); let metrics_auth = litellm_gateway::metrics_auth(&config, inference.secrets.clone()); let app = litellm_gateway::with_ops( - litellm_gateway::router_with_keys(inference, &config, ui, mcp_router, keys), + litellm_gateway::router_with_key_cache(inference, &config, ui, mcp_router, keys), metrics, draining.clone(), metrics_auth, diff --git a/litellm-rust/crates/gateway/tests/shared_state.rs b/litellm-rust/crates/gateway/tests/shared_state.rs new file mode 100644 index 00000000000..76aa77de3c0 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/shared_state.rs @@ -0,0 +1,248 @@ +//! Replicas sharing state through a live Redis. Set `LITELLM_TEST_REDIS_URL` (for example +//! `redis://127.0.0.1:6379`) to run them; without it they pass without checking anything. + +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Json, Router, + body::Body, + http::{Request, StatusCode}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway::coordination::{Coordination, RedisHealth, RedisSettings}; +use litellm_gateway_auth::keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, ModelAccess}; +use rstest::rstest; +use serde_json::json; +use tokio::net::TcpListener; +use tokio_util::sync::CancellationToken; +use tower::ServiceExt; + +fn redis_url() -> Option { + std::env::var("LITELLM_TEST_REDIS_URL").ok() +} + +fn namespace() -> String { + format!("test-{}", uuid::Uuid::new_v4()) +} + +/// An Anthropic-shaped upstream that answers 429 when `limited`, counting its requests. +async fn upstream(limited: bool) -> (String, Arc) { + let hits = Arc::new(AtomicUsize::new(0)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let counted = hits.clone(); + let app = Router::new().fallback(move || { + let hits = counted.clone(); + async move { + hits.fetch_add(1, Ordering::SeqCst); + if limited { + return ( + StatusCode::TOO_MANY_REQUESTS, + Json(json!({"type": "error", "error": {"type": "rate_limit_error", "message": "slow down"}})), + ) + .into_response(); + } + Json(json!({ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-x", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": {"input_tokens": 1, "output_tokens": 1}, + })) + .into_response() + } + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + (format!("http://{address}"), hits) +} + +/// One gateway replica sharing state under `namespace`, with its key cache when `keys` is set. +async fn replica( + url: &str, + namespace: &str, + config: &Config, + keys: Option>, + shutdown: &CancellationToken, +) -> (Router, Arc) { + let inference = litellm_gateway::build_inference(config).unwrap(); + let coordination = Coordination::new( + RedisSettings::new(url, Some(namespace.to_owned())) + .with_cooldown_read_interval(Duration::from_millis(50)), + ) + .unwrap(); + inference + .models + .share_health(Arc::new(RedisHealth(coordination.clone()))); + let keys = keys.map(|keys| litellm_gateway::key_cache(keys, config)); + coordination.start(inference.clone(), keys.clone(), shutdown.clone()); + let app = litellm_gateway::router_with_key_cache(inference, config, None, None, keys); + (app, coordination) +} + +async fn until(mut done: impl AsyncFnMut() -> bool) { + for _ in 0..100 { + if done().await { + return; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!("condition not met within 5 s"); +} + +async fn send(app: &Router, key: &str) -> StatusCode { + app.clone() + .oneshot( + Request::post("/v1/messages") + .header("authorization", format!("Bearer {key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({"model": "claude", "max_tokens": 1, "messages": [{"role": "user", "content": "hi"}]}) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap() + .status() +} + +#[rstest] +#[tokio::test] +async fn a_deployment_one_replica_cools_down_is_skipped_by_the_others() { + let Some(url) = redis_url() else { return }; + let namespace = namespace(); + let (limited, limited_hits) = upstream(true).await; + let (healthy, _) = upstream(false).await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: {{model: anthropic/claude-x, api_base: '{limited}', api_key: sk-upstream}} + - model_name: claude + litellm_params: {{model: anthropic/claude-x, api_base: '{healthy}', api_key: sk-upstream}} +router_settings: + num_retries: 1 + cooldown_time: 30 +general_settings: + master_key: sk-master +" + )) + .unwrap(); + let shutdown = CancellationToken::new(); + let (first, first_redis) = replica(&url, &namespace, &config, None, &shutdown).await; + let (second, second_redis) = replica(&url, &namespace, &config, None, &shutdown).await; + until(async || first_redis.connected() && second_redis.connected()).await; + + // The first replica meets the 429, cools that deployment and retries on the other. + until(async || { + assert_eq!(send(&first, "sk-master").await, StatusCode::OK); + limited_hits.load(Ordering::SeqCst) > 0 + }) + .await; + tokio::time::sleep(Duration::from_millis(300)).await; + + let hits = limited_hits.load(Ordering::SeqCst); + for _ in 0..50 { + assert_eq!(send(&second, "sk-master").await, StatusCode::OK); + } + assert_eq!(limited_hits.load(Ordering::SeqCst), hits); + shutdown.cancel(); +} + +/// Keys whose statuses a test changes, all on `team-1`. +#[derive(Default)] +struct Keys(Mutex>); + +impl Keys { + fn put(&self, token: &str, status: KeyStatus) { + self.0 + .lock() + .unwrap() + .insert(KeyHash::from_token(&SecretValue::new(token)), status); + } +} + +impl KeyLookup for Keys { + fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> { + Box::pin(async move { Ok(self.0.lock().unwrap().get(hash).cloned()) }) + } +} + +fn on_team() -> KeyStatus { + match KeyStatus::active(None, ModelAccess::default()) { + KeyStatus::Active { + expires_at, + models, + allowed_routes, + .. + } => KeyStatus::Active { + expires_at, + models, + allowed_routes, + team_id: Some("team-1".into()), + }, + status => status, + } +} + +#[rstest] +#[case::key(None)] +#[case::team(Some("team_id:team-1"))] +#[tokio::test] +async fn a_key_the_proxy_invalidates_is_reread_on_every_replica(#[case] cache_key: Option<&str>) { + let Some(url) = redis_url() else { return }; + let namespace = namespace(); + let (healthy, _) = upstream(false).await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: {{model: anthropic/claude-x, api_base: '{healthy}', api_key: sk-upstream}} +general_settings: + master_key: sk-master +" + )) + .unwrap(); + let keys = Arc::new(Keys::default()); + keys.put("sk-user", on_team()); + let shutdown = CancellationToken::new(); + let (app, redis) = replica(&url, &namespace, &config, Some(keys.clone()), &shutdown).await; + // Subscribing drops every cached key, so cache this one after it. + until(async || redis.subscribed()).await; + assert_eq!(send(&app, "sk-user").await, StatusCode::OK); + + keys.put("sk-user", KeyStatus::Blocked); + // Still served from the cache: the key TTL is a minute. + assert_eq!(send(&app, "sk-user").await, StatusCode::OK); + + // What the proxy's `publish_auth_cache_invalidation` sends after a key or team update. + let hash = KeyHash::from_token(&SecretValue::new("sk-user")); + let message = json!({"cache_key": cache_key.unwrap_or(hash.as_str())}).to_string(); + let channel = format!("{namespace}:litellm_proxy.auth_cache_invalidation"); + let mut publisher = redis::Client::open(url.as_str()) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + until(async || { + let _: i64 = redis::cmd("PUBLISH") + .arg(&channel) + .arg(&message) + .query_async(&mut publisher) + .await + .unwrap(); + send(&app, "sk-user").await == StatusCode::UNAUTHORIZED + }) + .await; + shutdown.cancel(); +} diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml index ea35556eec0..5d270ccf293 100644 --- a/litellm-rust/crates/router/Cargo.toml +++ b/litellm-rust/crates/router/Cargo.toml @@ -10,6 +10,7 @@ litellm-config.workspace = true litellm-inference-messages.workspace = true litellm-router-types.workspace = true rand.workspace = true +sha2.workspace = true strum.workspace = true thiserror.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/router/src/deployment.rs b/litellm-rust/crates/router/src/deployment.rs index 800b4931dd5..5d6c2995820 100644 --- a/litellm-rust/crates/router/src/deployment.rs +++ b/litellm-rust/crates/router/src/deployment.rs @@ -18,6 +18,9 @@ pub struct Deployment { /// `model_info.team_id`: the team that owns this deployment, which only that team's keys /// may use. `None` for a deployment every caller shares. pub team_id: Option, + /// `model_info.id`, which names the deployment in state shared with other replicas and + /// with Python routers on the same Redis. + pub id: Option, } impl Deployment { @@ -42,6 +45,7 @@ impl Deployment { litellm_params: params.clone(), shaping: MessagesShaping::default(), team_id: None, + id: None, }) } diff --git a/litellm-rust/crates/router/src/group.rs b/litellm-rust/crates/router/src/group.rs index 3b0ac81699f..0a74d3d30ab 100644 --- a/litellm-rust/crates/router/src/group.rs +++ b/litellm-rust/crates/router/src/group.rs @@ -6,13 +6,15 @@ use crate::{Deployment, RoutingStrategy, health::Health}; #[derive(Debug)] pub(crate) struct Member { + /// The deployment's id in shared state: `model_info.id`, or one derived from the config. + pub id: Box, pub deployment: Deployment, pub health: Arc, weights: [f64; 3], } impl Member { - pub(crate) fn new(deployment: Deployment) -> Self { + pub(crate) fn new(id: Box, deployment: Deployment) -> Self { let params = &deployment.litellm_params; let weights = [ params @@ -25,6 +27,7 @@ impl Member { ] .map(|weight| weight.max(0.0)); Self { + id, deployment, health: Arc::default(), weights, @@ -160,7 +163,13 @@ mod tests { let members = config .model_list .iter() - .map(|model| Member::new(Deployment::from_params(&model.litellm_params).unwrap())) + .enumerate() + .map(|(index, model)| { + Member::new( + index.to_string().into(), + Deployment::from_params(&model.litellm_params).unwrap(), + ) + }) .collect(); ModelGroup::new("group".into(), members) } @@ -168,9 +177,13 @@ mod tests { fn cooling_group(size: usize, cooling: &[usize]) -> ModelGroup { let group = group(&vec![""; size]); for index in cooling { - group.members[*index] - .health - .record_failure(0, Failure::Unavailable, &COOLDOWN, true); + group.members[*index].health.record_failure( + 0, + Failure::Unavailable, + &COOLDOWN, + true, + None, + ); } group } diff --git a/litellm-rust/crates/router/src/health.rs b/litellm-rust/crates/router/src/health.rs index 82c8ba18a01..393d17f2540 100644 --- a/litellm-rust/crates/router/src/health.rs +++ b/litellm-rust/crates/router/src/health.rs @@ -46,19 +46,22 @@ impl Health { /// Counts a failure and returns whether it put the deployment in cooldown. With /// `allowed_fails` set this is Python's v1 rule (more than `allowed_fails` failures in - /// the window); without it, Python's v2 rule (a 429 or a rejected call when the group - /// has alternatives, or a failure rate over the current minute). + /// the window, counted across replicas when `fleet_failures` carries the shared count); + /// without it, Python's v2 rule (a 429 or a rejected call when the group has + /// alternatives, or a failure rate over the current minute), which Python also keeps + /// per process. pub(crate) fn record_failure( &self, now: u64, failure: Failure, policy: &CooldownPolicy, alternatives: bool, + fleet_failures: Option, ) -> bool { self.roll(now, policy); let failures = self.failures.fetch_add(1, Relaxed).saturating_add(1); let cool = match policy.allowed_fails { - Some(allowed) => failures > allowed, + Some(allowed) => fleet_failures.unwrap_or(failures).max(failures) > allowed, None => { let total = failures.saturating_add(self.successes.load(Relaxed)); let failing = (failures == total && total >= SINGLE_DEPLOYMENT_MINIMUM_REQUESTS) @@ -75,12 +78,16 @@ impl Health { }; if cool { let cooldown = u64::try_from(policy.cooldown_time.as_millis()).unwrap_or(u64::MAX); - self.cooldown_until - .fetch_max(now.saturating_add(cooldown), Relaxed); + self.cool_until(now.saturating_add(cooldown)); } cool } + /// Cools the deployment until `until`, unless a longer cooldown is already running. + pub(crate) fn cool_until(&self, until: u64) { + self.cooldown_until.fetch_max(until, Relaxed); + } + fn roll(&self, now: u64, policy: &CooldownPolicy) { let length = u64::try_from(policy.window().as_millis()) .unwrap_or(u64::MAX) @@ -135,7 +142,7 @@ mod tests { fn fail(health: &Health, times: u32, failure: Failure, policy: &CooldownPolicy) -> bool { (0..times).fold(false, |_, _| { - health.record_failure(1_000, failure, policy, true) + health.record_failure(1_000, failure, policy, true, None) }) } @@ -173,7 +180,7 @@ mod tests { let health = Health::default(); fail(&health, 2, Failure::Unavailable, &V1); - assert!(!health.record_failure(12_000, Failure::Unavailable, &V1, true)); + assert!(!health.record_failure(12_000, Failure::Unavailable, &V1, true, None)); assert!(!health.cooling(12_000)); } @@ -190,7 +197,7 @@ mod tests { let health = Health::default(); assert_eq!( - health.record_failure(1_000, failure, &V2, alternatives), + health.record_failure(1_000, failure, &V2, alternatives, None), cooling ); } @@ -210,6 +217,33 @@ mod tests { assert_eq!(fail(&health, failures, Failure::Unavailable, &V2), cooling); } + #[rstest] + #[case::fleet_over_allowed(Some(3), true)] + #[case::fleet_within_allowed(Some(2), false)] + #[case::fleet_behind_this_replica(Some(0), false)] + fn allowed_fails_counts_the_fleet_when_shared( + #[case] fleet_failures: Option, + #[case] cooling: bool, + ) { + let health = Health::default(); + + assert_eq!( + health.record_failure(1_000, Failure::Unavailable, &V1, true, fleet_failures), + cooling + ); + } + + #[rstest] + fn a_shared_cooldown_never_shortens_a_local_one() { + let health = Health::default(); + fail(&health, 3, Failure::Unavailable, &V1); + health.cool_until(5_000); + + assert!(health.cooling(10_999)); + health.cool_until(20_000); + assert!(health.cooling(19_999)); + } + #[rstest] fn in_flight_counts_until_dropped() { let health = Arc::new(Health::default()); diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index 140510d6b96..b05c231a263 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -3,8 +3,14 @@ mod error; mod group; mod health; mod policy; +mod shared; -use std::{collections::HashMap, time::Instant}; +use std::{ + collections::HashMap, + fmt, + sync::{Arc, OnceLock}, + time::{Duration, Instant}, +}; use litellm_config::{Config, Model}; @@ -13,17 +19,32 @@ pub use error::Error; pub use group::ModelGroup; pub use health::{Failure, InFlight}; pub use policy::{CooldownPolicy, RoutingPolicy, RoutingStrategy}; +pub use shared::{CountFuture, SharedHealth}; use group::Member; +use health::Health; /// The model groups of a config and the health of each deployment in them. -#[derive(Debug)] pub struct Router { groups: HashMap, + /// Every deployment's shared-state id and health, in a fixed order. + deployments: Box<[(Box, Arc)]>, + shared: OnceLock>, policy: RoutingPolicy, epoch: Instant, } +impl fmt::Debug for Router { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("Router") + .field("groups", &self.groups) + .field("policy", &self.policy) + .field("shared", &self.shared.get().is_some()) + .finish_non_exhaustive() + } +} + impl Default for Router { fn default() -> Self { Self::new(std::iter::empty(), RoutingPolicy::default()) @@ -63,6 +84,7 @@ impl Router { Deployment { timeout: deployment.timeout.or(policy.timeout), team_id: owner_team(model), + id: model_id(model), ..deployment }, )) @@ -76,25 +98,60 @@ impl Router { policy: RoutingPolicy, ) -> Self { let mut members: HashMap> = HashMap::new(); - for (name, deployment) in deployments { - members - .entry(name) - .or_default() - .push(Member::new(deployment)); + let mut order = Vec::new(); + for (position, (name, deployment)) in deployments.into_iter().enumerate() { + let id = match &deployment.id { + Some(id) => id.as_str().into(), + None => derived_id(&name, position, &deployment), + }; + let group = members.entry(name.clone()).or_default(); + order.push((name, group.len())); + group.push(Member::new(id, deployment)); } + let groups: HashMap = members + .into_iter() + .map(|(name, members)| { + let group = ModelGroup::new(name.as_str().into(), members.into()); + (name, group) + }) + .collect(); + let deployments = order + .into_iter() + .map(|(name, index)| { + let member = groups[&name].member(index); + (member.id.clone(), member.health.clone()) + }) + .collect(); Self { - groups: members - .into_iter() - .map(|(name, members)| { - let group = ModelGroup::new(name.as_str().into(), members.into()); - (name, group) - }) - .collect(), + groups, + deployments, + shared: OnceLock::new(), policy, epoch: Instant::now(), } } + /// Shares cooldowns, and `allowed_fails` counts, with other replicas through `shared`. + /// Only the first call takes effect. + pub fn share_health(&self, shared: Arc) { + let _ = self.shared.set(shared); + } + + /// Every deployment's id in shared state, in config order; [`Router::cool_down`] takes a + /// position in it. + pub fn deployment_ids(&self) -> impl ExactSizeIterator { + self.deployments.iter().map(|(id, _)| &**id) + } + + /// Cools the deployment at `position` of [`Router::deployment_ids`] for `remaining`, as + /// another replica reported; a longer cooldown already running here is kept. + pub fn cool_down(&self, position: usize, remaining: Duration) { + if let Some((_, health)) = self.deployments.get(position) { + let remaining = u64::try_from(remaining.as_millis()).unwrap_or(u64::MAX); + health.cool_until(self.now().saturating_add(remaining)); + } + } + pub fn group(&self, model_name: &str) -> Option<&ModelGroup> { self.groups.get(model_name) } @@ -206,12 +263,25 @@ impl Router { let Some(failure) = classify(&error) else { return Err(error); }; + let shared = self.shared.get(); + let fleet_failures = match (shared, self.policy.cooldown.allowed_fails) { + (Some(shared), Some(_)) => { + shared + .count_failure(&member.id, self.policy.cooldown.window()) + .await + } + _ => None, + }; if member.health.record_failure( self.now(), failure, &self.policy.cooldown, group.len() > 1, + fleet_failures, ) { + if let Some(shared) = shared { + shared.cooled(&member.id, self.policy.cooldown.cooldown_time, failure); + } tracing::warn!( model_group = group.name(), deployment = index, @@ -241,6 +311,38 @@ impl Router { } } +/// `model_info.id` when it is a string or a number, as Python reads it. +fn model_id(model: &Model) -> Option { + match model.model_info.get("id")? { + litellm_config::Value::String(id) if !id.is_empty() => Some(id.clone()), + litellm_config::Value::Number(id) => Some(id.to_string()), + _ => None, + } +} + +/// An id for a deployment without `model_info.id`, the same on every replica that loads the +/// same config and different when the deployment's group, place or target changes. +fn derived_id(group: &str, position: usize, deployment: &Deployment) -> Box { + use sha2::{Digest, Sha256}; + + let mut hash = Sha256::new(); + for part in [ + group, + &position.to_string(), + &deployment.model, + deployment.api_base.as_deref().unwrap_or_default(), + ] { + hash.update(part.as_bytes()); + hash.update([0]); + } + let digest = hash.finalize(); + let hex: String = digest[..16] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect(); + format!("gateway-{hex}").into() +} + /// The indices of the deployments in `group` that `usable` rejects. fn excluded(group: &ModelGroup, usable: impl Fn(&Deployment) -> bool) -> Vec { (0..group.len()) @@ -529,6 +631,85 @@ mod tests { assert!(result.is_ok()); } + #[derive(Default)] + struct Fleet { + cooled: std::sync::Mutex>, + failures: Option, + } + + impl SharedHealth for Fleet { + fn cooled(&self, deployment: &str, cooldown: Duration, _: Failure) { + self.cooled + .lock() + .unwrap() + .push((deployment.to_owned(), cooldown)); + } + + fn count_failure<'a>(&'a self, _: &'a str, _: Duration) -> CountFuture<'a> { + Box::pin(async move { self.failures }) + } + } + + fn cooling_after(allowed_fails: u32) -> RoutingPolicy { + RoutingPolicy { + cooldown: CooldownPolicy { + allowed_fails: Some(allowed_fails), + cooldown_time: Duration::from_secs(60), + }, + ..retries(0) + } + } + + #[rstest] + #[case::fleet_over_allowed(Some(3), true)] + #[case::fleet_within_allowed(Some(1), false)] + #[case::fleet_unavailable(None, false)] + #[tokio::test] + async fn allowed_fails_counts_failures_across_replicas( + #[case] failures: Option, + #[case] cools: bool, + ) { + let router = router(&["a"], cooling_after(2)); + let fleet = Arc::new(Fleet { + failures, + ..Fleet::default() + }); + router.share_health(fleet.clone()); + + let _ = send(&router, |_| Err(Some(Failure::Unavailable))).await; + + let id = router.deployment_ids().next().unwrap().to_owned(); + let expected = cools.then(|| (id, Duration::from_secs(60))); + assert_eq!(fleet.cooled.lock().unwrap().first().cloned(), expected); + } + + #[rstest] + #[tokio::test] + async fn a_cooldown_from_another_replica_is_skipped_here() { + let router = router(&["a", "b"], retries(0)); + router.cool_down(0, Duration::from_secs(60)); + + for _ in 0..20 { + let (tried, _) = send(&router, |_| Ok("served")).await; + assert_eq!(tried[0].0, "b"); + } + } + + #[rstest] + fn deployments_are_named_by_model_info_id_or_a_stable_derived_id() { + let yaml = "model_list:\n - {model_name: m, litellm_params: {model: openai/a}, model_info: {id: shared-a}}\n - {model_name: m, litellm_params: {model: openai/b}}\n - {model_name: n, litellm_params: {model: openai/b}}\n"; + let ids = |yaml: &str| -> Vec { + let router = Router::from_config(&Config::from_yaml(yaml).unwrap()).unwrap(); + router.deployment_ids().map(str::to_owned).collect() + }; + + let first = ids(yaml); + assert_eq!(first[0], "shared-a"); + assert!(first[1].starts_with("gateway-") && first[2].starts_with("gateway-")); + assert_ne!(first[1], first[2]); + assert_eq!(ids(yaml), first); + } + #[rstest] #[tokio::test] async fn every_deployment_cooling_still_serves() { diff --git a/litellm-rust/crates/router/src/shared.rs b/litellm-rust/crates/router/src/shared.rs new file mode 100644 index 00000000000..7e45e183e03 --- /dev/null +++ b/litellm-rust/crates/router/src/shared.rs @@ -0,0 +1,18 @@ +use std::{future::Future, pin::Pin, time::Duration}; + +use crate::Failure; + +pub type CountFuture<'a> = Pin> + Send + 'a>>; + +/// Health shared with the other replicas of a gateway, so a deployment one replica finds +/// failing cools down on all of them. The router never waits on it for a call that succeeds. +pub trait SharedHealth: Send + Sync { + /// Tells the other replicas `deployment` cooled down for `cooldown`. Must not block: the + /// router calls it between a failed attempt and its retry. + fn cooled(&self, deployment: &str, cooldown: Duration, failure: Failure); + + /// Adds one failure to `deployment`'s count across replicas for `allowed_fails`, in a + /// window of `window`, and returns the new count; `None` when the count is unavailable, and + /// the router then counts this replica's failures alone. + fn count_failure<'a>(&'a self, deployment: &'a str, window: Duration) -> CountFuture<'a>; +} diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs index 3372685e3a1..bd75f2b61f3 100644 --- a/litellm-rust/crates/router/tests/router.rs +++ b/litellm-rust/crates/router/tests/router.rs @@ -116,6 +116,7 @@ fn programmatic_deployments_preserve_overrides_and_share_their_group() { ..Default::default() }, team_id: None, + id: None, }; let router = Router::from_iter([ ("public-model".into(), Deployment::default()), diff --git a/litellm-rust/docs/customer-proxy/shared_state-redis.svg b/litellm-rust/docs/customer-proxy/shared_state-redis.svg new file mode 100644 index 00000000000..7b1b9edfffc --- /dev/null +++ b/litellm-rust/docs/customer-proxy/shared_state-redis.svg @@ -0,0 +1,85 @@ + + + + + + Shared state across gateway replicas through the coordination Redis + Nothing here runs on a request that succeeds: the hot path reads only in-process atomics and the in-process key cache. + + + + Rust gateway replica A + + Router health (atomics per deployment) + in flight, window counts, cooldown_until; + read on every routing choice, never Redis + + Virtual key cache (moka, 60 s TTL) + answers tagged with key hash, team and user; + read on every request, never Redis + + Background tasks (coordination.rs) + cooldown poll: one MGET per second + invalidation subscriber: one SUBSCRIBE + only after a failed attempt: SET cooldown, + INCR allowed_fails (waits at most 250 ms) + + + + Coordination Redis + coordination_redis, cache_params or REDIS_* + [ns:]deployment:<id>:cooldown + {timestamp, cooldown_time, ...} PX cooldown, + Python CooldownCache's key and fields + [ns:]deployment:<id>:allowed_fails + fleet count for router_settings.allowed_fails + [ns:]litellm_proxy.auth_cache_invalidation + {"cache_key": "<hash>" | "team_id:<id>" | + "team_membership:<user>:<team>" | "<user_id>"} + If Redis is down, each replica keeps its own + state, as before; requests never wait on it. + + + + Rust gateway replicas B..N + Same as A. A cooldown read from Redis is + applied with fetch_max, so a longer local + cooldown is never shortened. + Deployments are named by model_info.id, else + gateway-<sha256 of group, position, model, + api_base>, the same on every replica that + loads the same config. + + + Python proxy (admin API, unchanged) + /key/update, /key/block, /key/delete, + /team/update, member changes: commit to + Postgres, then publish_auth_cache_invalidation + Its routers share the cooldown keys too + + + + + + + + + + + + PUBLISH + + + orange: written by a failed attempt + blue: read in the background and applied in process + + + + What each replica does with it + Cooldown: the replica that cools a deployment writes the key; every other replica's poll applies the remaining time within a second. + allowed_fails: a failed attempt adds to the fleet count and cools once the fleet passes the limit (Python's v1 rule). + The v2 failure-rate rule (no allowed_fails) stays per replica, as it is in Python. + Keys: an invalidation drops the named key, or the team's or user's keys, from the cache and the outage fallback on every replica at once. + A store read that overlapped an invalidation is not cached, and every (re)subscribe drops all fresh answers, + since messages published while unsubscribed were missed. +