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", } )