From 3df56d9366d69f0fbcb02a47180aabc679c7ab7e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 31 Oct 2026 19:33:41 -0700 Subject: [PATCH] refactor(redis): keep the selector prefetch and tier policy on the Python host The cooldown store no longer takes the v2 selector's DualCache and raw keys. The selector prefetch moves next to its caller in router_utils, and the store surface is model ids only, so a Rust-backed store never mixes a redis-py read into its own call The Rust tier policy now comes from the host. main reads the cooldown refresh interval and the batch throttle size from the environment and the usage refresh from litellm.default_redis_batch_cache_expiry, so a new router_tiers accessor in litellm/rust_bridge/host/settings.py reports them and both pyclasses build their store with those values. POLICY in redis-store stays as main's defaults Co-Authored-By: Claude Fable 5.1 --- .../python-bridge/src/python_settings.rs | 2 + .../crates/python-bridge/src/redis/mod.rs | 1 + .../crates/python-bridge/src/redis/policy.rs | 130 ++++++++++++++++++ .../src/redis/router_cooldowns.rs | 4 +- .../python-bridge/src/redis/router_usage.rs | 4 +- litellm-rust/crates/redis-store/AGENTS.md | 2 +- .../redis-store/src/router_cooldowns.rs | 19 ++- .../crates/redis-store/src/router_usage.rs | 10 +- .../redis-store/tests/router_cooldowns.rs | 22 +-- .../crates/redis-store/tests/router_usage.rs | 14 +- litellm/redis/stores/router_cooldowns.py | 28 ---- .../routing_prefetch.py} | 16 +-- litellm/router_utils/routing_read_batch.py | 33 +++-- litellm/rust_bridge/host/settings.py | 20 +++ .../test_request_redis_batch_pre_call.py | 2 +- .../redis/stores/test_router_cooldowns.py | 26 +--- tests/unit/rust_bridge/host/test_settings.py | 25 +++- tests/unit/test_internal_context.py | 2 +- 18 files changed, 248 insertions(+), 112 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/policy.rs rename litellm/{redis/stores/_routing_prefetch.py => router_utils/routing_prefetch.py} (89%) diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index eacec6859b5..0152fe8a5b9 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -13,6 +13,8 @@ pub(crate) enum PythonSettings { UrlPolicy, #[strum(serialize = "provider_defaults")] ProviderDefaults, + #[strum(serialize = "router_tiers")] + RouterTiers, #[strum(serialize = "secret_manager")] SecretManager, #[strum(serialize = "secret_manager_binding")] diff --git a/litellm-rust/crates/python-bridge/src/redis/mod.rs b/litellm-rust/crates/python-bridge/src/redis/mod.rs index 45c1c18a29a..e14e59861b4 100644 --- a/litellm-rust/crates/python-bridge/src/redis/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/mod.rs @@ -1,5 +1,6 @@ mod client; mod events; +mod policy; pub(crate) mod router_cooldowns; pub(crate) mod router_usage; diff --git a/litellm-rust/crates/python-bridge/src/redis/policy.rs b/litellm-rust/crates/python-bridge/src/redis/policy.rs new file mode 100644 index 00000000000..c21b8c39f47 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/policy.rs @@ -0,0 +1,130 @@ +use std::time::Duration; + +use litellm_cache::TierPolicy; +use litellm_redis_store::{router_cooldowns, router_usage}; +use pyo3::prelude::*; + +use crate::{ + coercion::{Field, FieldSpec, ProjectionError}, + python_settings::{PythonSettings, Snapshot}, +}; + +const COOLDOWN_REFRESH: FieldSpec = + FieldSpec::new("cooldown_redis_read_interval", seconds); +const USAGE_REFRESH: FieldSpec = FieldSpec::new("usage_redis_read_interval", seconds); +const TRACKED_READS: FieldSpec = FieldSpec::new("max_redis_batch_cache_size", count); + +/// The tier values this Python process runs with. `POLICY` in each family holds `main`'s +/// defaults, the host reports what its environment and `litellm` globals changed. +#[derive(Debug)] +pub(crate) struct RouterTierPolicies { + pub(crate) cooldowns: TierPolicy, + pub(crate) usage: TierPolicy, +} + +pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult { + Ok(project(&PythonSettings::RouterTiers.read(py)?)?) +} + +fn project(snapshot: &Snapshot<'_>) -> Result { + let tracked_reads = snapshot.read(&TRACKED_READS)?; + Ok(RouterTierPolicies { + cooldowns: TierPolicy { + refresh_after: snapshot.read(&COOLDOWN_REFRESH)?, + tracked_reads, + ..router_cooldowns::POLICY + }, + usage: TierPolicy { + refresh_after: snapshot.read(&USAGE_REFRESH)?, + tracked_reads, + ..router_usage::POLICY + }, + }) +} + +fn seconds(field: &Field<'_>) -> Result { + match field.value().extract::() { + Ok(seconds) if seconds.is_finite() && seconds >= 0.0 => { + Ok(Duration::from_secs_f64(seconds)) + } + _ => Err(field.invalid("a non-negative number of seconds")), + } +} + +fn count(field: &Field<'_>) -> Result { + field + .value() + .extract::() + .map_err(|_| field.invalid("a non-negative integer")) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use pyo3::{prelude::*, types::PyDict}; + + use super::project; + use crate::python_settings::PythonSettings; + + fn projected(py: Python<'_>, fields: &str) -> Result { + let locals = PyDict::new(py); + let source = std::ffi::CString::new(format!( + "import types\nsettings = types.SimpleNamespace({fields})" + )) + .unwrap(); + py.run(&source, Some(&locals), Some(&locals)).unwrap(); + let snapshot = + PythonSettings::RouterTiers.snapshot(locals.get_item("settings").unwrap().unwrap()); + project(&snapshot).map_err(|error| PyErr::from(error).to_string()) + } + + #[test] + fn the_host_values_replace_only_the_refresh_and_tracked_reads_of_mains_defaults() { + Python::initialize(); + Python::attach(|py| { + let policies = projected( + py, + "cooldown_redis_read_interval=2.5, usage_redis_read_interval=30, max_redis_batch_cache_size=42", + ) + .unwrap(); + assert_eq!( + policies.cooldowns.refresh_after, + Duration::from_millis(2500) + ); + assert_eq!(policies.usage.refresh_after, Duration::from_secs(30)); + assert_eq!(policies.cooldowns.tracked_reads, 42); + assert_eq!(policies.usage.tracked_reads, 42); + assert_eq!( + policies.cooldowns.backfill_ttl, + litellm_redis_store::router_cooldowns::POLICY.backfill_ttl + ); + assert_eq!( + policies.usage.backfill_ttl, + litellm_redis_store::router_usage::POLICY.backfill_ttl + ); + }); + } + + #[test] + fn a_value_python_would_not_run_with_is_a_configuration_error() { + Python::initialize(); + Python::attach(|py| { + let negative = projected( + py, + "cooldown_redis_read_interval=-1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000", + ) + .unwrap_err(); + let text = projected( + py, + "cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many'", + ) + .unwrap_err(); + assert!( + negative.contains("cooldown_redis_read_interval"), + "{negative}" + ); + assert!(text.contains("max_redis_batch_cache_size"), "{text}"); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs b/litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs index d5874d9e89b..264b8e99e81 100644 --- a/litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs +++ b/litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs @@ -10,6 +10,7 @@ use super::{ PyNumber, call, client::{PythonRedisSettings, shared_client}, memory_error, + policy::router_tier_policies, }; use crate::execution::{run_async_value, run_sync_value}; @@ -80,7 +81,8 @@ impl NativeRouterCooldowns { let client = redis_kwargs .map(|kwargs| shared_client(py, kwargs, namespace)) .transpose()?; - Ok(Self(Arc::new(RouterCooldowns::new(client)))) + let policy = router_tier_policies(py)?.cooldowns; + Ok(Self(Arc::new(RouterCooldowns::new(client, policy)))) } fn put( diff --git a/litellm-rust/crates/python-bridge/src/redis/router_usage.rs b/litellm-rust/crates/python-bridge/src/redis/router_usage.rs index 1a5f99b1543..8690f0bc89d 100644 --- a/litellm-rust/crates/python-bridge/src/redis/router_usage.rs +++ b/litellm-rust/crates/python-bridge/src/redis/router_usage.rs @@ -7,6 +7,7 @@ use super::{ PyNumber, call, client::{PythonRedisSettings, shared_client}, memory_error, + policy::router_tier_policies, }; use crate::execution::run_async_value; @@ -53,7 +54,8 @@ impl NativeRouterUsage { let client = redis_kwargs .map(|kwargs| shared_client(py, kwargs, namespace)) .transpose()?; - Ok(Self(Arc::new(RouterUsage::new(client)))) + let policy = router_tier_policies(py)?.usage; + Ok(Self(Arc::new(RouterUsage::new(client, policy)))) } fn increment<'py>( diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md index 28cf1d55a95..955c98f6768 100644 --- a/litellm-rust/crates/redis-store/AGENTS.md +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -1,6 +1,6 @@ - One module per Redis family (`router_cooldowns.rs`, `router_usage.rs`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys - Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec -- A read-through family builds its own `DualCache` from a `pub const POLICY` copied from `main`'s Python values. Callers pass a `RedisClient` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` +- A read-through family builds its own `DualCache` from a `TierPolicy` the host passes in. Its `pub const POLICY` is `main`'s defaults, and the Python host overrides what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `RedisClient` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` - Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not - `tests/fixtures/router_wire.json` is the wire contract, checked here against real Redis and in `tests/unit/redis/stores/` against Python's own codec. Change both sides together - Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages diff --git a/litellm-rust/crates/redis-store/src/router_cooldowns.rs b/litellm-rust/crates/redis-store/src/router_cooldowns.rs index 154785cc566..e8285679cae 100644 --- a/litellm-rust/crates/redis-store/src/router_cooldowns.rs +++ b/litellm-rust/crates/redis-store/src/router_cooldowns.rs @@ -88,9 +88,9 @@ fn allowed_fails_key(model_id: &str, suffix: Option<&str>) -> String { } } -fn memory(now: &Clock) -> InMemoryCache { +fn memory(policy: TierPolicy, now: &Clock) -> InMemoryCache { let now = now.clone(); - InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(POLICY.backfill_ttl), move || { + InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || { now() }) } @@ -107,14 +107,19 @@ pub struct RouterCooldowns { } impl RouterCooldowns { - pub fn new(client: Option) -> Self { - Self::with_clock(client, system_clock()) + pub fn new(client: Option, policy: TierPolicy) -> Self { + Self::with_clock(client, policy, system_clock()) } - pub fn with_clock(client: Option, now: Clock) -> Self { + pub fn with_clock(client: Option, policy: TierPolicy, now: Clock) -> Self { Self { - cooldowns: DualCache::read_through(memory(&now), client.clone(), POLICY, now.clone()), - fails: DualCache::read_through(memory(&now), client, POLICY, now), + cooldowns: DualCache::read_through( + memory(policy, &now), + client.clone(), + policy, + now.clone(), + ), + fails: DualCache::read_through(memory(policy, &now), client, policy, now), } } diff --git a/litellm-rust/crates/redis-store/src/router_usage.rs b/litellm-rust/crates/redis-store/src/router_usage.rs index bed0e07cfdd..aa6689c7dda 100644 --- a/litellm-rust/crates/redis-store/src/router_usage.rs +++ b/litellm-rust/crates/redis-store/src/router_usage.rs @@ -61,17 +61,17 @@ fn whole_seconds(ttl: Duration) -> Duration { pub struct RouterUsage(DualCache, RedisClient>); impl RouterUsage { - pub fn new(client: Option) -> Self { - Self::with_clock(client, system_clock()) + pub fn new(client: Option, policy: TierPolicy) -> Self { + Self::with_clock(client, policy, system_clock()) } - pub fn with_clock(client: Option, now: Clock) -> Self { + pub fn with_clock(client: Option, policy: TierPolicy, now: Clock) -> Self { let clock = now.clone(); let memory = - InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(POLICY.backfill_ttl), move || { + InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || { clock() }); - Self(DualCache::read_through(memory, client, POLICY, now)) + Self(DualCache::read_through(memory, client, policy, now)) } /// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL. diff --git a/litellm-rust/crates/redis-store/tests/router_cooldowns.rs b/litellm-rust/crates/redis-store/tests/router_cooldowns.rs index 2888dc77043..9ee945fb55b 100644 --- a/litellm-rust/crates/redis-store/tests/router_cooldowns.rs +++ b/litellm-rust/crates/redis-store/tests/router_cooldowns.rs @@ -5,7 +5,7 @@ use std::time::Duration; use litellm_redis_client::{Expiry, Ttl}; use litellm_redis_store::{ Number, - router_cooldowns::{Cooldown, RouterCooldowns}, + router_cooldowns::{Cooldown, POLICY, RouterCooldowns}, }; use litellm_redis_testing::{RedisDeployment, Topology}; use rstest::rstest; @@ -44,7 +44,7 @@ async fn put_writes_mains_str_dict_bytes_with_whole_second_ttl( ) { let deployment = RedisDeployment::start(topology).await.unwrap(); let raw = connect(&deployment); - let store = RouterCooldowns::new(Some(connect(&deployment))); + let store = RouterCooldowns::new(Some(connect(&deployment)), POLICY); for case in cases("cooldowns") { let ttl = Duration::from_secs_f64(case["ttl_seconds"].as_f64().unwrap()); @@ -74,7 +74,7 @@ async fn reads_cooldowns_python_wrote_as_str_dict_or_json( ) { let deployment = RedisDeployment::start(topology).await.unwrap(); let raw = connect(&deployment); - let store = RouterCooldowns::with_clock(Some(connect(&deployment)), fixture_clock()); + let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock()); for case in cases("cooldowns") { raw.set( @@ -100,7 +100,7 @@ async fn an_ended_cooldown_is_not_active_but_still_counts_toward_the_minimum( let deployment = RedisDeployment::start(topology).await.unwrap(); let raw = connect(&deployment); let late = std::sync::Arc::new(|| Duration::from_secs(1_730_000_100)); - let store = RouterCooldowns::with_clock(Some(connect(&deployment)), late); + let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, late); let models = cases("cooldowns") .iter() .map(|case| text(case, "model_id")) @@ -129,8 +129,8 @@ async fn allowed_fails_count_across_routers_and_keep_the_first_ttl( ) { let deployment = RedisDeployment::start(topology).await.unwrap(); let raw = connect(&deployment); - let first = RouterCooldowns::new(Some(connect(&deployment))); - let second = RouterCooldowns::new(Some(connect(&deployment))); + let first = RouterCooldowns::new(Some(connect(&deployment)), POLICY); + let second = RouterCooldowns::new(Some(connect(&deployment)), POLICY); for case in cases("allowed_fails") { let key = text(&case, "key"); @@ -163,7 +163,7 @@ async fn python_written_fail_counts_continue_in_rust( ) { let deployment = RedisDeployment::start(topology).await.unwrap(); let raw = connect(&deployment); - let store = RouterCooldowns::new(Some(connect(&deployment))); + let store = RouterCooldowns::new(Some(connect(&deployment)), POLICY); raw.incr_by( "deployment:dep-a:allowed_fails", 4, @@ -184,7 +184,7 @@ async fn python_written_fail_counts_continue_in_rust( #[rstest] #[tokio::test] async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() { - let store = RouterCooldowns::with_clock(None, fixture_clock()); + let store = RouterCooldowns::with_clock(None, POLICY, fixture_clock()); let case = &cases("cooldowns")[0]; store @@ -225,7 +225,7 @@ async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() { #[tokio::test] async fn a_redis_outage_serves_memory_and_counts_fails_locally() { let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); - let store = RouterCooldowns::with_clock(Some(connect(&deployment)), fixture_clock()); + let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock()); let case = &cases("cooldowns")[0]; store .put( @@ -277,7 +277,7 @@ async fn a_miss_is_read_from_redis_again_only_after_the_one_second_refresh( let deployment = RedisDeployment::start(topology).await.unwrap(); let raw = connect(&deployment); let time = StepClock::at(1_730_000_000.75); - let store = RouterCooldowns::with_clock(Some(connect(&deployment)), time.clock()); + let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, time.clock()); let case = &cases("cooldowns")[0]; let before = store.active(&ids(case)).await.unwrap(); @@ -312,7 +312,7 @@ async fn a_breaker_refusal_does_not_hold_back_the_next_read() { ); let time = StepClock::at(1_730_000_000.75); let guarded = support::connect_guarded(&deployment, breaker); - let store = RouterCooldowns::with_clock(Some(guarded.clone()), time.clock()); + let store = RouterCooldowns::with_clock(Some(guarded.clone()), POLICY, time.clock()); let case = &cases("cooldowns")[0]; deployment.stop().await.unwrap(); store.active(&ids(case)).await.unwrap(); diff --git a/litellm-rust/crates/redis-store/tests/router_usage.rs b/litellm-rust/crates/redis-store/tests/router_usage.rs index 89f5cf7729b..6b18c5d676b 100644 --- a/litellm-rust/crates/redis-store/tests/router_usage.rs +++ b/litellm-rust/crates/redis-store/tests/router_usage.rs @@ -5,7 +5,7 @@ use std::time::Duration; use litellm_redis_client::{Expiry, Ttl}; use litellm_redis_store::{ Number, - router_usage::{COUNTER_TTL, Counter, Kind, RouterUsage}, + router_usage::{COUNTER_TTL, Counter, Kind, POLICY, RouterUsage}, }; use litellm_redis_testing::{RedisDeployment, Topology}; use rstest::rstest; @@ -39,7 +39,7 @@ async fn reads_counters_python_incremented_under_mains_keys( ) { let deployment = RedisDeployment::start(topology).await.unwrap(); let raw = connect(&deployment); - let store = RouterUsage::new(Some(connect(&deployment))); + let store = RouterUsage::new(Some(connect(&deployment)), POLICY); let counters = cases("usage"); for (case, delta) in counters.iter().zip([7.75, 3.0]) { raw.incr_by_float(&text(case, "key"), delta, Expiry::Refresh(COUNTER_TTL)) @@ -63,8 +63,8 @@ async fn deployment_increments_write_incrbyfloat_totals_with_a_refreshed_minute_ ) { let deployment = RedisDeployment::start(topology).await.unwrap(); let raw = connect(&deployment); - let writer = RouterUsage::new(Some(connect(&deployment))); - let reader = RouterUsage::new(Some(connect(&deployment))); + let writer = RouterUsage::new(Some(connect(&deployment)), POLICY); + let reader = RouterUsage::new(Some(connect(&deployment)), POLICY); let case = &cases("usage")[0]; let model = case["model"].as_str(); let [tpm_key, rpm_key] = @@ -106,7 +106,7 @@ async fn a_request_increment_keeps_the_window_python_started( ) { let deployment = RedisDeployment::start(topology).await.unwrap(); let raw = connect(&deployment); - let store = RouterUsage::new(Some(connect(&deployment))); + let store = RouterUsage::new(Some(connect(&deployment)), POLICY); let case = &cases("usage")[1]; let key = text(case, "key"); raw.incr_by_float(&key, 1.0, Expiry::KeepExisting(Duration::from_secs(20))) @@ -128,7 +128,7 @@ async fn a_request_increment_keeps_the_window_python_started( #[rstest] #[tokio::test] async fn without_redis_counters_accumulate_in_memory_until_flushed() { - let store = RouterUsage::new(None); + let store = RouterUsage::new(None, POLICY); let counters = tpm_and_rpm(Some("model")); store @@ -153,7 +153,7 @@ async fn without_redis_counters_accumulate_in_memory_until_flushed() { #[tokio::test] async fn a_failed_pipeline_keeps_the_memory_counts_as_integers() { let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); - let store = RouterUsage::new(Some(connect(&deployment))); + let store = RouterUsage::new(Some(connect(&deployment)), POLICY); deployment.stop().await.unwrap(); store diff --git a/litellm/redis/stores/router_cooldowns.py b/litellm/redis/stores/router_cooldowns.py index fd91a517642..18c3bed077b 100644 --- a/litellm/redis/stores/router_cooldowns.py +++ b/litellm/redis/stores/router_cooldowns.py @@ -1,8 +1,6 @@ from __future__ import annotations -import asyncio from collections.abc import Awaitable, Callable, Mapping, Sequence -from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Protocol, TypeAlias if TYPE_CHECKING: @@ -13,9 +11,7 @@ from litellm._internal_context import service_target from litellm.caching.redis_cache import RedisCache from litellm.rust_bridge._native import NativeRouterCooldowns -from . import _routing_prefetch from ._connection import redis_kwargs -from ._routing_prefetch import TierRead class CooldownCacheValue(TypedDict): @@ -29,21 +25,11 @@ ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns" ActiveCooldowns: TypeAlias = Sequence[tuple[str, CooldownCacheValue]] -@dataclass(frozen=True, slots=True) -class RoutingRead: - cooldowns: ActiveCooldowns - selector_values: Sequence[object | None] | None - - class RouterCooldowns(Protocol): def attach_redis(self, connection: RedisCache | None) -> None: ... def flush_local_counters(self) -> None: ... def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None: ... def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int: ... - def arm_routing_read(self, model_ids: Sequence[str], selector: TierRead | None) -> None: ... - async def routing_read( - self, model_ids: Sequence[str], selector: TierRead | None, parent_otel_span: Span | None - ) -> RoutingRead: ... async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ... def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ... def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float: ... @@ -94,20 +80,6 @@ class RouterCooldownStore: with service_target(ROUTER_COOLDOWNS_TARGET): return self._handle.increment_allowed_fails(model_id, suffix, ttl) - def arm_routing_read(self, model_ids: Sequence[str], selector: TierRead | None) -> None: - if selector is not None: - _routing_prefetch.arm(selector) - - async def routing_read( - self, model_ids: Sequence[str], selector: TierRead | None, parent_otel_span: Span | None - ) -> RoutingRead: - with service_target(ROUTER_COOLDOWNS_TARGET): - cooldowns: Final = self._handle.aactive(model_ids) - if selector is None: - return RoutingRead(cooldowns=tuple(await cooldowns), selector_values=None) - active, selector_values = await asyncio.gather(cooldowns, _routing_prefetch.read(selector, parent_otel_span)) - return RoutingRead(cooldowns=tuple(active), selector_values=selector_values) - async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: with service_target(ROUTER_COOLDOWNS_TARGET): pending: Final = self._handle.aactive(model_ids) diff --git a/litellm/redis/stores/_routing_prefetch.py b/litellm/router_utils/routing_prefetch.py similarity index 89% rename from litellm/redis/stores/_routing_prefetch.py rename to litellm/router_utils/routing_prefetch.py index bcf142c4fb8..bcdfdc82e5a 100644 --- a/litellm/redis/stores/_routing_prefetch.py +++ b/litellm/router_utils/routing_prefetch.py @@ -14,7 +14,7 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.redis_batch import BatchResult, active_request_redis_batches _PREFETCH_SLOT: Final = "routing_read" -SELECTOR_TARGET: Final = "router_usage" +ROUTER_USAGE_TARGET: Final = "router_usage" @dataclass(frozen=True, slots=True) @@ -35,7 +35,7 @@ class ArmedRead: previous: Mapping[str, float | None] def release(self) -> None: - self.read.cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # store owns tier + self.read.cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # the selector owns this tier dict(self.previous) ) @@ -44,7 +44,7 @@ class ArmedRead: self.release() return try: - pending: Final = await self.read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # store owns tier + pending: Final = await self.read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier list(self.read.keys), local_only=True ) values: Final = future.result() @@ -53,7 +53,7 @@ class ArmedRead: for key, local in zip(self.read.keys, pending.result) if local is None and values.get(key) is not None } - await self.read.cache._apply_batch_get(pending, missed) # pyright: ignore[reportPrivateUsage] # store owns tier + await self.read.cache._apply_batch_get(pending, missed) # pyright: ignore[reportPrivateUsage] # the selector owns this tier except Exception: self.release() raise @@ -67,7 +67,7 @@ def arm(read: TierRead) -> None: due, previous = read.cache.reserve_redis_batch_reads(read.keys) if not due: return - with service_target(SELECTOR_TARGET): + with service_target(ROUTER_USAGE_TARGET): result: Final = request.batch(redis).mget(tuple(due)) armed_read: Final = ArmedRead(read=read, fetched=frozenset(due), result=result, previous=previous) result.on_settled(armed_read.settle) @@ -100,18 +100,18 @@ async def _read_prefetched(read: TierRead) -> Sequence[object | None] | None: except Exception as e: # noqa: BLE001 # the read falls back to the tier's own Redis read verbose_router_logger.debug("routing prefetch failed, reading again: %s", e) return None - pending: Final = await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # store owns tier + pending: Final = await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier list(read.keys), local_only=True ) if any(key not in prefetch.fetched for key, value in zip(read.keys, pending.result) if value is None): return None - return await read.cache._apply_batch_get( # pyright: ignore[reportPrivateUsage] # store owns tier + return await read.cache._apply_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier pending, {key: values.get(key) for key, value in zip(read.keys, pending.result) if value is None} ) async def read(read: TierRead, parent_otel_span: Span | None) -> Sequence[object | None] | None: - with service_target(SELECTOR_TARGET): + with service_target(ROUTER_USAGE_TARGET): prefetched: Final = await _read_prefetched(read) if prefetched is not None: return prefetched diff --git a/litellm/router_utils/routing_read_batch.py b/litellm/router_utils/routing_read_batch.py index c0bbaa0c859..1107f209328 100644 --- a/litellm/router_utils/routing_read_batch.py +++ b/litellm/router_utils/routing_read_batch.py @@ -1,21 +1,21 @@ -from collections.abc import Iterator +import asyncio +from collections.abc import Iterator, Sequence from contextlib import contextmanager from contextvars import ContextVar from typing import TYPE_CHECKING, Final from litellm._logging import verbose_router_logger -from litellm.redis.stores import _routing_prefetch -from litellm.redis.stores._routing_prefetch import TierRead, selector_read from litellm.router_deprecated import renamed_router_param from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage +from litellm.router_utils import routing_prefetch +from litellm.router_utils.routing_prefetch import ROUTER_USAGE_TARGET, TierRead, selector_read if TYPE_CHECKING: from opentelemetry.trace import Span from litellm.router_api import RoutingContext - -ROUTER_USAGE_TARGET: Final = "router_usage" +__all__ = ("ROUTER_USAGE_TARGET", "RoutingPrefetch", "RoutingReadBatch") def _selector_keys(selector: LowestTPMLoggingHandler_v2, deployments: list) -> list[str]: @@ -27,17 +27,21 @@ def _selector_read(selector: LowestTPMLoggingHandler_v2 | None, deployments: lis return None if selector is None else selector_read(selector.router_cache, _selector_keys(selector, deployments)) +async def _nothing() -> Sequence[object | None] | None: + return None + + class RoutingPrefetch: @staticmethod @renamed_router_param("litellm_router_instance") def arm(ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: list) -> None: - ctx.health().cooldowns.cooldown_store.arm_routing_read( - ctx.registry().names(kind="ids"), _selector_read(usage_selector, deployments) - ) + read: Final = _selector_read(usage_selector, deployments) + if read is not None: + routing_prefetch.arm(read) @staticmethod def armed() -> bool: - return _routing_prefetch.armed() + return routing_prefetch.armed() _active_routing_read_batch: Final[ContextVar["RoutingReadBatch | None"]] = ContextVar( @@ -74,13 +78,14 @@ class RoutingReadBatch: self, ctx: "RoutingContext", healthy_deployments: list, parent_otel_span: "Span | None" ) -> list[str]: usage_read: Final = _selector_read(self.usage_selector, healthy_deployments) - read: Final = await ctx.health().cooldowns.cooldown_store.routing_read( - ctx.registry().names(kind="ids"), usage_read, parent_otel_span + cooldowns, selector_values = await asyncio.gather( + ctx.health().cooldowns.cooldown_store.active(ctx.registry().names(kind="ids"), parent_otel_span), + _nothing() if usage_read is None else routing_prefetch.read(usage_read, parent_otel_span), ) if usage_read is not None: self.prefetched_usage = PrefetchedUsage( keys=frozenset(usage_read.keys), - values=None if read.selector_values is None else dict(zip(usage_read.keys, read.selector_values)), + values=None if selector_values is None else dict(zip(usage_read.keys, selector_values)), ) - verbose_router_logger.debug("retrieve cooldown models: %s", read.cooldowns) - return [model_id for model_id, _ in read.cooldowns] + verbose_router_logger.debug("retrieve cooldown models: %s", cooldowns) + return [model_id for model_id, _ in cooldowns] diff --git a/litellm/rust_bridge/host/settings.py b/litellm/rust_bridge/host/settings.py index 60a58f1aaba..4932a788221 100644 --- a/litellm/rust_bridge/host/settings.py +++ b/litellm/rust_bridge/host/settings.py @@ -31,6 +31,15 @@ class ProviderDefaults: enable_azure_ad_token_refresh: object +@dataclass(frozen=True, slots=True) +class RouterTierSettings: + """The memory-over-Redis tier values `main`'s router reads from the environment and `litellm` globals.""" + + cooldown_redis_read_interval: float + usage_redis_read_interval: float + max_redis_batch_cache_size: int + + @dataclass(frozen=True, slots=True) class SecretManager: readable: bool @@ -108,6 +117,17 @@ def secret_manager_binding() -> SecretManagerBinding: ) +def router_tiers() -> RouterTierSettings: + import litellm + from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE + + return RouterTierSettings( + cooldown_redis_read_interval=DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, + usage_redis_read_interval=litellm.default_redis_batch_cache_expiry or 10, + max_redis_batch_cache_size=DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, + ) + + def provider_defaults() -> ProviderDefaults: import litellm diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index 845acdb80c9..96652b99007 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -29,8 +29,8 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( RateLimitUnverifiableError, ) from litellm.proxy.utils import InternalUsageCache -from litellm.redis.stores._routing_prefetch import ArmedRead, TierRead from litellm.router_api import RoutingContext +from litellm.router_utils.routing_prefetch import ArmedRead, TierRead from litellm.router_utils.routing_read_batch import ( ROUTER_USAGE_TARGET, RoutingPrefetch, diff --git a/tests/unit/redis/stores/test_router_cooldowns.py b/tests/unit/redis/stores/test_router_cooldowns.py index 0dec26f9859..a80c3c5c053 100644 --- a/tests/unit/redis/stores/test_router_cooldowns.py +++ b/tests/unit/redis/stores/test_router_cooldowns.py @@ -1,21 +1,18 @@ -import asyncio import ast +import asyncio import json from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path from typing import Final +import pytest from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict -import pytest - from litellm._internal_context import current_service_target -from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCache from litellm.redis.stores._connection import redis_kwargs -from litellm.redis.stores._routing_prefetch import selector_read from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore _WIRE: Final = Path(__file__).parents[4] / "litellm-rust/crates/redis-store/tests/fixtures/router_wire.json" @@ -133,25 +130,6 @@ def test_minimum_falls_back_to_the_default_when_nothing_or_zero_is_stored( assert store.minimum(["dep"], 60.0, None) == expected -@pytest.mark.asyncio -async def test_the_routing_read_pairs_native_cooldowns_with_the_selector_counters() -> None: - handles: Final = Handles() - store: Final = RouterCooldownStore(None, handle=handles) - store.put("dep", _entry(), 60.0) - selector: Final = DualCache() - await selector.async_set_cache( # pyright: ignore[reportUnknownMemberType] # legacy DualCache is untyped - "dep:model:tpm:12-34", 7 - ) - - read: Final = await store.routing_read(["dep"], selector_read(selector, ["dep:model:tpm:12-34", "missing"]), None) - alone: Final = await store.routing_read(["dep"], None, None) - - assert read.cooldowns == (("dep", _entry()),) - assert list(read.selector_values or ()) == [7, None] - assert alone.selector_values is None - assert handles.built[0].targets[-1] == "router_cooldowns" - - def test_the_wire_fixture_is_what_mains_codec_writes_and_reads() -> None: wire: Final = TypeAdapter(_Wire).validate_json(_WIRE.read_text()) for case in wire["cooldowns"]: diff --git a/tests/unit/rust_bridge/host/test_settings.py b/tests/unit/rust_bridge/host/test_settings.py index 5d3486d9a64..a73eb41bd9b 100644 --- a/tests/unit/rust_bridge/host/test_settings.py +++ b/tests/unit/rust_bridge/host/test_settings.py @@ -4,6 +4,7 @@ import httpx import pytest import litellm +from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.custom_httpx.http_handler import default_user_agent from litellm.rust_bridge.host import settings @@ -46,6 +47,26 @@ def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch ) +def test_router_tiers_read_the_env_constants_and_the_litellm_global(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.constants + + monkeypatch.setattr(litellm.constants, "DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", 2.5) + monkeypatch.setattr(litellm.constants, "DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 42) + monkeypatch.setattr(litellm, "default_redis_batch_cache_expiry", 30.0) + + assert settings.router_tiers() == settings.RouterTierSettings( + cooldown_redis_read_interval=2.5, + usage_redis_read_interval=30.0, + max_redis_batch_cache_size=42, + ) + + +def test_router_tiers_fall_back_to_dual_caches_ten_second_refresh(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "default_redis_batch_cache_expiry", None) + + assert settings.router_tiers().usage_redis_read_interval == DualCache().redis_batch_cache_expiry + + def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_USER_AGENT", "operator/1") monkeypatch.setenv("SSL_VERIFY", "false") @@ -114,9 +135,7 @@ def test_secret_manager_is_native_only_when_a_readable_system_is_configured( monkeypatch.setattr(litellm, "_key_management_system", system) monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode)) - assert settings.secret_manager() == settings.SecretManager( - readable=access_mode != "write_only", native=native - ) + assert settings.secret_manager() == settings.SecretManager(readable=access_mode != "write_only", native=native) def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 4b6780fd058..af621a29567 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -32,7 +32,7 @@ _CACHE_PARAMETER_HELPERS: Final = frozenset( "litellm/llms/anthropic/prompt_cache_prediction.py", "litellm/proxy/common_utils/cache_coordinator.py", "litellm/proxy/common_utils/user_api_key_cache.py", - "litellm/redis/stores/_routing_prefetch.py", + "litellm/router_utils/routing_prefetch.py", "litellm/utils.py", } )