From b5a7064fafad09f75a776d15c6bd25c613a231dd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 10 Nov 2026 16:45:58 +0900 Subject: [PATCH] refactor(router): count io-token itpm and otpm on the native router_usage store The io-token check reserves, reconciles and refunds its global_router itpm and otpm counters through the native router_usage store, so ModelRateLimitingCheck keeps no DualCache. A reservation stashes its counter's fields in the request metadata instead of a Redis key, so nothing parses a key back into ids. The Redis keys and values are unchanged, pinned by contract tests on both sides. Co-Authored-By: Claude Opus 5.5 --- .../src/redis/store/router_usage.rs | 12 + .../crates/python-bridge/src/router/mod.rs | 16 - .../redis-store/src/router_usage/mod.rs | 71 +- .../crates/redis-store/tests/router_usage.rs | 61 ++ .../internal_call_metadata.py | 6 +- litellm/redis/stores/router_usage/__init__.py | 9 + litellm/redis/stores/router_usage/handle.py | 1 + litellm/router_rust/core.py | 6 +- litellm/router_rust/invoke/endpoints.py | 27 +- .../router_rust/invoke/prompt_management.py | 2 +- litellm/router_rust/invoke/shaping.py | 8 +- litellm/router_rust/invoke/websocket.py | 2 +- litellm/router_rust/steps/callbacks.py | 2 +- litellm/router_rust/surface/legacy_members.py | 2 +- litellm/router_rust/surface/model_info.py | 46 +- .../io_token_rate_limit_check.py | 277 ++++---- .../pre_call_checks/model_rate_limit_check.py | 22 +- litellm/rust_bridge/_native.pyi | 2 +- .../unit/redis/stores/router_usage_double.py | 6 + .../test_responses_websocket_all_providers.py | 81 ++- ...st_registry_characterization_model_info.py | 15 +- .../test_enforce_model_rate_limits.py | 3 +- .../test_router/test_io_token_rate_limits.py | 618 +++++++++--------- .../test_router/test_router_helper_utils.py | 39 +- 24 files changed, 678 insertions(+), 656 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs index 7d8c11830dc..b13ea3cf302 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs @@ -301,6 +301,18 @@ impl NativeRouterUsage { }) } + fn shared_or_local<'py>( + &self, + py: Python<'py>, + counter: PyCounter, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_async(py, "async_get_cache", kind_error, async move { + Ok(store.shared_or_local(&counter).await?.map(PyNumber)) + }) + } + fn current_sync( &self, py: Python<'_>, diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs index c10fa2b9c1b..8e8e10d20c5 100644 --- a/litellm-rust/crates/python-bridge/src/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -385,22 +385,6 @@ impl RouterCoreHandle { .collect() } - /// The key deployment `id` serving litellm `model` counts `metric` (`tpm`, `rpm`, `itpm` or - /// `otpm`) under in `minute` (`HH-MM`), as the usage store spells it. - fn usage_key(&self, id: &str, model: &str, metric: &str, minute: &str) -> PyResult { - let kind: Kind = metric - .parse() - .map_err(|_| PyValueError::new_err(format!("`{metric}` is not a usage metric")))?; - Ok(Counter { - scope: Scope::GlobalRouter, - deployment_id: id.to_owned(), - model: Some(model.to_owned()), - window: minute.to_owned(), - kind, - } - .key()) - } - /// Keep the providers `other`'s build resolved for the deployments this router serves too. fn adopt_providers(&self, other: &Self) { self.router.adopt_providers(&other.router); diff --git a/litellm-rust/crates/redis-store/src/router_usage/mod.rs b/litellm-rust/crates/redis-store/src/router_usage/mod.rs index 077069cb06d..e049b5b8675 100644 --- a/litellm-rust/crates/redis-store/src/router_usage/mod.rs +++ b/litellm-rust/crates/redis-store/src/router_usage/mod.rs @@ -219,28 +219,21 @@ impl RouterUsage { if local.is_some_and(|local| local.as_f64() >= limit.as_f64()) { return Ok(local); } - let Some(client) = self.usage.redis().client() else { - return Ok(local); - }; - let key = counter.key(); - match client.send(commands::get::>(&key)).await { - Err(RedisError::BreakerOpen) => Ok(local), - Err(error) => { - redis_failure!( - &error, - "litellm.caching.caching: get() - Got exception from REDIS" - ); - Ok(None) - } - Ok(raw) => Ok(raw - .as_deref() - .map(codec::decode) - .transpose() - .unwrap_or_else(|error| { - tracing::warn!(%error, key, "a usage counter did not decode"); - None - })), - } + Ok(match self.shared(counter).await { + Shared::Unasked => local, + Shared::Failed => None, + Shared::Read(count) => count, + }) + } + + /// The io-token check's usage: what every pod counted in Redis, else this router's memory + /// count when Redis has none, is absent, refuses or fails. + pub async fn shared_or_local(&self, counter: &Counter) -> Result, Error> { + let local = self.get_local(counter)?; + Ok(match self.shared(counter).await { + Shared::Read(Some(count)) => Some(count), + Shared::Unasked | Shared::Failed | Shared::Read(None) => local, + }) } pub fn flush_local_counters(&self) -> Result<(), Error> { @@ -252,6 +245,40 @@ impl RouterUsage { pub fn snapshot(&self) -> Result, Error> { Ok(self.usage.memory().entries()?) } + + /// `redis_cache.get_cache` of one counter, which `RedisCache` answers with `None` on a failure + /// and its breaker refuses. + async fn shared(&self, counter: &Counter) -> Shared { + let Some(client) = self.usage.redis().client() else { + return Shared::Unasked; + }; + let key = counter.key(); + match client.send(commands::get::>(&key)).await { + Err(RedisError::BreakerOpen) => Shared::Unasked, + Err(error) => { + redis_failure!( + &error, + "litellm.caching.caching: get() - Got exception from REDIS" + ); + Shared::Failed + } + Ok(raw) => match raw.as_deref().map(codec::decode).transpose() { + Ok(count) => Shared::Read(count), + Err(error) => { + tracing::warn!(%error, key, "a usage counter did not decode"); + Shared::Failed + } + }, + } + } +} + +/// What one counter's Redis read came back with. +enum Shared { + /// No Redis, or its breaker refused the read. + Unasked, + Failed, + Read(Option), } fn keys(counters: &[Counter]) -> Vec { diff --git a/litellm-rust/crates/redis-store/tests/router_usage.rs b/litellm-rust/crates/redis-store/tests/router_usage.rs index 9418c03ee7b..640b686b71e 100644 --- a/litellm-rust/crates/redis-store/tests/router_usage.rs +++ b/litellm-rust/crates/redis-store/tests/router_usage.rs @@ -658,3 +658,64 @@ async fn a_failed_deployment_pipeline_keeps_the_memory_counts_as_integers() { [Some(Number::Int(5)), Some(Number::Int(1))] ); } + +// The io-token check stashes a counter's fields and the store spells the key, so the key must stay +// main's RouterCacheEnum.ITPM / OTPM spelling +#[rstest] +#[case::itpm(Kind::Itpm, "global_router:dep-1:bedrock/claude:itpm:12-34")] +#[case::otpm(Kind::Otpm, "global_router:dep-1:bedrock/claude:otpm:12-34")] +#[tokio::test] +async fn io_token_counters_increment_under_mains_keys(#[case] kind: Kind, #[case] key: &str) { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let counter = Counter { + scope: Scope::GlobalRouter, + deployment_id: "dep-1".to_owned(), + model: Some("bedrock/claude".to_owned()), + window: "12-34".to_owned(), + kind, + }; + + store + .increment(&counter, Number::Int(8), COUNTER_TTL) + .await + .unwrap(); + + let written: Option = raw.send(commands::get(key)).await.unwrap(); + assert_eq!(written, Some("8".to_owned())); +} + +#[rstest] +#[tokio::test] +async fn shared_usage_prefers_what_every_pod_counted_and_falls_back_to_memory() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let [rpm, tpd] = deployment_usage(); + for case in [&rpm, &tpd] { + store + .increment_local(&increment(&case.counter, Number::Int(2), COUNTER_TTL)) + .unwrap(); + } + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(9.0), + Some(Expiry::Refresh(COUNTER_TTL)), + )) + .await + .unwrap(); + + let counted = store.shared_or_local(&rpm.counter).await.unwrap(); + let absent = store.shared_or_local(&tpd.counter).await.unwrap(); + deployment.stop().await.unwrap(); + let failed = store.shared_or_local(&rpm.counter).await.unwrap(); + + assert_eq!(counted, Some(Number::Int(9))); + assert_eq!( + absent, + Some(Number::Int(2)), + "Redis has no count, so memory answers" + ); + assert_eq!(failed, Some(Number::Int(2))); +} diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 138a01c6524..a866c09c36d 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -44,9 +44,9 @@ BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" -ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" -OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key" -IO_TOKEN_RESERVATION_KEYS: Final = frozenset({ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY}) +ITPM_COUNTER_KEY: Final = "_litellm_itpm_counter" +OTPM_COUNTER_KEY: Final = "_litellm_otpm_counter" +IO_TOKEN_RESERVATION_KEYS: Final = frozenset({ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_COUNTER_KEY, OTPM_COUNTER_KEY}) MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups" """Where auth records the model access groups that authorized the request, for the spend writer. diff --git a/litellm/redis/stores/router_usage/__init__.py b/litellm/redis/stores/router_usage/__init__.py index b09874e3512..719d28bde43 100644 --- a/litellm/redis/stores/router_usage/__init__.py +++ b/litellm/redis/stores/router_usage/__init__.py @@ -66,6 +66,7 @@ class RouterUsage(Protocol): def set_local(self, counter: UsageCounter, value: float) -> None: ... async def current(self, counter: UsageCounter, limit: float) -> float | None: ... def current_sync(self, counter: UsageCounter, limit: float) -> float | None: ... + async def shared_or_local(self, counter: UsageCounter) -> float | None: ... def snapshot(self) -> Mapping[str, float]: ... @@ -197,5 +198,13 @@ class RouterUsageStore: def current_sync(self, counter: UsageCounter, limit: float) -> float | None: return self._handle.current_sync(counter, limit) + @with_service_target(ROUTER_USAGE_TARGET) + async def shared_or_local(self, counter: UsageCounter) -> float | None: + """What every pod counted in Redis, else this router's memory count when Redis has none, is absent, refuses + or fails.""" + if not self.has_redis: + return self._handle.get_local(counter) + return await self._handle.shared_or_local(counter) + def snapshot(self) -> Mapping[str, float]: return self._handle.snapshot() diff --git a/litellm/redis/stores/router_usage/handle.py b/litellm/redis/stores/router_usage/handle.py index cb812721d99..71478e32f05 100644 --- a/litellm/redis/stores/router_usage/handle.py +++ b/litellm/redis/stores/router_usage/handle.py @@ -42,5 +42,6 @@ class RouterUsageHandle(Protocol): def set_local(self, counter: UsageCounter, value: float) -> None: ... def current(self, counter: UsageCounter, limit: float) -> Awaitable[float | None]: ... def current_sync(self, counter: UsageCounter, limit: float) -> float | None: ... + def shared_or_local(self, counter: UsageCounter) -> Awaitable[float | None]: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> Mapping[str, float]: ... diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 64c93a6c3e2..e866f3d26bf 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -1290,6 +1290,10 @@ class _EndpointPort: def routing_cache(self) -> "DualCache": return self._core._cache # pyright: ignore[reportPrivateUsage] # see the class docstring + @property + def router_usage(self) -> RouterUsage: + return self._core._usage # pyright: ignore[reportPrivateUsage] # see the class docstring + @property def selectors(self) -> StrategySelectors: return self._core._selectors # pyright: ignore[reportPrivateUsage] # see the class docstring @@ -5100,7 +5104,7 @@ class _Driver: # logging object (shared across attempts, it logs only the first failure). with contextlib.suppress(Exception): await async_io_token_refund_failure( - self._core._cache, # pyright: ignore[reportPrivateUsage] # the driver serves its own core + self._core._usage, # pyright: ignore[reportPrivateUsage] # the driver serves its own core call.kwargs, # pyright: ignore[reportArgumentType] # the call's kwargs dict parent_otel_span=get_parent_otel_span_from_kwargs(call.kwargs), ) diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index e2c621ccb7d..e22e05c6d57 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -56,6 +56,7 @@ from litellm.router_utils.endpoint_eligibility import endpoint_for_call, endpoin if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache + from litellm.redis.stores.router_usage import RouterUsage from litellm.router_api import RouterRegistry from litellm.router_rust.core import RustCore from litellm.router_rust.steps.selection import StrategySelectors @@ -93,6 +94,9 @@ class EndpointCore(Protocol): @property def routing_cache(self) -> "DualCache": ... + @property + def router_usage(self) -> "RouterUsage": ... + @property def selectors(self) -> "StrategySelectors": ... @@ -851,7 +855,6 @@ def attempt_call( named: Final = {name: kwargs.pop(name) for name in plan.named if name in kwargs} data: Final = as_dict(deployment.get("litellm_params")) model_name: Final = str(data.get("model")) - cache: Final = core.routing_cache caching: Final = shaping.cache_responses # A fallback hop routes the target group in a copy of the metadata (`run_async_fallback`). if fallback_level > 0: @@ -868,7 +871,9 @@ def attempt_call( start_silent_experiments(core, plan, silent_model, named.get("messages"), kwargs) if "messages" in named: kwargs.setdefault("messages", named["messages"]) - update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, splits_stream_timeout=True) + update_kwargs_with_deployment( + shaping, deployment, kwargs, usage=core.router_usage, splits_stream_timeout=True + ) kwargs.pop("silent_model", None) final = without_router_only_kwargs( # rebind-ok: rebound as the Python router rebinds it { @@ -882,7 +887,7 @@ def attempt_call( return AttemptCall(plan.host_op, final, kwargs, str(params.get("model")), group) case "text" | "embedding" | "image" | "transcription" | "rerank" | "adapter": update_kwargs_with_deployment( - shaping, deployment, kwargs, cache=cache, splits_stream_timeout=plan.splits_stream_timeout + shaping, deployment, kwargs, usage=core.router_usage, splits_stream_timeout=plan.splits_stream_timeout ) refused: Final = embedding_batch_refusal(plan, deployment, named.get("input"), kwargs, requested) if refused is not None: @@ -896,7 +901,7 @@ def attempt_call( } return AttemptCall(plan.host_op, without_router_only_kwargs(final), kwargs, model_name, group) case "speech": - update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache) + update_kwargs_with_deployment(shaping, deployment, kwargs, usage=core.router_usage) voice = kwargs.pop("voice", None) # rebind-ok: rebound as the Python router rebinds it final = { # rebind-ok: rebound as the Python router rebinds it **data, @@ -910,7 +915,9 @@ def attempt_call( if plan.kind == "cancel_batch": data.update(_resolved_credentials(core, deployment, requested)) data.pop("litellm_credential_name", None) - update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, function_name=plan.function_name) + update_kwargs_with_deployment( + shaping, deployment, kwargs, usage=core.router_usage, function_name=plan.function_name + ) final = { # rebind-ok: rebound as the Python router rebinds it **data, "custom_llm_provider": _provider(core, data), @@ -926,7 +933,7 @@ def attempt_call( shaping, deployment, kwargs, - cache=cache, + usage=core.router_usage, function_name="generic_api_call", splits_stream_timeout=invoked_handler_name(core, plan) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, ) @@ -942,7 +949,7 @@ def attempt_call( } return AttemptCall(plan.host_op, final, kwargs, model_name, group) case "moderation": - update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache) + update_kwargs_with_deployment(shaping, deployment, kwargs, usage=core.router_usage) kwargs["model"] = model_name kwargs.update(data) return AttemptCall(plan.host_op, kwargs, kwargs, model_name, group) @@ -973,7 +980,7 @@ def _generic_attempt( core.shaping, deployment, kwargs, - cache=core.routing_cache, + usage=core.router_usage, function_name="_ageneric_api_call_with_fallbacks", splits_stream_timeout=handler in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, opens_responses_socket=socket, @@ -2205,7 +2212,7 @@ class RustEndpoints: shaping, deployment, new_kwargs, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values - cache=self._core.routing_cache, + usage=self._core.router_usage, function_name="aretrieve_batch", ) bucket: Final = metadata_bucket(new_kwargs, metadata_variable_name) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values @@ -2289,7 +2296,7 @@ class RustEndpoints: self._core.shaping, deployment, kwargs_copy, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values - cache=self._core.routing_cache, + usage=self._core.router_usage, function_name="acreate_file", ) data: Final = as_dict(deployment.get("litellm_params")) diff --git a/litellm/router_rust/invoke/prompt_management.py b/litellm/router_rust/invoke/prompt_management.py index 62bd77e1b48..b83d518acd9 100644 --- a/litellm/router_rust/invoke/prompt_management.py +++ b/litellm/router_rust/invoke/prompt_management.py @@ -76,7 +76,7 @@ async def prompt_management_acompletion( core.shaping, deployment, kwargs, # pyright: ignore[reportUnknownArgumentType] # the request kwargs - cache=core.routing_cache, + usage=core.router_usage, splits_stream_timeout=True, ) params: Final = as_dict(deployment.get("litellm_params")) diff --git a/litellm/router_rust/invoke/shaping.py b/litellm/router_rust/invoke/shaping.py index 258f919c9f1..b3dc7914f5a 100644 --- a/litellm/router_rust/invoke/shaping.py +++ b/litellm/router_rust/invoke/shaping.py @@ -52,7 +52,7 @@ from litellm.types.router import Deployment, GenericLiteLLMParams, LiteLLM_Param from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN if TYPE_CHECKING: - from litellm.caching.dual_cache import DualCache + from litellm.redis.stores.router_usage import RouterUsage from litellm.router_python_host import PythonHost # The SDK-client and max_parallel_requests reads of the routing cache: process-local, so the family @@ -403,7 +403,7 @@ def update_kwargs_with_deployment( deployment: Mapping[str, object], kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place *, - cache: "DualCache | None", + usage: "RouterUsage | None", function_name: str | None = None, splits_stream_timeout: bool = False, opens_responses_socket: bool = False, @@ -466,8 +466,8 @@ def update_kwargs_with_deployment( if isinstance(other_bucket_metadata, Mapping) and KUBERNETES_POD_ROUTING_KEY in other_bucket_metadata: kwargs[other_bucket] = {**as_dict(other_bucket_metadata), KUBERNETES_POD_ROUTING_KEY: routing} # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them - if cache is not None: - refund_stale_reservation_before_retry(cache, kwargs) # pyright: ignore[reportArgumentType] # the request's kwargs dict + if usage is not None: + refund_stale_reservation_before_retry(usage, kwargs) # pyright: ignore[reportArgumentType] # the request's kwargs dict set_io_token_rate_limit_request_kwargs( kwargs, # pyright: ignore[reportArgumentType] # the request's kwargs dict store_in_context=deployment_has_io_token_limits(deployment), # pyright: ignore[reportArgumentType] # a deployment record diff --git a/litellm/router_rust/invoke/websocket.py b/litellm/router_rust/invoke/websocket.py index e7baeb92345..18d57527586 100644 --- a/litellm/router_rust/invoke/websocket.py +++ b/litellm/router_rust/invoke/websocket.py @@ -64,7 +64,7 @@ async def admit_deployment_turn( io_token_admission: Final[Mapping[str, object]] = ( MappingProxyType({}) if rate_limit_check is None - else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) # pyright: ignore[reportArgumentType] # the configured deployment dict + else await async_admit_io_token_turn(rate_limit_check.usage, deployment, get_utc_datetime()) # pyright: ignore[reportArgumentType] # the configured deployment dict ) budget_limiter: Final[object] = _budget_limiter(core) if budget_limiter is not None: diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index b566f358ddf..25fb067c206 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -306,7 +306,7 @@ class OptionalChecks: ) self.router_budget_logger = callback elif name == "enforce_model_rate_limits": - callback = ModelRateLimitingCheck(usage=self._usage, dual_cache=self._cache) + callback = ModelRateLimitingCheck(usage=self._usage) if callback is None: continue if self.optional_callbacks is not None and any( diff --git a/litellm/router_rust/surface/legacy_members.py b/litellm/router_rust/surface/legacy_members.py index fe34b99758f..e4633dc32cb 100644 --- a/litellm/router_rust/surface/legacy_members.py +++ b/litellm/router_rust/surface/legacy_members.py @@ -421,7 +421,7 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): except Exception: with contextlib.suppress(Exception): await async_io_token_refund_failure( - self._cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs) + self._usage, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs) ) raise await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) diff --git a/litellm/router_rust/surface/model_info.py b/litellm/router_rust/surface/model_info.py index d16b3d19d9b..c1a39035bce 100644 --- a/litellm/router_rust/surface/model_info.py +++ b/litellm/router_rust/surface/model_info.py @@ -7,8 +7,8 @@ host registry. router-core holds the rules: a deployment's model info (``get_router_model_info``, ``get_deployment_model_info``: ``router-core/src/model_info.rs``), a model group's info and the configured-value reads (``_set_model_group_info``, ``get_model_listing_info``, the configured token -limits, mode, display name and service tiers: ``router-core/src/group_info.rs``), and the usage -counters' keys (``RouterCoreHandle.usage_key``). What lives in this process is answered here as a +limits, mode, display name and service tiers: ``router-core/src/group_info.rs``); the usage +counters' keys are the ``router_usage`` store's. What lives in this process is answered here as a query asks for it, at query time: the cost map (``get_model_info``, ``litellm.model_cost``), provider resolution, the wildcard pattern routes, what the background refresh discovered, and a group deployment's validated params, supported params, endpoints, fast mode and reasoning efforts. @@ -33,7 +33,6 @@ from pydantic import BaseModel, TypeAdapter from typing_extensions import overload import litellm -from litellm._internal_context import service_target from litellm._logging import verbose_router_logger from litellm.litellm_core_utils.add_retry_fallback_headers import ( apply_quality_router_decision_headers, @@ -54,7 +53,7 @@ from litellm.litellm_core_utils.reasoning_effort_capability import ( from litellm.litellm_core_utils.routing_common_utils import team_may_use_deployment from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.bedrock.common_utils import get_bedrock_base_model -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter +from litellm.redis.stores.router_usage import RouterUsage, UsageCounter from litellm.router_python_host import python_host from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import build_io_token_rate_limit_headers from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata @@ -797,30 +796,31 @@ class HostModelInfo: """ Returns current ITPM/OTPM usage for a model group (sum across deployments). """ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") - itpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged - otpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged - + current_minute: Final = get_utc_datetime().strftime("%H-%M") model_list: Final = self._core_get_model_list(model_name=model_group) if model_list is None: return None, None - for model in model_list: - model_id: str | None = model.get("model_info", {}).get("id") - litellm_model: str | None = model["litellm_params"].get("model") - if model_id is None or litellm_model is None: - continue - itpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "itpm", current_minute)) - otpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "otpm", current_minute)) + named: Final[tuple[tuple[object, object], ...]] = tuple( + (model.get("model_info", {}).get("id"), model["litellm_params"].get("model")) for model in model_list + ) + deployments: Final = tuple( + (f"{model_id}", f"{litellm_model}") + for model_id, litellm_model in named + if model_id is not None and litellm_model is not None + ) + itpm_counters: Final = tuple( + UsageCounter("global_router", model_id, litellm_model, current_minute, "itpm") + for model_id, litellm_model in deployments + ) + otpm_counters: Final = tuple( + UsageCounter("global_router", model_id, litellm_model, current_minute, "otpm") + for model_id, litellm_model in deployments + ) + combined_values: Final = await self._core_usage_store.get_many((*itpm_counters, *otpm_counters)) - with service_target(ROUTER_USAGE_TARGET): - combined_values: Final = await self._core_cache.async_batch_get_cache(keys=itpm_keys + otpm_keys) - if combined_values is None: - return None, None - - itpm_values: Final = combined_values[: len(itpm_keys)] - otpm_values: Final = combined_values[len(itpm_keys) :] + itpm_values: Final = combined_values[: len(itpm_counters)] + otpm_values: Final = combined_values[len(itpm_counters) :] total_itpm: int | None = None # rebind-ok: ported from the Python core unchanged for value in itpm_values: diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 846cc9f3d61..9917e4f7ea7 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -14,36 +14,35 @@ from __future__ import annotations import contextlib import contextvars from collections.abc import Mapping, MutableMapping -from dataclasses import dataclass +from dataclasses import asdict, dataclass, replace from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger -from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION from litellm.litellm_core_utils.core_helpers import coerce_token_limit from litellm.litellm_core_utils.internal_call_metadata import ( IO_TOKEN_RESERVATION_KEYS, - ITPM_CACHE_KEY, + ITPM_COUNTER_KEY, ITPM_RESERVED_KEY, - OTPM_CACHE_KEY, + OTPM_COUNTER_KEY, OTPM_RESERVED_KEY, ) from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET -from litellm.types.router import RouterCacheEnum, RouterErrors +from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter +from litellm.types.router import RouterErrors from litellm.utils import get_utc_datetime _MODEL_INFO_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_STASHED_COUNTER: Final[TypeAdapter[UsageCounter]] = TypeAdapter(UsageCounter) if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -144,16 +143,17 @@ def _without_io_token_limit_keys(section: object) -> object: return {key: value for key, value in section.items() if key not in _IO_TOKEN_LIMIT_KEYS} -def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | None: - model_id: Final = deployment.get("model_info", {}).get("id") - deployment_name: Final = deployment.get("litellm_params", {}).get("model") +def _get_counters(deployment: dict, current_minute: str) -> tuple[UsageCounter, UsageCounter] | None: + model_id: Final[object] = deployment.get("model_info", {}).get("id") + deployment_name: Final[object] = deployment.get("litellm_params", {}).get("model") # Without both a deployment id and model name the key would collapse to a # shared "None:None" bucket across misconfigured deployments, so bail out. if model_id is None or deployment_name is None: return None - itpm_key = RouterCacheEnum.ITPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute) - otpm_key = RouterCacheEnum.OTPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute) - return itpm_key, otpm_key + return ( + UsageCounter("global_router", f"{model_id}", f"{deployment_name}", current_minute, "itpm"), + UsageCounter("global_router", f"{model_id}", f"{deployment_name}", current_minute, "otpm"), + ) def _served_model(deployment: object) -> str: @@ -354,14 +354,15 @@ def _resolve_reconcile_usage_tokens( def _reservation( - itpm_reserved: int, otpm_reserved: int, itpm_cache_key: str | None, otpm_cache_key: str | None + itpm_reserved: int, otpm_reserved: int, itpm: UsageCounter | None, otpm: UsageCounter | None ) -> Mapping[str, object]: + """The reservation as plain fields, which survive the copies logging makes of the request metadata.""" return MappingProxyType( { ITPM_RESERVED_KEY: itpm_reserved, OTPM_RESERVED_KEY: otpm_reserved, - ITPM_CACHE_KEY: itpm_cache_key, - OTPM_CACHE_KEY: otpm_cache_key, + ITPM_COUNTER_KEY: None if itpm is None else asdict(itpm), + OTPM_COUNTER_KEY: None if otpm is None else asdict(otpm), } ) @@ -372,29 +373,40 @@ def _stash_reservation_in_metadata(request: _IoTokenRateLimitRequest | None, res metadata.update(reservation) -def _key_for_this_minute(reserved_key: str) -> str: - counter, _, _ = reserved_key.rpartition(":") - return f"{counter}:{get_utc_datetime().strftime('%H-%M')}" +def _this_minute(counter: UsageCounter) -> UsageCounter: + return replace(counter, window=get_utc_datetime().strftime("%H-%M")) + + +_Reservation = tuple[int, int, UsageCounter | None, UsageCounter | None] def _reconcile_increments( - reservation: tuple[int, int, str | None, str | None], billable_input: int, completion_tokens: int -) -> tuple[tuple[str, int], ...]: - itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation - deltas: Final = ((itpm_key, billable_input - itpm_reserved), (otpm_key, completion_tokens - otpm_reserved)) + reservation: _Reservation, billable_input: int, completion_tokens: int +) -> tuple[tuple[UsageCounter, int], ...]: + itpm_reserved, otpm_reserved, itpm, otpm = reservation + deltas: Final = ((itpm, billable_input - itpm_reserved), (otpm, completion_tokens - otpm_reserved)) return tuple( - (key if delta < 0 else _key_for_this_minute(key), delta) for key, delta in deltas if key is not None and delta + (counter if delta < 0 else _this_minute(counter), delta) + for counter, delta in deltas + if counter is not None and delta ) -def _extract_reservation(reservation: Mapping[str, int | str | None]) -> tuple[int, int, str | None, str | None]: - itpm_cache_key: Final = reservation.get(ITPM_CACHE_KEY) - otpm_cache_key: Final = reservation.get(OTPM_CACHE_KEY) +def _stashed_counter(value: object) -> UsageCounter | None: + if value is None: + return None + try: + return _STASHED_COUNTER.validate_python(value) + except ValidationError: + return None + + +def _extract_reservation(reservation: Mapping[str, object]) -> _Reservation: return ( - int(reservation.get(ITPM_RESERVED_KEY, 0) or 0), - int(reservation.get(OTPM_RESERVED_KEY, 0) or 0), - itpm_cache_key if isinstance(itpm_cache_key, str) else None, - otpm_cache_key if isinstance(otpm_cache_key, str) else None, + int(reservation.get(ITPM_RESERVED_KEY, 0) or 0), # pyright: ignore[reportArgumentType] # the int the router stashed + int(reservation.get(OTPM_RESERVED_KEY, 0) or 0), # pyright: ignore[reportArgumentType] # see above + _stashed_counter(reservation.get(ITPM_COUNTER_KEY)), + _stashed_counter(reservation.get(OTPM_COUNTER_KEY)), ) @@ -421,7 +433,7 @@ def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object, return top_level + from_params + from_logging_object -def _read_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> tuple[int, int, str | None, str | None]: +def _read_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> _Reservation: for channel_dict in _reservation_channels(kwargs): if isinstance(channel_dict, dict) and ITPM_RESERVED_KEY in channel_dict: return _extract_reservation(channel_dict) @@ -476,8 +488,8 @@ def _rate_limit_error(limit_label: str, limit: int, current: float) -> litellm.R @with_service_target(ROUTER_USAGE_TARGET) def _sync_increment_with_rollback( - dual_cache: DualCache, - key: str, + usage: RouterUsage, + counter: UsageCounter, value: int, limit: int | None, *, @@ -485,24 +497,16 @@ def _sync_increment_with_rollback( ) -> None: if value <= 0 or limit is None: return - current: Final = dual_cache.increment_cache( - key=key, - value=value, - ttl=RoutingArgsTTL, - ) + current: Final = usage.increment_sync(counter, value, RoutingArgsTTL) if current is not None and current > limit: - dual_cache.increment_cache( - key=key, - value=-value, - ttl=RoutingArgsTTL, - ) + usage.increment_sync(counter, -value, RoutingArgsTTL) raise _rate_limit_error(limit_label, limit, current) @with_service_target(ROUTER_USAGE_TARGET) async def _increment_with_rollback( - dual_cache: DualCache, - key: str, + usage: RouterUsage, + counter: UsageCounter, value: int, limit: int | None, *, @@ -511,25 +515,15 @@ async def _increment_with_rollback( ) -> None: if value <= 0 or limit is None: return - current: Final = await dual_cache.async_increment_cache( - key=key, - value=value, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) + current: Final = await usage.increment(counter, value, RoutingArgsTTL) if current is not None and current > limit: - await dual_cache.async_increment_cache( - key=key, - value=-value, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) + await usage.increment(counter, -value, RoutingArgsTTL) raise _rate_limit_error(limit_label, limit, current) @with_service_target(ROUTER_USAGE_TARGET) def io_token_pre_call_check( - dual_cache: DualCache, + usage: RouterUsage, deployment: dict, ) -> dict | None: itpm_limit, otpm_limit = get_deployment_io_token_limits(deployment) @@ -546,10 +540,10 @@ def io_token_pre_call_check( dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") - cache_keys: Final = _get_cache_keys(deployment, current_minute) - if cache_keys is None: + counters: Final = _get_counters(deployment, current_minute) + if counters is None: return deployment - itpm_key, otpm_key = cache_keys + itpm_counter, otpm_counter = counters itpm_reserved = 0 otpm_reserved = 0 @@ -557,8 +551,8 @@ def io_token_pre_call_check( if itpm_limit is not None: itpm_reserved = _reservation_value(estimated_input, itpm_limit) _sync_increment_with_rollback( - dual_cache, - itpm_key, + usage, + itpm_counter, itpm_reserved, itpm_limit, limit_label="ITPM", @@ -568,19 +562,15 @@ def io_token_pre_call_check( otpm_reserved = otpm_to_reserve try: _sync_increment_with_rollback( - dual_cache, - otpm_key, + usage, + otpm_counter, otpm_reserved, otpm_limit, limit_label="OTPM", ) except Exception: if itpm_reserved > 0: - dual_cache.increment_cache( - key=itpm_key, - value=-itpm_reserved, - ttl=RoutingArgsTTL, - ) + usage.increment_sync(itpm_counter, -itpm_reserved, RoutingArgsTTL) raise _stash_reservation_in_metadata( @@ -588,8 +578,8 @@ def io_token_pre_call_check( _reservation( itpm_reserved, otpm_reserved, - itpm_key if itpm_limit is not None else None, - otpm_key if otpm_limit is not None else None, + itpm_counter if itpm_limit is not None else None, + otpm_counter if otpm_limit is not None else None, ), ) return deployment @@ -597,7 +587,7 @@ def io_token_pre_call_check( @with_service_target(ROUTER_USAGE_TARGET) async def async_io_token_pre_call_check( - dual_cache: DualCache, + usage: RouterUsage, deployment: dict, parent_otel_span: Span | None = None, ) -> dict | None: @@ -610,7 +600,7 @@ async def async_io_token_pre_call_check( if request is not None and request.opens_responses_socket: _stash_reservation_in_metadata( request, - await async_admit_io_token_turn(dual_cache, deployment, get_utc_datetime(), parent_otel_span), + await async_admit_io_token_turn(usage, deployment, get_utc_datetime(), parent_otel_span), ) return deployment _model, served_model = _tokenizer_and_served_models(deployment) @@ -621,10 +611,10 @@ async def async_io_token_pre_call_check( dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") - cache_keys: Final = _get_cache_keys(deployment, current_minute) - if cache_keys is None: + counters: Final = _get_counters(deployment, current_minute) + if counters is None: return deployment - itpm_key, otpm_key = cache_keys + itpm_counter, otpm_counter = counters itpm_reserved = 0 otpm_reserved = 0 @@ -632,8 +622,8 @@ async def async_io_token_pre_call_check( if itpm_limit is not None: itpm_reserved = _reservation_value(estimated_input, itpm_limit) await _increment_with_rollback( - dual_cache, - itpm_key, + usage, + itpm_counter, itpm_reserved, itpm_limit, parent_otel_span=parent_otel_span, @@ -644,8 +634,8 @@ async def async_io_token_pre_call_check( otpm_reserved = otpm_to_reserve try: await _increment_with_rollback( - dual_cache, - otpm_key, + usage, + otpm_counter, otpm_reserved, otpm_limit, parent_otel_span=parent_otel_span, @@ -656,12 +646,7 @@ async def async_io_token_pre_call_check( # release the ITPM reservation already made, or it stays inflated # until the TTL expires. if itpm_reserved > 0: - await dual_cache.async_increment_cache( - key=itpm_key, - value=-itpm_reserved, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) + await usage.increment(itpm_counter, -itpm_reserved, RoutingArgsTTL) raise _stash_reservation_in_metadata( @@ -669,42 +654,26 @@ async def async_io_token_pre_call_check( _reservation( itpm_reserved, otpm_reserved, - itpm_key if itpm_limit is not None else None, - otpm_key if otpm_limit is not None else None, + itpm_counter if itpm_limit is not None else None, + otpm_counter if otpm_limit is not None else None, ), ) return deployment -async def _async_redis_usage(redis_cache: RedisCache | None, key: str, parent_otel_span: Span | None) -> int | None: - if redis_cache is None: - return None - with contextlib.suppress(RedisCircuitBreakerOpenError): - return await redis_cache.async_get_cache(key=key, parent_otel_span=parent_otel_span) - return None - - -@with_service_target(ROUTER_USAGE_TARGET) -async def _async_shared_usage(dual_cache: DualCache, key: str, parent_otel_span: Span | None) -> int | None: - shared: Final = await _async_redis_usage(dual_cache.redis_cache, key, parent_otel_span) - if shared is not None: - return shared - return await dual_cache.async_get_cache(key=key, local_only=True) - - async def _raise_if_at_limit( - dual_cache: DualCache, key: str, limit: int | None, limit_label: str, parent_otel_span: Span | None + usage: RouterUsage, counter: UsageCounter, limit: int | None, limit_label: str, parent_otel_span: Span | None ) -> None: if limit is None: return - usage: Final = await _async_shared_usage(dual_cache, key, parent_otel_span) - if usage is not None and usage >= limit: - raise _rate_limit_error(limit_label, limit, usage) + count: Final = await usage.shared_or_local(counter) + if count is not None and count >= limit: + raise _rate_limit_error(limit_label, limit, count) @with_service_target(ROUTER_USAGE_TARGET) async def async_admit_io_token_turn( - dual_cache: DualCache, + usage: RouterUsage, deployment: Mapping[str, object], now: datetime, parent_otel_span: Span | None = None, @@ -713,36 +682,36 @@ async def async_admit_io_token_turn( itpm_limit, otpm_limit = get_deployment_io_token_limits(configured) if itpm_limit is None and otpm_limit is None: return MappingProxyType({}) - cache_keys: Final = _get_cache_keys(configured, now.strftime("%H-%M")) - if cache_keys is None: + counters: Final = _get_counters(configured, now.strftime("%H-%M")) + if counters is None: return MappingProxyType({}) - itpm_key, otpm_key = cache_keys - await _raise_if_at_limit(dual_cache, itpm_key, itpm_limit, "ITPM", parent_otel_span) - await _raise_if_at_limit(dual_cache, otpm_key, otpm_limit, "OTPM", parent_otel_span) + itpm_counter, otpm_counter = counters + await _raise_if_at_limit(usage, itpm_counter, itpm_limit, "ITPM", parent_otel_span) + await _raise_if_at_limit(usage, otpm_counter, otpm_limit, "OTPM", parent_otel_span) return _reservation( - 0, 0, itpm_key if itpm_limit is not None else None, otpm_key if otpm_limit is not None else None + 0, 0, itpm_counter if itpm_limit is not None else None, otpm_counter if otpm_limit is not None else None ) @with_service_target(ROUTER_USAGE_TARGET) def io_token_reconcile_success( - dual_cache: DualCache, + usage: RouterUsage, kwargs: Mapping[str, object] | None, response_obj: object, ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj reservation: Final = _read_reservation_from_kwargs(request_kwargs) - itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation - if itpm_key is None and otpm_key is None: + itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = reservation + if itpm_counter is None and otpm_counter is None: return billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response) try: if usage_resolved: - for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): - dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL) + for counter, delta in _reconcile_increments(reservation, billable_input, completion_tokens): + usage.increment_sync(counter, delta, RoutingArgsTTL) else: verbose_router_logger.debug( "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", @@ -764,7 +733,7 @@ def io_token_reconcile_success( @with_service_target(ROUTER_USAGE_TARGET) async def async_io_token_reconcile_success( - dual_cache: DualCache, + usage: RouterUsage, kwargs: Mapping[str, object] | None, response_obj: object, *, @@ -773,8 +742,8 @@ async def async_io_token_reconcile_success( request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj reservation: Final = _read_reservation_from_kwargs(request_kwargs) - itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation - if itpm_key is None and otpm_key is None: + itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = reservation + if itpm_counter is None and otpm_counter is None: return billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response) @@ -785,10 +754,8 @@ async def async_io_token_reconcile_success( # event can't re-process it. try: if usage_resolved: - for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): - await dual_cache.async_increment_cache( - key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span - ) + for counter, delta in _reconcile_increments(reservation, billable_input, completion_tokens): + await usage.increment(counter, delta, RoutingArgsTTL) else: verbose_router_logger.debug( "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", @@ -810,30 +777,22 @@ async def async_io_token_reconcile_success( @with_service_target(ROUTER_USAGE_TARGET) def io_token_refund_failure( - dual_cache: DualCache, + usage: RouterUsage, kwargs: Mapping[str, object] | None, ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs - itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs) - if itpm_key is None and otpm_key is None: + itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = _read_reservation_from_kwargs(request_kwargs) + if itpm_counter is None and otpm_counter is None: return - if itpm_key is not None and itpm_reserved > 0: - dual_cache.increment_cache( - key=itpm_key, - value=-itpm_reserved, - ttl=RoutingArgsTTL, - ) - if otpm_key is not None and otpm_reserved > 0: - dual_cache.increment_cache( - key=otpm_key, - value=-otpm_reserved, - ttl=RoutingArgsTTL, - ) + if itpm_counter is not None and itpm_reserved > 0: + usage.increment_sync(itpm_counter, -itpm_reserved, RoutingArgsTTL) + if otpm_counter is not None and otpm_reserved > 0: + usage.increment_sync(otpm_counter, -otpm_reserved, RoutingArgsTTL) _clear_reservation_from_kwargs(request_kwargs) verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved) -def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping[str, object] | None) -> None: +def refund_stale_reservation_before_retry(usage: RouterUsage, kwargs: Mapping[str, object] | None) -> None: """ Synchronously refund and clear any reservation a previous deployment attempt stashed in ``kwargs``, before it's overwritten for the next @@ -857,34 +816,24 @@ def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping """ if not kwargs: return - io_token_refund_failure(dual_cache, kwargs) + io_token_refund_failure(usage, kwargs) @with_service_target(ROUTER_USAGE_TARGET) async def async_io_token_refund_failure( - dual_cache: DualCache, + usage: RouterUsage, kwargs: Mapping[str, object] | None, *, parent_otel_span: Span | None = None, ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs - itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs) - if itpm_key is None and otpm_key is None: + itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = _read_reservation_from_kwargs(request_kwargs) + if itpm_counter is None and otpm_counter is None: return - if itpm_key is not None and itpm_reserved > 0: - await dual_cache.async_increment_cache( - key=itpm_key, - value=-itpm_reserved, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) - if otpm_key is not None and otpm_reserved > 0: - await dual_cache.async_increment_cache( - key=otpm_key, - value=-otpm_reserved, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) + if itpm_counter is not None and itpm_reserved > 0: + await usage.increment(itpm_counter, -itpm_reserved, RoutingArgsTTL) + if otpm_counter is not None and otpm_reserved > 0: + await usage.increment(otpm_counter, -otpm_reserved, RoutingArgsTTL) _clear_reservation_from_kwargs(request_kwargs) verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved) diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index a4c3098a71c..bbeae46582b 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -22,7 +22,6 @@ from pydantic import BaseModel, ConfigDict import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger -from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, CounterIncrement, RouterUsage, UsageCounter from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( @@ -138,11 +137,8 @@ class ModelRateLimitingCheck(CustomLogger): this check actively enforces those limits across ALL routing strategies. """ - def __init__( - self, usage: RouterUsage, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime - ) -> None: + def __init__(self, usage: RouterUsage, clock: Callable[[], datetime] = get_utc_datetime) -> None: self.usage = usage - self.dual_cache = dual_cache self.clock = clock # model_ids already warned about itpm/otpm + tpm/rpm on the same deployment, # so the warning is logged once per deployment rather than per request. @@ -200,7 +196,7 @@ class ModelRateLimitingCheck(CustomLogger): def _refund_io_token_reservation_if_any(self) -> None: request_kwargs: Final = get_io_token_rate_limit_request_kwargs() if request_kwargs is not None: - io_token_refund_failure(self.dual_cache, request_kwargs) + io_token_refund_failure(self.usage, request_kwargs) async def _async_refund_io_token_reservation_if_any( self, @@ -209,7 +205,7 @@ class ModelRateLimitingCheck(CustomLogger): request_kwargs: Final = get_io_token_rate_limit_request_kwargs() if request_kwargs is not None: await async_io_token_refund_failure( - self.dual_cache, + self.usage, request_kwargs, parent_otel_span=parent_otel_span, ) @@ -280,7 +276,7 @@ class ModelRateLimitingCheck(CustomLogger): if deployment_has_io_token_limits(deployment): self._warn_io_token_and_tpm_rpm_coexist_once(deployment) io_token_pre_call_check( - self.dual_cache, + self.usage, deployment, ) io_reservation_made = True @@ -364,7 +360,7 @@ class ModelRateLimitingCheck(CustomLogger): if deployment_has_io_token_limits(deployment): self._warn_io_token_and_tpm_rpm_coexist_once(deployment) await async_io_token_pre_call_check( - self.dual_cache, + self.usage, deployment, parent_otel_span=parent_otel_span, ) @@ -453,7 +449,7 @@ class ModelRateLimitingCheck(CustomLogger): try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") await async_io_token_reconcile_success( - self.dual_cache, + self.usage, kwargs, response_obj, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), @@ -500,7 +496,7 @@ class ModelRateLimitingCheck(CustomLogger): # Never fail the primary logging pipeline over an io-token refund error. with contextlib.suppress(Exception): await async_io_token_refund_failure( - self.dual_cache, + self.usage, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), ) @@ -513,7 +509,7 @@ class ModelRateLimitingCheck(CustomLogger): """ try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") - io_token_reconcile_success(self.dual_cache, kwargs, response_obj) + io_token_reconcile_success(self.usage, kwargs, response_obj) if standard_logging_object is None: return @@ -545,6 +541,6 @@ class ModelRateLimitingCheck(CustomLogger): def log_failure_event(self, kwargs, response_obj, start_time, end_time): with contextlib.suppress(Exception): io_token_refund_failure( - self.dual_cache, + self.usage, kwargs, ) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 0248a3a4ebf..204ee52185c 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -184,6 +184,7 @@ class NativeRouterUsage: def set_local(self, counter: _UsageCounter, value: float) -> None: ... def current(self, counter: _UsageCounter, limit: float) -> Future[float | None]: ... def current_sync(self, counter: _UsageCounter, limit: float) -> float | None: ... + def shared_or_local(self, counter: _UsageCounter) -> Future[float | None]: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... @@ -1089,7 +1090,6 @@ class RouterCoreHandle: @staticmethod def configured_service_tiers(model_infos: list[object]) -> list[object]: ... # The routing-cache key a deployment counts `metric` (tpm, rpm, itpm, otpm) under in `minute`. - def usage_key(self, id: str, model: str, metric: str, minute: str) -> str: ... def registry_version(self) -> int: ... # The version the deployments before the newest appends last changed at. def registry_rewritten(self) -> int: ... diff --git a/tests/unit/redis/stores/router_usage_double.py b/tests/unit/redis/stores/router_usage_double.py index 416db232e65..558b87d7a03 100644 --- a/tests/unit/redis/stores/router_usage_double.py +++ b/tests/unit/redis/stores/router_usage_double.py @@ -100,6 +100,12 @@ class FakeUsageHandle: def current(self, counter: UsageCounter, limit: float) -> asyncio.Future[float | None]: return _done(self.current_sync(counter, limit)) + def shared_or_local(self, counter: UsageCounter) -> asyncio.Future[float | None]: + self.calls.append(("shared_or_local", (counter,))) + redis: Final = self._reachable() + shared: Final = None if redis is None else redis.get(counter) + return _done(self.local.get(counter) if shared is None else shared) + def _read(self, counters: Sequence[UsageCounter]) -> list[float | None]: redis: Final = self._reachable() return [self.local.get(c, None if redis is None else redis.get(c)) for c in counters] diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index afeea55291c..65bac717011 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -48,10 +48,11 @@ from litellm.llms.volcengine.responses.transformation import ( VolcEngineResponsesAPIConfig, ) from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.types.utils import GenericStreamingChunk if TYPE_CHECKING: - from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.litellm_logging import Logging from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck @@ -1701,8 +1702,8 @@ def _guard_recording(bill_refusal: MagicMock): return guard -_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-34" -_TURNED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-35" +_OPENED_ON: Final = UsageCounter("global_router", "ws-dep", "openai/gpt-5-mini", "12-34", "otpm") +_TURNED_ON: Final = UsageCounter("global_router", "ws-dep", "openai/gpt-5-mini", "12-35", "otpm") _OPENED_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) _TURNED_AT: Final = datetime(2026, 1, 1, 12, 35, 5, tzinfo=timezone.utc) _WS_DEPLOYMENT: Final = { @@ -1728,12 +1729,12 @@ def _turn_completed() -> str: ) -async def _socket_opened_on(dual_cache: "DualCache") -> Mapping[str, object]: +async def _socket_opened_on(usage: RouterUsageStore) -> Mapping[str, object]: from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn - await dual_cache.async_increment_cache(key=_OPENED_ON, value=_OTHER_USAGE, ttl=60) - await dual_cache.async_increment_cache(key=_TURNED_ON, value=_OTHER_USAGE, ttl=60) - return await async_admit_io_token_turn(dual_cache, _WS_DEPLOYMENT, _OPENED_AT) + usage.increment_local(CounterIncrement(_OPENED_ON, _OTHER_USAGE, 60)) + usage.increment_local(CounterIncrement(_TURNED_ON, _OTHER_USAGE, 60)) + return await async_admit_io_token_turn(usage, _WS_DEPLOYMENT, _OPENED_AT) def _turns_finish_at(monkeypatch: pytest.MonkeyPatch, finished_at: datetime) -> None: @@ -1742,13 +1743,13 @@ def _turns_finish_at(monkeypatch: pytest.MonkeyPatch, finished_at: datetime) -> ) -def _later_turns_admitted_through(dual_cache: "DualCache") -> "ResponsesWsTurnAdmission": +def _later_turns_admitted_through(usage: RouterUsageStore) -> "ResponsesWsTurnAdmission": from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn - return partial(async_admit_io_token_turn, dual_cache, _WS_DEPLOYMENT, _TURNED_AT) + return partial(async_admit_io_token_turn, usage, _WS_DEPLOYMENT, _TURNED_AT) -_ITPM_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:itpm:12-34" +_ITPM_OPENED_ON: Final = UsageCounter("global_router", "ws-dep", "openai/gpt-5-mini", "12-34", "itpm") _ITPM_DEPLOYMENT: Final = { "litellm_params": {"model": "openai/gpt-5-mini", "itpm": 40_000}, "model_info": {"id": "ws-dep"}, @@ -1775,12 +1776,14 @@ def _turn_completed_reading_cache(cached_tokens: int) -> str: ) -def _rate_limit_check(dual_cache: "DualCache") -> "ModelRateLimitingCheck": - from litellm.redis.stores.connection import connect - from litellm.redis.stores.router_usage import RouterUsageStore +def _memory_usage() -> RouterUsageStore: + return RouterUsageStore(connect(None).router_usage()) + + +def _rate_limit_check(usage: RouterUsageStore) -> "ModelRateLimitingCheck": from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck - return ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + return ModelRateLimitingCheck(usage=usage) def _connection_logged_through(rate_limits: "ModelRateLimitingCheck", opened_with: Mapping[str, object]) -> "Logging": @@ -2012,13 +2015,11 @@ class TestWebSocketTurnAdmission: ) -> None: import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - from litellm.caching.dual_cache import DualCache from litellm.responses.streaming_iterator import drain_logging_tasks - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck _turns_finish_at(monkeypatch, _TURNED_AT) - dual_cache: Final = DualCache() - opened_with: Final = await _socket_opened_on(dual_cache) + usage: Final = _memory_usage() + opened_with: Final = await _socket_opened_on(usage) first, second, third = _turn_frames() backend_ws: Final = MagicMock() backend_ws.send = AsyncMock() @@ -2029,18 +2030,18 @@ class TestWebSocketTurnAdmission: websocket=_ClientSocket(second, third), backend_ws=backend_ws, request_data={}, - logging_obj=_connection_logged_through(_rate_limit_check(dual_cache), opened_with), + logging_obj=_connection_logged_through(_rate_limit_check(usage), opened_with), first_message=first, - turn_admission=_later_turns_admitted_through(dual_cache), + turn_admission=_later_turns_admitted_through(usage), ) await handler.client_to_backend() await handler.backend_to_client() await drain_logging_tasks() - assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert usage.get_local(_OPENED_ON) == _OTHER_USAGE assert ( - await dual_cache.async_get_cache(key=_TURNED_ON) + usage.get_local(_TURNED_ON) == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN ) @@ -2057,14 +2058,12 @@ class TestWebSocketTurnAdmission: ) -> None: import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - from litellm.caching.dual_cache import DualCache from litellm.responses.streaming_iterator import drain_logging_tasks from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck _turns_finish_at(monkeypatch, _OPENED_AT) - dual_cache: Final = DualCache() - opened_with: Final = await async_admit_io_token_turn(dual_cache, _ITPM_DEPLOYMENT, _OPENED_AT) + usage: Final = _memory_usage() + opened_with: Final = await async_admit_io_token_turn(usage, _ITPM_DEPLOYMENT, _OPENED_AT) first, _, _ = _turn_frames() backend_ws: Final = MagicMock() backend_ws.send = AsyncMock() @@ -2073,7 +2072,7 @@ class TestWebSocketTurnAdmission: websocket=_ClientSocket(), backend_ws=backend_ws, request_data={}, - logging_obj=_connection_logged_through(_rate_limit_check(dual_cache), opened_with), + logging_obj=_connection_logged_through(_rate_limit_check(usage), opened_with), first_message=first, ) @@ -2081,7 +2080,7 @@ class TestWebSocketTurnAdmission: await handler.backend_to_client() await drain_logging_tasks() - assert await dual_cache.async_get_cache(key=_ITPM_OPENED_ON) == counted_input_tokens + assert usage.get_local(_ITPM_OPENED_ON) == counted_input_tokens @pytest.mark.asyncio @pytest.mark.parametrize( @@ -2094,25 +2093,23 @@ class TestWebSocketTurnAdmission: async def test_managed_relay_counts_every_turn_on_the_minute_it_finishes( self, monkeypatch: pytest.MonkeyPatch, failing_turns: frozenset[int], first_turn_tokens: int ) -> None: - from litellm.caching.dual_cache import DualCache - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck _turns_finish_at(monkeypatch, _TURNED_AT) - dual_cache: Final = DualCache() - opened_with: Final = await _socket_opened_on(dual_cache) + usage: Final = _memory_usage() + opened_with: Final = await _socket_opened_on(usage) await _run_managed_relay( monkeypatch, _TwoOutputTokenTurns(failing_turns), - _rate_limit_check(dual_cache), + _rate_limit_check(usage), _turn_frames(), - _later_turns_admitted_through(dual_cache), + _later_turns_admitted_through(usage), opened_with, ) - assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert usage.get_local(_OPENED_ON) == _OTHER_USAGE assert ( - await dual_cache.async_get_cache(key=_TURNED_ON) + usage.get_local(_TURNED_ON) == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN ) @@ -2121,26 +2118,24 @@ class TestWebSocketTurnAdmission: self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter ) -> None: import litellm - from litellm.caching.dual_cache import DualCache - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() _turns_finish_at(monkeypatch, _TURNED_AT) provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=_provider_turn_stream) client_metadata: Final = {"session": "abc", "team": "search"} - dual_cache: Final = DualCache() + usage: Final = _memory_usage() await _relay_turns_to( monkeypatch, "openai/gpt-5-mini", - _rate_limit_check(dual_cache), + _rate_limit_check(usage), tuple( json.dumps({"type": "response.create", "input": word, "metadata": client_metadata}) for word in ("one", "two") ), - _later_turns_admitted_through(dual_cache), - await _socket_opened_on(dual_cache), + _later_turns_admitted_through(usage), + await _socket_opened_on(usage), api_key="sk-fake", ) @@ -2148,7 +2143,7 @@ class TestWebSocketTurnAdmission: client_metadata, client_metadata, ] - assert await dual_cache.async_get_cache(key=_TURNED_ON) == _OTHER_USAGE + 2 * _OUTPUT_TOKENS_PER_TURN + assert usage.get_local(_TURNED_ON) == _OTHER_USAGE + 2 * _OUTPUT_TOKENS_PER_TURN @pytest.mark.asyncio async def test_native_relay_sends_the_provider_only_the_metadata_the_socket_opened_with( diff --git a/tests/unit/test_router/core/test_registry_characterization_model_info.py b/tests/unit/test_router/core/test_registry_characterization_model_info.py index d79aebf75b9..db6c544a6a5 100644 --- a/tests/unit/test_router/core/test_registry_characterization_model_info.py +++ b/tests/unit/test_router/core/test_registry_characterization_model_info.py @@ -4,12 +4,15 @@ limits/mode/display name/service tiers/listing info, group usage this minute, re the rate limit headers ``set_response_headers`` attaches. """ +from typing import Final, Literal + import pytest import litellm from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.router_core import RouterCore -from litellm.types.router import Deployment, RouterCacheEnum +from litellm.redis.stores.router_usage import UsageCounter +from litellm.types.router import Deployment from litellm.utils import get_utc_datetime from tests.unit.test_router.core.registry_helpers import MakeCore, typed, unique_id from tests.unit.test_router.core.scripted_host import completion_response, deployment @@ -45,9 +48,9 @@ async def _serve(core: RouterCore, deployment_id: str, model: str, total_tokens: ) -async def _count(core: RouterCore, kind: RouterCacheEnum, deployment_id: str, model: str, value: int) -> None: - key = kind.value.format(id=deployment_id, model=model, current_minute=_minute()) - await core.health().store.async_set_cache(key, value) +async def _count(core: RouterCore, kind: Literal["itpm", "otpm"], deployment_id: str, model: str, value: int) -> None: + usage: Final = core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the core counts io tokens in + usage.set_local(UsageCounter("global_router", deployment_id, model, _minute(), kind), value) # -- model group info ----------------------------------------------------------------------------- @@ -464,7 +467,7 @@ async def test_group_usage_sums_each_deployments_counters_for_this_minute(make_c async def test_io_token_usage_sums_input_and_output_counters(make_core: MakeCore) -> None: core = _usage_core(make_core) assert await core.get_model_group_io_token_usage("io") == (None, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core - await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7) + await _count(core, "itpm", "i", "openai/gpt-4o", 7) assert await core.get_model_group_io_token_usage("io") == (7, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core @@ -489,7 +492,7 @@ async def test_remaining_usage_of_an_io_limited_group_reports_input_and_output_t make_core: MakeCore, ) -> None: core = _usage_core(make_core) - await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7) + await _count(core, "itpm", "i", "openai/gpt-4o", 7) remaining = await core.health().usage("io", kind="remaining") diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index ee103e2df3c..9edaf885e4f 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -14,7 +14,6 @@ import pytest import litellm from litellm import Router -from litellm.caching.dual_cache import DualCache from litellm.redis.stores.connection import connect from litellm.redis.stores.router_usage import CounterIncrement, CounterKind, RouterUsage, RouterUsageStore, UsageCounter from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( @@ -38,7 +37,7 @@ def _memory_usage() -> RouterUsageStore: def _check( usage: RouterUsage | None = None, clock: Callable[[], datetime] = get_utc_datetime ) -> ModelRateLimitingCheck: - return ModelRateLimitingCheck(usage=usage or _memory_usage(), dual_cache=DualCache(), clock=clock) + return ModelRateLimitingCheck(usage=usage or _memory_usage(), clock=clock) def _counter(deployment: dict, kind: CounterKind, window: str) -> UsageCounter: diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index 7960dbe7073..9ec1a25b97a 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -4,28 +4,23 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). import asyncio from collections.abc import Awaitable, Callable, Mapping +from dataclasses import asdict from datetime import datetime, timezone -from typing import Final -from unittest.mock import AsyncMock, MagicMock +from typing import Final, Literal import pytest -from fakeredis import FakeRedis, FakeServer -from fakeredis.aioredis import FakeRedis as FakeAsyncRedis import litellm from litellm import Router -from litellm.caching.dual_cache import DualCache -from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.redis.stores.connection import connect from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - ITPM_CACHE_KEY, + ITPM_COUNTER_KEY, ITPM_RESERVED_KEY, - OTPM_CACHE_KEY, + OTPM_COUNTER_KEY, OTPM_RESERVED_KEY, _explicit_max_tokens, _reservation_value, @@ -44,9 +39,9 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse -from litellm.types.router import RouterCacheEnum from litellm.types.utils import ModelResponse, Usage from litellm.utils import get_utc_datetime +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store MANTLE_OPUS = "bedrock_mantle/anthropic.claude-opus-4-7" MANTLE_OPUS_MAX_OUTPUT = litellm.get_model_info(model=MANTLE_OPUS)["max_output_tokens"] @@ -56,6 +51,33 @@ AZURE_BASE_MODEL = "azure/gpt-5" AZURE_BASE_MODEL_MAX_OUTPUT = litellm.get_model_info(model=AZURE_BASE_MODEL)["max_output_tokens"] +def _memory_usage() -> RouterUsageStore: + return RouterUsageStore(connect(None).router_usage()) + + +def _io(deployment_id: str, model: str, kind: Literal["itpm", "otpm"], window: str) -> UsageCounter: + return UsageCounter("global_router", deployment_id, model, window, kind) + + +class _FailingKind(FakeUsageHandle): + """A usage store whose increments of one counter kind fail, as a transient Redis error would.""" + + def __init__(self, failing: str) -> None: + super().__init__() + self.failing: Final = failing + + def increment_sync(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float: + if counter.kind == self.failing: + raise RuntimeError("transient cache error") + return super().increment_sync(counter, value, ttl_seconds) + + +def _stashed(value: object) -> UsageCounter: + """The counter a reservation stashed in the request metadata.""" + assert isinstance(value, dict) + return UsageCounter(**value) + + class TestIOTokenRateLimitHelpers: def test_deployment_has_io_token_limits(self): assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}}) @@ -107,8 +129,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_itpm_reservation_and_reconcile(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -128,8 +150,8 @@ class TestModelRateLimitingCheckIOTokens: await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + itpm_key = _io("io-test-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) + otpm_key = _io("io-test-id", "bedrock_mantle/anthropic.claude-opus-4-7", "otpm", minute) kwargs = { "standard_logging_object": { @@ -151,16 +173,16 @@ class TestModelRateLimitingCheckIOTokens: ) await check.async_log_success_event(kwargs, response, None, None) - current_itpm = await dual_cache.async_get_cache(key=itpm_key) - current_otpm = await dual_cache.async_get_cache(key=otpm_key) + current_itpm = usage.get_local(itpm_key) + current_otpm = usage.get_local(otpm_key) # ITPM tracks input tokens only (billable prompt tokens), not output. assert current_itpm == 5 assert current_otpm == 3 @pytest.mark.asyncio async def test_itpm_limit_raises_429(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -195,8 +217,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) otpm_limit = 10 max_tokens = 4 deployment = { @@ -228,8 +250,8 @@ class TestModelRateLimitingCheckIOTokens: successes = sum(1 for r in results if r) minute = get_utc_datetime().strftime("%H-%M") - otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - current_otpm = await dual_cache.async_get_cache(key=otpm_key) + otpm_key = _io("io-otpm-race-id", "bedrock_mantle/anthropic.claude-opus-4-7", "otpm", minute) + current_otpm = usage.get_local(otpm_key) # Atomic reservation must never let concurrent requests overshoot the limit. assert current_otpm is not None @@ -249,8 +271,8 @@ class TestModelRateLimitingCheckIOTokens: """ from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) itpm_limit = 5 deployment = { "litellm_params": { @@ -281,8 +303,8 @@ class TestModelRateLimitingCheckIOTokens: successes = sum(1 for r in results if r) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - current_itpm = await dual_cache.async_get_cache(key=itpm_key) + itpm_key = _io("io-itpm-estimate-fail-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) + current_itpm = usage.get_local(itpm_key) # A minimal 1-token reservation per request lets itpm_limit concurrent # requests through, instead of a single request starving the rest. @@ -294,11 +316,11 @@ class TestModelRateLimitingCheckIOTokens: async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + itpm_key = _io("io-lp-id", "bedrock_mantle/test", "itpm", minute) + usage.increment_local(CounterIncrement(itpm_key, 20, 60)) # Production kwargs commonly carry litellm_params.metadata; the stashed # reservation lives in the top-level metadata and must still be found. @@ -308,20 +330,20 @@ class TestModelRateLimitingCheckIOTokens: "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, "metadata": {}, }, - "metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 20, ITPM_COUNTER_KEY: asdict(itpm_key)}, "litellm_params": {"metadata": {"user_api_key_hash": "abc123"}}, } await check.async_log_failure_event(kwargs, None, None, None) - current = await dual_cache.async_get_cache(key=itpm_key) + current = usage.get_local(itpm_key) assert current == 0 @pytest.mark.asyncio async def test_reconcile_tracks_actual_usage_when_estimate_zero(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -336,10 +358,10 @@ class TestModelRateLimitingCheckIOTokens: await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + itpm_key = _io("io-zero-est-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) # A failed/zero estimate reserves a minimal 1 token, not the full # itpm limit, so it doesn't starve concurrent requests. - assert await dual_cache.async_get_cache(key=itpm_key) == 1 + assert usage.get_local(itpm_key) == 1 kwargs = { "standard_logging_object": { @@ -355,7 +377,7 @@ class TestModelRateLimitingCheckIOTokens: ) await check.async_log_success_event(kwargs, response, None, None) - assert await dual_cache.async_get_cache(key=itpm_key) == 7 + assert usage.get_local(itpm_key) == 7 @pytest.mark.asyncio async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self): @@ -364,8 +386,8 @@ class TestModelRateLimitingCheckIOTokens: full itpm limit, so up to itpm_limit such calls are allowed concurrently instead of the first one claiming the entire bucket. """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -389,8 +411,8 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -407,9 +429,9 @@ class TestModelRateLimitingCheckIOTokens: set_io_token_rate_limit_request_kwargs(zero_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) - zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] + zero_output_otpm_key = _stashed(zero_output_kwargs["metadata"][OTPM_COUNTER_KEY]) assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0 - assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0 + assert (usage.get_local(zero_output_otpm_key) or 0) == 0 normal_output_kwargs = { "messages": [{"role": "user", "content": "hello"}], @@ -419,8 +441,8 @@ class TestModelRateLimitingCheckIOTokens: set_io_token_rate_limit_request_kwargs(normal_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) - normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] - assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 + normal_output_otpm_key = _stashed(normal_output_kwargs["metadata"][OTPM_COUNTER_KEY]) + assert usage.get_local(normal_output_otpm_key) == 5 @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) @@ -436,10 +458,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_request_without_max_tokens_is_admitted_below_model_max_output( self, model: str, otpm_limit: int, use_sync_check: bool ): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache - ) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": model, "otpm": otpm_limit}, "model_info": {"id": f"io-no-max-tokens-{otpm_limit}"}, @@ -456,7 +476,7 @@ class TestModelRateLimitingCheckIOTokens: reserved: Final = request_kwargs["metadata"][OTPM_RESERVED_KEY] assert isinstance(reserved, int) assert 0 < reserved < otpm_limit - assert await dual_cache.async_get_cache(key=request_kwargs["metadata"][OTPM_CACHE_KEY]) == reserved + assert usage.get_local(_stashed(request_kwargs["metadata"][OTPM_COUNTER_KEY])) == reserved @pytest.mark.asyncio @pytest.mark.parametrize( @@ -466,10 +486,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_concurrent_requests_without_max_tokens_share_otpm_without_overshoot( self, model: str, otpm_limit: int ): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache - ) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": model, "otpm": otpm_limit}, "model_info": {"id": "io-no-max-tokens-race-id"}, @@ -486,12 +504,12 @@ class TestModelRateLimitingCheckIOTokens: return request_kwargs["metadata"] admitted: Final = [m for m in await asyncio.gather(*[_attempt() for _ in range(8)]) if m is not None] - otpm_key: Final = admitted[0][OTPM_CACHE_KEY] + otpm_key: Final = _stashed(admitted[0][OTPM_COUNTER_KEY]) assert len(admitted) == round(1 / NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) - assert {m[OTPM_CACHE_KEY] for m in admitted} == {otpm_key} - assert await dual_cache.async_get_cache(key=otpm_key) == sum(m[OTPM_RESERVED_KEY] for m in admitted) - assert await dual_cache.async_get_cache(key=otpm_key) <= otpm_limit + assert {_stashed(m[OTPM_COUNTER_KEY]) for m in admitted} == {otpm_key} + assert usage.get_local(otpm_key) == sum(m[OTPM_RESERVED_KEY] for m in admitted) + assert usage.get_local(otpm_key) <= otpm_limit @pytest.mark.asyncio @pytest.mark.parametrize( @@ -501,10 +519,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_request_without_max_tokens_reserves_model_max_output_when_otpm_allows( self, model: str, model_max_output: int ): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache - ) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": model, "otpm": model_max_output * 10}, "model_info": {"id": "io-no-max-tokens-roomy-id"}, @@ -523,9 +539,7 @@ class TestModelRateLimitingCheckIOTokens: catalog_max_output: Final = litellm.get_model_info(model=backend_model)["max_output_tokens"] assert isinstance(catalog_max_output, int) otpm_limit: Final = int(max(catalog_max_output, 987) / NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) + 1 - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() - ) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) async def _reserved(deployment_id: str, model_info: dict[str, object]) -> int: metadata: Final[dict[str, object]] = {} @@ -552,9 +566,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_base_model_does_not_change_the_default_otpm_reservation(self, use_sync_check: bool): - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() - ) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) otpm_limit: Final = AZURE_BASE_MODEL_MAX_OUTPUT * 10 async def _reserved(model_info: dict[str, str]) -> int: @@ -579,10 +591,8 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_unparseable_max_tokens_leaves_itpm_uncharged(self, use_sync_check: bool): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache - ) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": MANTLE_OPUS, "itpm": 2000, "otpm": 2000}, "model_info": {"id": "io-bad-max-tokens-id"}, @@ -604,15 +614,13 @@ class TestModelRateLimitingCheckIOTokens: } await _pre_call(valid_kwargs) - itpm_counter: Final = await dual_cache.async_get_cache(key=valid_kwargs["metadata"][ITPM_CACHE_KEY]) + itpm_counter: Final = usage.get_local(_stashed(valid_kwargs["metadata"][ITPM_COUNTER_KEY])) assert itpm_counter == valid_kwargs["metadata"][ITPM_RESERVED_KEY] @pytest.mark.asyncio async def test_explicit_max_tokens_above_default_share_is_reserved_in_full(self): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache - ) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": MANTLE_OPUS, "otpm": 100}, "model_info": {"id": "io-explicit-max-tokens-id"}, @@ -633,9 +641,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_deployment_max_tokens_sets_the_reservation_unless_the_request_sets_one(self, use_sync_check: bool): - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() - ) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) deployment: Final = { "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, "max_tokens": 100}, "model_info": {"id": "io-deployment-max-tokens-id"}, @@ -656,9 +662,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_deployment_cap_sets_the_reservation_only_through_the_param_its_api_sends(self): - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() - ) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) chat: Final = {"messages": [{"role": "user", "content": "hello"}]} responses: Final = {"input": "hello"} @@ -688,9 +692,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_deployment_cap_above_the_default_reserves_the_default(self, use_sync_check: bool): - check: Final = ModelRateLimitingCheck( - usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() - ) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) otpm_limit: Final = UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT * 5 async def reserved(deployment_cap: Mapping[str, int]) -> int: @@ -715,8 +717,8 @@ class TestModelRateLimitingCheckIOTokens: def test_sync_io_pre_call_reserves_and_reconciles(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -735,8 +737,8 @@ class TestModelRateLimitingCheckIOTokens: check.pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + itpm_key = _io("io-sync-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) + otpm_key = _io("io-sync-id", "bedrock_mantle/anthropic.claude-opus-4-7", "otpm", minute) kwargs = { "standard_logging_object": { "model_id": "io-sync-id", @@ -751,15 +753,15 @@ class TestModelRateLimitingCheckIOTokens: ) check.log_success_event(kwargs, response, None, None) - assert dual_cache.get_cache(key=itpm_key) == 5 - assert dual_cache.get_cache(key=otpm_key) == 3 + assert usage.get_local(itpm_key) == 5 + assert usage.get_local(otpm_key) == 3 @pytest.mark.asyncio async def test_reconcile_runs_via_success_event_without_model_id(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) - itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) + itpm_key = _io("io-noid", "bedrock_mantle/test", "itpm", "12-34") + usage.increment_local(CounterIncrement(itpm_key, 10, 60)) # standard_logging_object has no model_id (only the TPM path needs it); # IO reconciliation must still run off the stashed cache key. @@ -768,7 +770,7 @@ class TestModelRateLimitingCheckIOTokens: "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, "metadata": {}, }, - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_COUNTER_KEY: asdict(itpm_key)}, } response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], @@ -776,21 +778,20 @@ class TestModelRateLimitingCheckIOTokens: ) await check.async_log_success_event(kwargs, response, None, None) - assert await dual_cache.async_get_cache(key=itpm_key) == 3 + assert usage.get_local(itpm_key) == 3 @pytest.mark.asyncio async def test_failure_clears_reservation_so_retry_is_not_poisoned(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - usage = RouterUsageStore(connect(None).router_usage()) - check = ModelRateLimitingCheck(usage=usage, dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + itpm_key = _io("io-first", "bedrock_mantle/test", "itpm", minute) + usage.increment_local(CounterIncrement(itpm_key, 8, 60)) # Shared request metadata carrying the first (IO) deployment's reservation. - metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key} + metadata = {ITPM_RESERVED_KEY: 8, ITPM_COUNTER_KEY: asdict(itpm_key)} fail_kwargs = { "metadata": metadata, "standard_logging_object": { @@ -801,9 +802,9 @@ class TestModelRateLimitingCheckIOTokens: } await check.async_log_failure_event(fail_kwargs, None, None, None) - assert await dual_cache.async_get_cache(key=itpm_key) == 0 + assert usage.get_local(itpm_key) == 0 assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata + assert ITPM_COUNTER_KEY not in metadata # Retry succeeds on a non-IO fallback deployment reusing the same metadata. retry_kwargs = { @@ -822,7 +823,7 @@ class TestModelRateLimitingCheckIOTokens: await check.async_log_success_event(retry_kwargs, response, None, None) # The first deployment's ITPM counter is not driven negative... - assert await dual_cache.async_get_cache(key=itpm_key) == 0 + assert usage.get_local(itpm_key) == 0 # ...and the non-IO deployment's TPM usage is tracked normally. tpm = UsageCounter.for_deployment("non-io-second", "openai/gpt-4o-mini", "tpm", minute) assert usage.get_local(tpm) == 12 @@ -841,67 +842,67 @@ class TestModelRateLimitingCheckIOTokens: """ from litellm.utils import get_utc_datetime - dual_cache = DualCache() + usage = _memory_usage() minute = get_utc_datetime().strftime("%H-%M") - itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60) + itpm_key_a = _io("io-retry-a", "bedrock_mantle/test-a", "itpm", minute) + usage.increment_local(CounterIncrement(itpm_key_a, 9, 60)) # Deployment A's still-unreconciled reservation, stashed on the shared # kwargs dict the retry loop reuses. - shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}} + shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_COUNTER_KEY: asdict(itpm_key_a)}} # Router calls this before overwriting kwargs for deployment B's attempt - # simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs. - refund_stale_reservation_before_retry(dual_cache, shared_kwargs) + refund_stale_reservation_before_retry(usage, shared_kwargs) # A's reservation is refunded immediately, not left stranded for a # background failure task that may run arbitrarily later (or never, # if the sentinels get cleared out from under it first). - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + assert usage.get_local(itpm_key_a) == 0 assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"] - assert ITPM_CACHE_KEY not in shared_kwargs["metadata"] + assert ITPM_COUNTER_KEY not in shared_kwargs["metadata"] # A's own (now-late) failure event finds nothing left to refund and # is a safe no-op, since the sentinels were already cleared above. - io_token_refund_failure(dual_cache, shared_kwargs) - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + io_token_refund_failure(usage, shared_kwargs) + assert usage.get_local(itpm_key_a) == 0 # The retry proceeds to stash deployment B's own reservation on the # same dict; it starts clean, unaffected by A's cleared sentinels. set_io_token_rate_limit_request_kwargs(shared_kwargs, metadata_channel="metadata") - itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" + itpm_key_b = _io("io-retry-b", "bedrock_mantle/test-b", "itpm", minute) shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 - shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b - await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60) - assert await dual_cache.async_get_cache(key=itpm_key_b) == 4 + shared_kwargs["metadata"][ITPM_COUNTER_KEY] = asdict(itpm_key_b) + usage.increment_local(CounterIncrement(itpm_key_b, 4, 60)) + assert usage.get_local(itpm_key_b) == 4 @pytest.mark.asyncio async def test_client_supplied_reservation_keys_are_stripped(self): # metadata is caller-controlled; the server-only reservation sentinels # must be removed before the router captures the request kwargs. forged = { - "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"}, + "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_COUNTER_KEY: asdict(_io("attacker", "key", "itpm", "00-00"))}, "litellm_metadata": {OTPM_RESERVED_KEY: 7}, - "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, + "litellm_params": {"metadata": {OTPM_COUNTER_KEY: asdict(_io("attacker", "key", "otpm", "00-00"))}}, } set_io_token_rate_limit_request_kwargs(forged, metadata_channel="metadata") stored = get_io_token_rate_limit_request_kwargs() assert ITPM_RESERVED_KEY not in stored["metadata"] - assert ITPM_CACHE_KEY not in stored["metadata"] + assert ITPM_COUNTER_KEY not in stored["metadata"] assert OTPM_RESERVED_KEY not in stored["litellm_metadata"] - assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"] + assert OTPM_COUNTER_KEY not in stored["litellm_params"]["metadata"] @pytest.mark.asyncio async def test_forged_reservation_cannot_decrement_counter(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) - victim_key = "global_router:victim:model:itpm:00-00" - await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) + victim_key = _io("victim", "model", "itpm", "00-00") + usage.increment_local(CounterIncrement(victim_key, 100, 60)) # A caller forges a reservation pointing at another deployment's counter. kwargs = { - "metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key}, + "metadata": {ITPM_RESERVED_KEY: 100, ITPM_COUNTER_KEY: asdict(victim_key)}, "standard_logging_object": { "model_id": "m", "hidden_params": {"litellm_model_name": "model"}, @@ -918,19 +919,13 @@ class TestModelRateLimitingCheckIOTokens: await check.async_log_success_event(kwargs, response, None, None) # The forged reservation was stripped, so the victim counter is untouched. - assert await dual_cache.async_get_cache(key=victim_key) == 100 + assert usage.get_local(victim_key) == 100 @pytest.mark.asyncio async def test_otpm_reservation_error_rolls_back_itpm(self): from litellm.utils import get_utc_datetime - class _OtpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":otpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _OtpmFailCache() + usage = usage_store(_FailingKind("otpm")) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -950,23 +945,17 @@ class TestModelRateLimitingCheckIOTokens: ) with pytest.raises(RuntimeError): - await async_io_token_pre_call_check(dual_cache, deployment) + await async_io_token_pre_call_check(usage, deployment) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + itpm_key = _io("io-rollback-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) # A transient OTPM error must release the ITPM reservation, not leak it. - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + assert (usage.get_local(itpm_key) or 0) == 0 @pytest.mark.asyncio async def test_reconcile_clears_stash_even_when_increment_errors(self): - class _ItpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":itpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _ItpmFailCache() - metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"} + usage = usage_store(_FailingKind("itpm")) + metadata = {ITPM_RESERVED_KEY: 5, ITPM_COUNTER_KEY: asdict(_io("x", "model", "itpm", "00-00"))} kwargs = {"metadata": metadata} response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], @@ -974,18 +963,18 @@ class TestModelRateLimitingCheckIOTokens: ) with pytest.raises(RuntimeError): - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) # The stash is cleared even though reconciliation raised, so a duplicate # success event can't re-process it. assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata + assert ITPM_COUNTER_KEY not in metadata @pytest.mark.asyncio async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): import logging - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache()) + check = ModelRateLimitingCheck(usage=_memory_usage()) deployment = { "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, "model_info": {}, @@ -1000,7 +989,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_missing_deployment_id_skips_io_reservation(self): - dual_cache = DualCache() + usage = _memory_usage() deployment = { "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100}, "model_info": {}, # no id -> cannot build a per-deployment cache key @@ -1013,7 +1002,7 @@ class TestModelRateLimitingCheckIOTokens: } set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") - result = await async_io_token_pre_call_check(dual_cache, deployment) + result = await async_io_token_pre_call_check(usage, deployment) assert result is deployment # No reservation is stashed, so nothing lands in a shared None:None bucket. @@ -1021,32 +1010,31 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_reconcile_uses_reservation_minute_key(self): - dual_cache = DualCache() + usage = _memory_usage() # Reservation was made on a fixed minute key; a call that finishes in a # later minute must reconcile against that same key, never a key built # from the response-time minute. - itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + itpm_key = _io("io-min-id", "bedrock_mantle/test", "itpm", "99-99") + usage.increment_local(CounterIncrement(itpm_key, 10, 60)) - kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}} + kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_COUNTER_KEY: asdict(itpm_key)}} response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), ) - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) - assert await dual_cache.async_get_cache(key=itpm_key) == 4 + assert usage.get_local(itpm_key) == 4 @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_reconcile", [False, True]) async def test_reconcile_charges_usage_above_the_reservation_to_the_minute_it_finished_in( self, use_sync_reconcile: bool ): - dual_cache: Final = DualCache() - counter: Final = "global_router:io-late-id:bedrock_mantle/test:otpm" - reserved_minute_key: Final = f"{counter}:99-99" - await dual_cache.async_increment_cache(key=reserved_minute_key, value=10, ttl=60) - kwargs: Final = {"metadata": {ITPM_RESERVED_KEY: 0, OTPM_RESERVED_KEY: 10, OTPM_CACHE_KEY: reserved_minute_key}} + usage: Final = _memory_usage() + reserved_minute_key: Final = _io("io-late-id", "bedrock_mantle/test", "otpm", "99-99") + usage.increment_local(CounterIncrement(reserved_minute_key, 10, 60)) + kwargs: Final = {"metadata": {ITPM_RESERVED_KEY: 0, OTPM_RESERVED_KEY: 10, OTPM_COUNTER_KEY: asdict(reserved_minute_key)}} response: Final = ModelResponse( choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], usage=Usage(prompt_tokens=4, completion_tokens=25, total_tokens=29), @@ -1054,39 +1042,41 @@ class TestModelRateLimitingCheckIOTokens: minute_before: Final = get_utc_datetime().strftime("%H-%M") if use_sync_reconcile: - io_token_reconcile_success(dual_cache, kwargs, response) + io_token_reconcile_success(usage, kwargs, response) else: - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) finished_minutes: Final = {minute_before, get_utc_datetime().strftime("%H-%M")} - charged: Final = [await dual_cache.async_get_cache(key=f"{counter}:{minute}") for minute in finished_minutes] - assert await dual_cache.async_get_cache(key=reserved_minute_key) == 10 + charged: Final = [ + usage.get_local(_io("io-late-id", "bedrock_mantle/test", "otpm", minute)) for minute in finished_minutes + ] + assert usage.get_local(reserved_minute_key) == 10 assert sum(value or 0 for value in charged) == 25 - 10 @pytest.mark.asyncio async def test_reconcile_missing_usage_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-missing-usage", "bedrock_mantle/test", "itpm", "00-00") + otpm_key = _io("io-missing-usage", "bedrock_mantle/test", "otpm", "00-00") + usage.increment_local(CounterIncrement(itpm_key, 8, 60)) + usage.increment_local(CounterIncrement(otpm_key, 5, 60)) kwargs = { "metadata": { ITPM_RESERVED_KEY: 8, OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, + ITPM_COUNTER_KEY: asdict(itpm_key), + OTPM_COUNTER_KEY: asdict(otpm_key), } } response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], ) - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 + assert usage.get_local(itpm_key) == 8 + assert usage.get_local(otpm_key) == 5 assert ITPM_RESERVED_KEY not in kwargs["metadata"] @pytest.mark.asyncio @@ -1097,50 +1087,50 @@ class TestModelRateLimitingCheckIOTokens: same as missing usage: keep the reservation instead of resolving to (0, 0) and refunding it in full. """ - dual_cache = DualCache() - itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-total-only", "bedrock_mantle/test", "itpm", "00-00") + otpm_key = _io("io-total-only", "bedrock_mantle/test", "otpm", "00-00") + usage.increment_local(CounterIncrement(itpm_key, 8, 60)) + usage.increment_local(CounterIncrement(otpm_key, 5, 60)) kwargs = { "metadata": { ITPM_RESERVED_KEY: 8, OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, + ITPM_COUNTER_KEY: asdict(itpm_key), + OTPM_COUNTER_KEY: asdict(otpm_key), } } response = {"type": "message", "usage": {"total_tokens": 13}} - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 + assert usage.get_local(itpm_key) == 8 + assert usage.get_local(otpm_key) == 5 def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00" - dual_cache.set_cache(key=itpm_key, value=10, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-slo-total-only", "bedrock_mantle/test", "itpm", "00-00") + usage.set_local(itpm_key, 10) kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_COUNTER_KEY: asdict(itpm_key)}, "standard_logging_object": {"total_tokens": 4}, } response = {"type": "message", "role": "assistant", "content": []} - io_token_reconcile_success(dual_cache, kwargs, response) + io_token_reconcile_success(usage, kwargs, response) - assert dual_cache.get_cache(key=itpm_key) == 10 + assert usage.get_local(itpm_key) == 10 @pytest.mark.asyncio async def test_reconcile_falls_back_to_standard_logging_object(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-slo-fallback", "bedrock_mantle/test", "itpm", "00-00") + usage.increment_local(CounterIncrement(itpm_key, 10, 60)) kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_COUNTER_KEY: asdict(itpm_key)}, "standard_logging_object": { "prompt_tokens": 4, "completion_tokens": 0, @@ -1149,23 +1139,23 @@ class TestModelRateLimitingCheckIOTokens: } response = {"type": "message", "role": "assistant", "content": []} - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) - assert await dual_cache.async_get_cache(key=itpm_key) == 4 + assert usage.get_local(itpm_key) == 4 def test_sync_reconcile_anthropic_dict_usage_counts_uncached_input_once(self): - dual_cache = DualCache() - itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" - dual_cache.set_cache(key=itpm_key, value=10, ttl=60) - dual_cache.set_cache(key=otpm_key, value=4, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-anthropic", "bedrock_mantle/test", "itpm", "00-00") + otpm_key = _io("io-anthropic", "bedrock_mantle/test", "otpm", "00-00") + usage.set_local(itpm_key, 10) + usage.set_local(otpm_key, 4) kwargs = { "metadata": { ITPM_RESERVED_KEY: 10, OTPM_RESERVED_KEY: 4, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, + ITPM_COUNTER_KEY: asdict(itpm_key), + OTPM_COUNTER_KEY: asdict(otpm_key), } } response = { @@ -1178,10 +1168,10 @@ class TestModelRateLimitingCheckIOTokens: }, } - io_token_reconcile_success(dual_cache, kwargs, response) + io_token_reconcile_success(usage, kwargs, response) - assert dual_cache.get_cache(key=itpm_key) == 3 + 4 - assert dual_cache.get_cache(key=otpm_key) == 2 + assert usage.get_local(itpm_key) == 3 + 4 + assert usage.get_local(otpm_key) == 2 @pytest.mark.asyncio async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog): @@ -1189,9 +1179,8 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime - dual_cache = DualCache() - usage = RouterUsageStore(connect(None).router_usage()) - check = ModelRateLimitingCheck(usage=usage, dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) model_id = "io-mixed-id" deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" deployment = { @@ -1206,7 +1195,7 @@ class TestModelRateLimitingCheckIOTokens: minute = get_utc_datetime().strftime("%H-%M") rpm = UsageCounter.for_deployment(model_id, deployment_name, "rpm", minute) - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + itpm_key = _io(model_id, deployment_name, "itpm", minute) usage.increment_local(CounterIncrement(rpm, 5, 60)) request_kwargs = { @@ -1221,7 +1210,7 @@ class TestModelRateLimitingCheckIOTokens: await check.async_pre_call_check(deployment) assert usage.get_local(rpm) == 5, "the rejected request hands its rpm count back" - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + assert (usage.get_local(itpm_key) or 0) == 0 assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] assert any("both limit types are enforced" in record.message for record in caplog.records) @@ -1234,18 +1223,17 @@ class TestModelRateLimitingCheckIOTokens: """ from litellm.utils import get_utc_datetime - dual_cache = DualCache() - usage = RouterUsageStore(connect(None).router_usage()) - check = ModelRateLimitingCheck(usage=usage, dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) model_id = "io-tpm-mixed-id" deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + itpm_key = _io(model_id, deployment_name, "itpm", minute) tpm = UsageCounter.for_deployment(model_id, deployment_name, "tpm", minute) - await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) + usage.increment_local(CounterIncrement(itpm_key, 5, 60)) kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_COUNTER_KEY: asdict(itpm_key)}, "standard_logging_object": { "model_id": model_id, "total_tokens": 7, @@ -1260,25 +1248,24 @@ class TestModelRateLimitingCheckIOTokens: await check.async_log_success_event(kwargs, response, None, None) # ITPM reconciled down from the 5-token reservation to actual usage (3). - assert await dual_cache.async_get_cache(key=itpm_key) == 3 + assert usage.get_local(itpm_key) == 3 # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. assert usage.get_local(tpm) == 7 def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - usage = RouterUsageStore(connect(None).router_usage()) - check = ModelRateLimitingCheck(usage=usage, dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) model_id = "io-tpm-mixed-sync-id" deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + itpm_key = _io(model_id, deployment_name, "itpm", minute) tpm = UsageCounter.for_deployment(model_id, deployment_name, "tpm", minute) - dual_cache.set_cache(key=itpm_key, value=5, ttl=60) + usage.set_local(itpm_key, 5) kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_COUNTER_KEY: asdict(itpm_key)}, "standard_logging_object": { "model_id": model_id, "total_tokens": 7, @@ -1292,20 +1279,20 @@ class TestModelRateLimitingCheckIOTokens: check.log_success_event(kwargs, response, None, None) - assert dual_cache.get_cache(key=itpm_key) == 3 + assert usage.get_local(itpm_key) == 3 assert usage.get_local(tpm) == 7 @pytest.mark.asyncio async def test_failure_refunds_itpm_reservation(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + itpm_key = _io("io-refund-id", "bedrock_mantle/test", "itpm", minute) + usage.increment_local(CounterIncrement(itpm_key, 20, 60)) - reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key} + reservation = {ITPM_RESERVED_KEY: 20, ITPM_COUNTER_KEY: asdict(itpm_key)} kwargs = { "standard_logging_object": { "model_id": "io-refund-id", @@ -1316,7 +1303,7 @@ class TestModelRateLimitingCheckIOTokens: } await check.async_log_failure_event(kwargs, None, None, None) - current = await dual_cache.async_get_cache(key=itpm_key) + current = usage.get_local(itpm_key) assert current == 0 @@ -1341,17 +1328,11 @@ def _current_minute() -> str: return get_utc_datetime().strftime("%H-%M") -def _counter(router: Router, template: RouterCacheEnum, minutes: frozenset[str]) -> int: +def _counter(router: Router, kind: Literal["itpm", "otpm"], minutes: frozenset[str]) -> int: deployment: Final = router.model_list[0] + usage: Final = router._core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the router counts io tokens in return sum( - router.cache.get_cache( - key=template.value.format( - id=deployment["model_info"]["id"], - model=deployment["litellm_params"]["model"], - current_minute=minute, - ) - ) - or 0 + int(usage.get_local(_io(deployment["model_info"]["id"], deployment["litellm_params"]["model"], kind, minute)) or 0) for minute in minutes ) @@ -1360,7 +1341,7 @@ async def _io_counters_since(router: Router, first_minute: str) -> tuple[int, in await asyncio.sleep(0) await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) minutes: Final = frozenset({first_minute, _current_minute()}) - return _counter(router, RouterCacheEnum.ITPM, minutes), _counter(router, RouterCacheEnum.OTPM, minutes) + return _counter(router, "itpm", minutes), _counter(router, "otpm", minutes) async def _sdk_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]: @@ -1561,12 +1542,12 @@ class TestContextSlotRetention: """ kwargs = { "messages": [{"role": "user", "content": "x" * 1000}], - "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, + "metadata": {ITPM_RESERVED_KEY: 999, ITPM_COUNTER_KEY: asdict(_io("forged", "model", "itpm", "00-00"))}, } set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is None assert ITPM_RESERVED_KEY not in kwargs["metadata"] - assert ITPM_CACHE_KEY not in kwargs["metadata"] + assert ITPM_COUNTER_KEY not in kwargs["metadata"] set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is kwargs @@ -1585,7 +1566,7 @@ async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_tok "metadata": metadata, } set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") - await async_io_token_pre_call_check(DualCache(), deployment) + await async_io_token_pre_call_check(_memory_usage(), deployment) return metadata[ITPM_RESERVED_KEY] named = await reserved_input( @@ -1644,21 +1625,21 @@ async def _deliver_success( async def test_success_settles_a_reservation_held_only_in_litellm_params_metadata_exactly_once( deliveries: tuple[str, ...], ) -> None: - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) - itpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:itpm:12-34" - otpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:otpm:12-34" - await dual_cache.async_increment_cache(key=itpm_key, value=40, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=4096, ttl=60) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) + itpm_key: Final = _io("proxy-shape", "openai/gpt-5-mini", "itpm", "12-34") + otpm_key: Final = _io("proxy-shape", "openai/gpt-5-mini", "otpm", "12-34") + usage.increment_local(CounterIncrement(itpm_key, 40, 60)) + usage.increment_local(CounterIncrement(otpm_key, 4096, 60)) kwargs: Final = _proxy_success_kwargs( - {ITPM_RESERVED_KEY: 40, OTPM_RESERVED_KEY: 4096, ITPM_CACHE_KEY: itpm_key, OTPM_CACHE_KEY: otpm_key} + {ITPM_RESERVED_KEY: 40, OTPM_RESERVED_KEY: 4096, ITPM_COUNTER_KEY: asdict(itpm_key), OTPM_COUNTER_KEY: asdict(otpm_key)} ) for delivery in deliveries: await _deliver_success(check, delivery, kwargs, _reply_using(prompt_tokens=7, completion_tokens=3)) - assert await dual_cache.async_get_cache(key=itpm_key) == 7 - assert await dual_cache.async_get_cache(key=otpm_key) == 3 + assert usage.get_local(itpm_key) == 7 + assert usage.get_local(otpm_key) == 3 _TURN_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) @@ -1669,8 +1650,8 @@ def _turn_deployment(kind: str) -> dict[str, object]: return {"litellm_params": {"model": "openai/gpt-5-mini", kind: 100}, "model_info": {"id": f"ws-turn-{kind}"}} -def _turn_key(kind: str, minute: str = "12-34") -> str: - return f"global_router:ws-turn-{kind}:openai/gpt-5-mini:{kind}:{minute}" +def _turn_key(kind: Literal["itpm", "otpm"], minute: str = "12-34") -> UsageCounter: + return _io(f"ws-turn-{kind}", "openai/gpt-5-mini", kind, minute) @pytest.fixture @@ -1684,11 +1665,11 @@ def settled_on_the_turn_s_minute(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.parametrize("kind", ["itpm", "otpm"]) @pytest.mark.parametrize("used", [100, 101, 250]) async def test_turn_admission_refuses_once_the_minutes_tokens_reach_the_limit(kind: str, used: int) -> None: - dual_cache: Final = DualCache() - await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + usage: Final = _memory_usage() + usage.increment_local(CounterIncrement(_turn_key(kind), used, 60)) with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage={used}"): - await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) @pytest.mark.asyncio @@ -1696,19 +1677,18 @@ async def test_turn_admission_refuses_once_the_minutes_tokens_reach_the_limit(ki @pytest.mark.parametrize("kind", ["itpm", "otpm"]) @pytest.mark.parametrize("used", [0, 1, 99]) async def test_an_admitted_turn_reserves_nothing_and_counts_its_actual_tokens(kind: str, used: int) -> None: - dual_cache: Final = DualCache() - await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + usage: Final = _memory_usage() + usage.increment_local(CounterIncrement(_turn_key(kind), used, 60)) actual: Final = _ACTUAL_TOKENS - reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) - counter_key: Final = reservation[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] - assert isinstance(counter_key, str) - assert await dual_cache.async_get_cache(key=counter_key) == used + reservation: Final = await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) + counter_key: Final = _stashed(reservation[ITPM_COUNTER_KEY if kind == "itpm" else OTPM_COUNTER_KEY]) + assert usage.get_local(counter_key) == used await async_io_token_reconcile_success( - dual_cache, _proxy_success_kwargs(reservation), _reply_using(actual["itpm"], actual["otpm"]) + usage, _proxy_success_kwargs(reservation), _reply_using(actual["itpm"], actual["otpm"]) ) - assert await dual_cache.async_get_cache(key=counter_key) == used + actual[kind] + assert usage.get_local(counter_key) == used + actual[kind] @pytest.mark.asyncio @@ -1722,7 +1702,7 @@ async def test_an_admitted_turn_reserves_nothing_and_counts_its_actual_tokens(ki async def test_a_turn_without_io_token_counters_is_admitted_with_nothing_to_settle( deployment: dict[str, object], ) -> None: - assert await async_admit_io_token_turn(DualCache(), deployment, _TURN_AT) == {} + assert await async_admit_io_token_turn(_memory_usage(), deployment, _TURN_AT) == {} @pytest.mark.asyncio @@ -1732,44 +1712,32 @@ async def test_a_turn_without_io_token_counters_is_admitted_with_nothing_to_sett async def test_a_socket_open_reserves_nothing_and_its_first_turn_counts_its_actual_tokens_on_its_own_minute( kind: str, used_by_others: int ) -> None: - dual_cache: Final = DualCache() - await dual_cache.async_increment_cache(key=_turn_key(kind), value=used_by_others, ttl=60) + usage: Final = _memory_usage() + usage.increment_local(CounterIncrement(_turn_key(kind), used_by_others, 60)) socket_open: Final[dict[str, object]] = {"metadata": {}} set_io_token_rate_limit_request_kwargs(socket_open, opens_responses_socket=True, metadata_channel="metadata") - assert await async_io_token_pre_call_check(dual_cache, _turn_deployment(kind)) == _turn_deployment(kind) - assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + assert await async_io_token_pre_call_check(usage, _turn_deployment(kind)) == _turn_deployment(kind) + assert usage.get_local(_turn_key(kind)) == used_by_others opened_with: Final = socket_open["metadata"] assert isinstance(opened_with, dict) await async_io_token_reconcile_success( - dual_cache, _proxy_success_kwargs(opened_with), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]) + usage, _proxy_success_kwargs(opened_with), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]) ) - assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + _ACTUAL_TOKENS[kind] - - -class _RedisCacheOnFakeServer(RedisCache): - def __init__(self, server: FakeServer) -> None: - super().__init__(host="io-token-turn-redis", connection_pool=FakeRedis(server=server).connection_pool) - self._fake_async_client: Final = FakeAsyncRedis(server=server) - - def init_async_client(self) -> FakeAsyncRedis: - return self._fake_async_client + assert usage.get_local(_turn_key(kind)) == used_by_others + _ACTUAL_TOKENS[kind] @pytest.mark.asyncio @pytest.mark.parametrize("kind", ["itpm", "otpm"]) async def test_a_turn_is_refused_on_this_instance_s_count_when_reading_the_shared_counter_fails( - kind: str, + kind: Literal["itpm", "otpm"], ) -> None: - server: Final = FakeServer() - dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnFakeServer(server)) - server.connected = False - await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) + usage: Final = usage_store(FakeUsageHandle(local={_turn_key(kind): 100}, redis={}, redis_down=True)) with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): - await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) @pytest.mark.asyncio @@ -1782,17 +1750,17 @@ async def test_a_turn_finishing_in_the_next_minute_counts_its_actual_tokens_on_t "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: datetime(2026, 1, 1, 12, 35, 10, tzinfo=timezone.utc), ) - dual_cache: Final = DualCache() - reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + usage: Final = _memory_usage() + reservation: Final = await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) await _deliver_success( - ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache), + ModelRateLimitingCheck(usage=usage), delivery, _proxy_success_kwargs(reservation), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]), ) - assert [await dual_cache.async_get_cache(key=_turn_key(kind, minute)) for minute in ("12-34", "12-35")] == [ + assert [usage.get_local(_turn_key(kind, minute)) for minute in ("12-34", "12-35")] == [ None, _ACTUAL_TOKENS[kind], ] @@ -1801,16 +1769,13 @@ async def test_a_turn_finishing_in_the_next_minute_counts_its_actual_tokens_on_t @pytest.mark.asyncio @pytest.mark.parametrize("kind", ["itpm", "otpm"]) async def test_a_turn_is_admitted_once_the_shared_counter_drops_below_the_limit_this_instance_last_read( - kind: str, + kind: Literal["itpm", "otpm"], ) -> None: - redis_cache: Final = _RedisCacheOnFakeServer(FakeServer()) - dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) - await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) - await redis_cache.async_increment(key=_turn_key(kind), value=99) + usage: Final = usage_store(FakeUsageHandle(local={_turn_key(kind): 100}, redis={_turn_key(kind): 99})) - admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + admitted: Final = await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) - assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) + assert _stashed(admitted[ITPM_COUNTER_KEY if kind == "itpm" else OTPM_COUNTER_KEY]) == _turn_key(kind) @pytest.mark.asyncio @@ -1846,36 +1811,51 @@ async def test_a_turn_is_admitted_once_the_shared_counter_drops_below_the_limit_ ) @pytest.mark.usefixtures("settled_on_the_turn_s_minute") async def test_a_turn_counts_only_its_uncached_responses_input_toward_itpm(responses_reply: object) -> None: - dual_cache: Final = DualCache() - await dual_cache.async_increment_cache(key=_turn_key("itpm"), value=10, ttl=60) + usage: Final = _memory_usage() + usage.increment_local(CounterIncrement(_turn_key("itpm"), 10, 60)) - reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment("itpm"), _TURN_AT) - await async_io_token_reconcile_success(dual_cache, _proxy_success_kwargs(reservation), responses_reply) + reservation: Final = await async_admit_io_token_turn(usage, _turn_deployment("itpm"), _TURN_AT) + await async_io_token_reconcile_success(usage, _proxy_success_kwargs(reservation), responses_reply) - assert await dual_cache.async_get_cache(key=_turn_key("itpm")) == 10 + 50 - 30 + assert usage.get_local(_turn_key("itpm")) == 10 + 50 - 30 -def _shared_counter_answering(answer: int | Exception) -> DualCache: - redis_cache: Final = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(side_effect=[answer]) - return DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) @pytest.mark.asyncio @pytest.mark.parametrize("kind", ["itpm", "otpm"]) -async def test_a_later_turn_is_refused_when_another_instance_filled_the_shared_minute(kind: str) -> None: +async def test_a_later_turn_is_refused_when_another_instance_filled_the_shared_minute(kind: Literal["itpm", "otpm"]) -> None: with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): - await async_admit_io_token_turn(_shared_counter_answering(100), _turn_deployment(kind), _TURN_AT) + await async_admit_io_token_turn( + usage_store(FakeUsageHandle(redis={_turn_key(kind): 100})), _turn_deployment(kind), _TURN_AT + ) @pytest.mark.asyncio @pytest.mark.parametrize("kind", ["itpm", "otpm"]) async def test_a_later_turn_is_admitted_on_this_instance_s_count_when_the_shared_counter_is_unreachable( - kind: str, + kind: Literal["itpm", "otpm"], ) -> None: - dual_cache: Final = _shared_counter_answering(RedisCircuitBreakerOpenError("circuit open")) - await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=99) + usage: Final = usage_store(FakeUsageHandle(local={_turn_key(kind): 99}, redis={}, breaker_open=True)) - admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + admitted: Final = await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) - assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) + assert _stashed(admitted[ITPM_COUNTER_KEY if kind == "itpm" else OTPM_COUNTER_KEY]) == _turn_key(kind) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +async def test_a_stashed_reservation_counts_under_mains_redis_keys() -> None: + usage: Final = _memory_usage() + deployment: Final = { + "litellm_params": {"model": "bedrock/claude", "itpm": 100, "otpm": 100}, + "model_info": {"id": "dep-1"}, + } + + reservation: Final = await async_admit_io_token_turn(usage, deployment, _TURN_AT) + await async_io_token_reconcile_success(usage, _proxy_success_kwargs(reservation), _reply_using(5, 2)) + + assert usage.snapshot() == { + "global_router:dep-1:bedrock/claude:itpm:12-34": 5, + "global_router:dep-1:bedrock/claude:otpm:12-34": 2, + } diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 9463424a48b..f4f5243ba5f 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -20,6 +20,7 @@ from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.redis.stores.router_usage import UsageCounter from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook from litellm.types.litellm_params import RoutingStrategyName from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo @@ -347,8 +348,9 @@ def _io_limited_router(kind: str, **litellm_params: object) -> Router: ) -def _minutes_counter_key(kind: str, minute: datetime, model: str = "openai/ws-io-unpriced") -> str: - return f"global_router:ws-io-dep:{model}:{kind}:{minute.strftime('%H-%M')}" +def _minutes_count(router: Router, kind: str, minute: datetime, model: str = "openai/ws-io-unpriced") -> object: + usage: Final = router._core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the router counts io tokens in + return usage.get_local(UsageCounter("global_router", "ws-io-dep", model, minute.strftime("%H-%M"), kind)) _PROVIDER_RESPONSE: Final = { @@ -404,7 +406,7 @@ async def test_responses_call_sends_the_provider_only_the_client_metadata_and_se assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [client_metadata] assert [ - await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned)) + _minutes_count(router, kind, router_minute_pinned) for kind in ("itpm", "otpm") ] == [10, 5] @@ -509,7 +511,7 @@ async def test_retried_call_settles_deployment_io_tokens_to_the_successful_attem assert provider.call_count == 2 assert [ - await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned, model)) + _minutes_count(router, kind, router_minute_pinned, model) for kind in ("itpm", "otpm") ] == [10, 5] @@ -925,7 +927,6 @@ async def test_get_model_group_io_token_usage_sums_across_deployments(): in the model group (not just the first), reading the same per-deployment cache keys the pre-call reservation writes to. """ - from litellm.types.router import RouterCacheEnum from litellm.utils import get_utc_datetime router = Router( @@ -952,26 +953,14 @@ async def test_get_model_group_io_token_usage_sums_across_deployments(): ) minute = get_utc_datetime().strftime("%H-%M") - keys_and_values = [ - ( - RouterCacheEnum.ITPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), - 30, - ), - ( - RouterCacheEnum.OTPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), - 10, - ), - ( - RouterCacheEnum.ITPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), - 70, - ), - ( - RouterCacheEnum.OTPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), - 20, - ), - ] - for key, value in keys_and_values: - await router.cache.async_increment_cache(key=key, value=value, ttl=60) + usage = router._core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the router counts io tokens in + for deployment_id, model, kind, value in ( + ("io-usage-dep-1", "openai/gpt-4o-mini", "itpm", 30), + ("io-usage-dep-1", "openai/gpt-4o-mini", "otpm", 10), + ("io-usage-dep-2", "openai/gpt-4o", "itpm", 70), + ("io-usage-dep-2", "openai/gpt-4o", "otpm", 20), + ): + usage.set_local(UsageCounter("global_router", deployment_id, model, minute, kind), value) current_itpm, current_otpm = await router.get_model_group_io_token_usage("opus")