From 4b0ca0385b5e1ef08b42b2b525bb250b706d5d7c Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 01:49:48 +0000 Subject: [PATCH 01/12] fix(router): count a deployment's requests and tokens once for usage-based-routing-v2 and enforce_model_rate_limits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_rust/core.py | 14 +- litellm/router_rust/invoke/endpoints.py | 9 +- litellm/router_rust/steps/callbacks.py | 9 +- litellm/router_rust/steps/selection.py | 43 +++-- litellm/router_rust/surface/legacy_members.py | 12 +- litellm/router_strategy/lowest_tpm_rpm_v2.py | 53 +----- .../router_utils/deployment_usage_counter.py | 148 +++++++++++++++ .../pre_call_checks/model_rate_limit_check.py | 69 +++---- .../router_strategy/test_lowest_tpm_rpm_v2.py | 63 +++++-- .../test_enforce_model_rate_limits.py | 174 +++++++++--------- 10 files changed, 376 insertions(+), 218 deletions(-) create mode 100644 litellm/router_utils/deployment_usage_counter.py diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index c6ca4a1f399..647e183ffac 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -133,6 +133,7 @@ from litellm.router_rust.steps.selection import ( pick_sync, picked_id, routing_pre_call_checks, + routing_pre_call_checks_sync, ) from litellm.router_rust.surface.alerting import initialize_alerting from litellm.router_rust.surface.fallbacks import ( @@ -206,6 +207,7 @@ from litellm.router_rust.surface.settings import ( from litellm.router_strategy.least_busy import CALL_IDENTITY, call_id_of, owning_picks, pick_reserves, picks_reserve from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET from litellm.router_utils.cooldown_handlers import get_cooldown_deployment_info, mark_deployment_selection_failure +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter from litellm.router_utils.endpoint_eligibility import endpoint_for_call, selected_endpoint from litellm.router_utils.kubernetes_pod_discovery import ( KubernetesPodDiscovery, @@ -1314,6 +1316,9 @@ class _EndpointPort: def routing_callbacks(self) -> Iterable[object]: return self._core._checks.routing_callbacks() # pyright: ignore[reportPrivateUsage] # see the class docstring + def usage_counter(self) -> DeploymentUsageCounter: + return self._core._checks.usage_counter # pyright: ignore[reportPrivateUsage] # see the class docstring + def resolve_alias(self, model: str) -> str | None: return self._core._hosted._core_get_model_from_alias(model) # pyright: ignore[reportPrivateUsage] # see the class docstring @@ -3358,7 +3363,7 @@ class RustCore( selector = self._override_selector(step, request) if selector is not None: for record in self._records_of(step.get("candidates")): - await selector.async_pre_call_check(dict(record), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + await routing_pre_call_checks((selector,), record, self._checks.usage_counter) return step["candidates"] if step.get("stage") == "provider_request": return _provider_request(step) @@ -3392,7 +3397,7 @@ class RustCore( ) ) if picked is not None and selector is not None and self._selectors.is_override(strategy, selector): - await selector.async_pre_call_check(dict(as_str_mapping(picked) or {}), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + await routing_pre_call_checks((selector,), as_str_mapping(picked) or {}, self._checks.usage_counter) if picked is not None and timed: self.service_logged(".async_get_available_deployments", start_time, clock(), request) return _pick_answer(step, candidates, picked) @@ -3419,7 +3424,7 @@ class RustCore( selector = self._override_selector(step, request) if selector is not None: for record in self._records_of(step.get("candidates")): - selector.pre_call_check(dict(record)) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + routing_pre_call_checks_sync((selector,), record, self._checks.usage_counter) return step["candidates"] if step.get("stage") == "provider_request": return _provider_request(step) @@ -3441,7 +3446,7 @@ class RustCore( reserve=pick_reserves(default=False) and bool(request), ) if picked is not None and self._selectors.is_override(strategy, selector): - selector.pre_call_check(dict(as_str_mapping(picked) or {})) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # see above + routing_pre_call_checks_sync((selector,), as_str_mapping(picked) or {}, self._checks.usage_counter) return _pick_answer(step, candidates, picked) # -- cache steps ----------------------------------------------------------------------------- @@ -5173,6 +5178,7 @@ class _Driver: await routing_pre_call_checks( self._core._checks.routing_callbacks(), # pyright: ignore[reportPrivateUsage] # see above as_str_mapping(step.get("deployment")) or {}, + self._core._checks.usage_counter, # pyright: ignore[reportPrivateUsage] # see above ) except Exception as refused: # noqa: BLE001 # handed to the Rust call loop as the attempt's failure # Only `_acompletion` hands the check its logging object. diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index e2c621ccb7d..154f278291d 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -52,6 +52,7 @@ from litellm.router_strategy.complexity_router.context_compaction import ( initialize_compaction_state, surface_for_call, ) +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter from litellm.router_utils.endpoint_eligibility import endpoint_for_call, endpoint_selection if TYPE_CHECKING: @@ -115,6 +116,8 @@ class EndpointCore(Protocol): def routing_callbacks(self) -> Iterable[object]: ... + def usage_counter(self) -> DeploymentUsageCounter: ... + def resolve_alias(self, model: str) -> str | None: ... def group_deployments(self, model_name: str) -> Sequence[Mapping[str, object]] | None: @@ -2447,16 +2450,14 @@ class RustEndpoints: def sync_pre_call_checks(self, plan: Plan, deployment: Mapping[str, object], request: Mapping[str, object]) -> None: """`routing_strategy_pre_call_checks` on the sync path: the process callbacks' pre-call checks. Sync completion runs them for a group but not for a deployment id.""" - from litellm.integrations.custom_logger import CustomLogger + from litellm.router_rust.steps.selection import routing_pre_call_checks_sync if plan.kind == "chat": model: Final = str(request.get("model")) names: Final = {as_dict(r).get("model_name") for r in self._core.handle.deployments()} if model not in names and self._core.handle.deployment(model) is not None: return - for callback in tuple(self._core.routing_callbacks()): - if isinstance(callback, CustomLogger): - callback.pre_call_check(dict(deployment)) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + routing_pre_call_checks_sync(self._core.routing_callbacks(), deployment, self._core.usage_counter()) _SINGLE_ATTEMPT: Final[Mapping[str, object]] = MappingProxyType({"num_retries": 0, "disable_fallbacks": True}) diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index 75afe38548e..b5e597ce42b 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -32,6 +32,7 @@ from litellm.litellm_core_utils.routing_common_utils import get_request_team_id from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.least_busy import LeastBusyLoggingHandler +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, reads_deployment_usage from litellm.router_utils.global_callback_holders import ( live_check_owners, routers_holding_global_callbacks, @@ -115,6 +116,7 @@ class OptionalChecks: self._cache: Final = cache self._hosted: Final = hosted self._ctx: Final = ctx + self.usage_counter: Final = DeploymentUsageCounter(cache, charges=self.reads_deployment_usage) self.optional_callbacks: list[object] | None = None # mutable-ok: the router's live list, as Python keeps it self.router_budget_logger: RouterBudgetLimiting | None = None self.provider_budget_config: object = arguments.get("provider_budget_config") @@ -144,7 +146,9 @@ class OptionalChecks: def build(self, checks: object, model_list: object) -> None: """The constructor's checks: a budget appends ``router_budget_limiting`` to the list given (in place, as Python does), the list is added, a per-group affinity config adds the affinity - check, and ``apply_default_settings`` adds the empty default list.""" + check, and ``apply_default_settings`` adds the empty default list. The usage counter the + checks and selectors read is owned first.""" + self._own(self.usage_counter) listed: list[object] | None = checks if isinstance(checks, list) else None # pyright: ignore[reportUnknownVariableType] # mutable-ok: appended in place, as Python does # rebind-ok: the argument as passed if RouterBudgetLimiting.should_init_router_budget_limiter( # pyright: ignore[reportUnknownMemberType] # untyped model_list model_list=model_list if isinstance(model_list, list) else None, # pyright: ignore[reportUnknownArgumentType] # the argument as passed @@ -212,6 +216,9 @@ class OptionalChecks: session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, # pyright: ignore[reportArgumentType] # same ) + def reads_deployment_usage(self) -> bool: + return reads_deployment_usage(self.routing_callbacks()) + def _own(self, callback: CustomLogger) -> None: if self.optional_callbacks is None: self.optional_callbacks = [] diff --git a/litellm/router_rust/steps/selection.py b/litellm/router_rust/steps/selection.py index 183350bc05f..d634d764b08 100644 --- a/litellm/router_rust/steps/selection.py +++ b/litellm/router_rust/steps/selection.py @@ -30,6 +30,7 @@ from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, reads_deployment_usage from litellm.types.router import RoutingGroup, RoutingStrategy if TYPE_CHECKING: @@ -347,21 +348,41 @@ def pick_sync( return None -async def routing_pre_call_checks(callbacks: Iterable[object], deployment: Mapping[str, object]) -> None: - """``async_routing_strategy_pre_call_checks``: every ``CustomLogger`` among the router's routing - callbacks (``_routing_callbacks``) checks the picked deployment before the call - (usage-based-routing-v2 counts its rpm here). A check's error fails the attempt; the core - records that failure like any other.""" - for callback in tuple(callbacks): - if isinstance(callback, CustomLogger): +async def routing_pre_call_checks( + callbacks: Iterable[object], deployment: Mapping[str, object], counter: DeploymentUsageCounter +) -> None: + """``async_routing_strategy_pre_call_checks``: the attempt's request is counted against the picked + deployment once, then every ``CustomLogger`` among the router's routing callbacks + (``_routing_callbacks``) checks it before the call. A check's error hands the request back and + fails the attempt; the core records that failure like any other.""" + checks: Final = tuple(callback for callback in callbacks if isinstance(callback, CustomLogger)) + counted: Final = reads_deployment_usage(checks) + if counted: + _ = await counter.charge_request(deployment) + try: + for callback in checks: await callback.async_pre_call_check(dict(deployment), None) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict + except Exception: + if counted: + await counter.refund_request(deployment) + raise -def routing_pre_call_checks_sync(callbacks: Iterable[object], deployment: Mapping[str, object]) -> None: - """``routing_strategy_pre_call_checks``: the sync check of each ``CustomLogger`` callback.""" - for callback in tuple(callbacks): - if isinstance(callback, CustomLogger): +def routing_pre_call_checks_sync( + callbacks: Iterable[object], deployment: Mapping[str, object], counter: DeploymentUsageCounter +) -> None: + """``routing_strategy_pre_call_checks``: the sync twin, with the sync check of each ``CustomLogger``.""" + checks: Final = tuple(callback for callback in callbacks if isinstance(callback, CustomLogger)) + counted: Final = reads_deployment_usage(checks) + if counted: + _ = counter.charge_request_sync(deployment) + try: + for callback in checks: callback.pre_call_check(dict(deployment)) # pyright: ignore[reportUnknownMemberType] # see above + except Exception: + if counted: + counter.refund_request_sync(deployment) + raise def picked_id(picked: object) -> str | None: diff --git a/litellm/router_rust/surface/legacy_members.py b/litellm/router_rust/surface/legacy_members.py index 6c30557eb4d..f5773720b1c 100644 --- a/litellm/router_rust/surface/legacy_members.py +++ b/litellm/router_rust/surface/legacy_members.py @@ -36,6 +36,7 @@ from litellm.router_api import RoutingContext from litellm.router_rust.invoke.websocket import admit_deployment_turn from litellm.router_rust.steps.callbacks import RustOptionalChecks from litellm.router_rust.steps.prerouting import TIER_PARAMS_NEVER_DROPPED as _TIER_PARAMS_NEVER_DROPPED +from litellm.router_rust.steps.selection import routing_pre_call_checks_sync from litellm.router_rust.surface.members import RustRegistryMembers from litellm.router_strategy.auto_router.auto_router_model_naming import ( GatedAutoRouterCapability, @@ -331,10 +332,9 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): self, deployment: dict[str, object], # mutable-ok: CustomLogger takes the deployment dict ) -> None: - """``routing_strategy_pre_call_checks``: the sync check of each ``CustomLogger`` callback.""" - for callback in self._opt.routing_callbacks(): - if isinstance(callback, CustomLogger): - callback.pre_call_check(deployment) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict + """``routing_strategy_pre_call_checks``: the request counted once, then the sync check of each + ``CustomLogger`` callback.""" + routing_pre_call_checks_sync(self._opt.routing_callbacks(), deployment, self._opt.usage_counter) async def async_routing_strategy_pre_call_checks( self, @@ -346,11 +346,14 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): under the Python rules (``set_cooldown_deployments``, cooldown event included).""" from litellm import exceptions as litellm_exceptions + counter: Final = self._opt.usage_counter + _ = await counter.charge_request(deployment, parent_otel_span) for callback in self._opt.routing_callbacks(): if isinstance(callback, CustomLogger): try: await callback.async_pre_call_check(deployment, parent_otel_span) # pyright: ignore[reportUnknownMemberType] # see above except litellm_exceptions.RateLimitError as e: + await counter.refund_request(deployment, parent_otel_span) _dispatch_failure(logging_obj, e) _ = set_cooldown_deployments( ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the core's RoutingContext @@ -361,6 +364,7 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): ) raise except Exception as e: + await counter.refund_request(deployment, parent_otel_span) _dispatch_failure(logging_obj, e) raise diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 0574c3d7f69..978daef72c8 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -16,7 +16,6 @@ from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type -from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models from litellm.litellm_core_utils.token_counter import embedding_aware_token_count from litellm.types.router import RouterErrors @@ -133,7 +132,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if deployment_rpm is None: deployment_rpm = float("inf") - if local_result is not None and local_result >= deployment_rpm: + if local_result is not None and local_result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={local_result}", llm_provider="", @@ -154,9 +153,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ), ) else: - # if local result below limit, check redis ## prevent unnecessary redis checks - - result: Final = self.router_cache.increment_cache(key=rpm_key, value=1, ttl=self.routing_args.ttl) + result: Final = self.router_cache.get_cache(key=rpm_key) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", @@ -213,7 +210,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): deployment_rpm = deployment.get("model_info", {}).get("rpm") if deployment_rpm is None: deployment_rpm = float("inf") - if local_result is not None and local_result >= deployment_rpm: + if local_result is not None and local_result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={local_result}", llm_provider="", @@ -230,8 +227,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): num_retries=deployment.get("num_retries"), ) else: - # if local result below limit, check redis ## prevent unnecessary redis checks - result = await self._increment_value_in_current_window(key=rpm_key, value=1, ttl=self.routing_args.ttl) + result = await self.router_cache.async_get_cache(key=rpm_key) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", @@ -259,9 +255,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: - """ - Update TPM/RPM usage on success - """ + """The router's DeploymentUsageCounter records the tokens; the success is only noted here.""" standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") if standard_logging_object is None: raise ValueError("standard_logging_object not passed in.") @@ -273,22 +267,6 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): elif isinstance(id, int): id = str(id) - total_tokens: Final = standard_logging_object.get("total_tokens") - - # ------------ - # Setup values - # ------------ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - - tpm_key: Final = f"{id}:{model}:tpm:{current_minute}" - # ------------ - # Update usage - # ------------ - # update cache - - ## TPM - self.router_cache.increment_cache(key=tpm_key, value=total_tokens, ttl=self.routing_args.ttl) ### TESTING ### if self.test_flag: self.logged_success += 1 @@ -309,32 +287,11 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if standard_logging_object is None: raise ValueError("standard_logging_object not passed in.") model_group: Final = standard_logging_object.get("model_group") - model: Final = standard_logging_object["hidden_params"]["litellm_model_name"] id = standard_logging_object.get("model_id") if model_group is None or id is None: return elif isinstance(id, int): id = str(id) - total_tokens: Final = standard_logging_object.get("total_tokens") - # ------------ - # Setup values - # ------------ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - - tpm_key: Final = f"{id}:{model}:tpm:{current_minute}" - # ------------ - # Update usage - # ------------ - # update cache - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - ## TPM - await self.router_cache.async_increment_cache_post_call( - key=tpm_key, - value=total_tokens, - ttl=self.routing_args.ttl, - parent_otel_span=parent_otel_span, - ) ### TESTING ### if self.test_flag: diff --git a/litellm/router_utils/deployment_usage_counter.py b/litellm/router_utils/deployment_usage_counter.py new file mode 100644 index 00000000000..b550e971dab --- /dev/null +++ b/litellm/router_utils/deployment_usage_counter.py @@ -0,0 +1,148 @@ +"""The one writer of the per-deployment ``{id}:{model}:rpm|tpm:{HH-MM}`` usage counters that +usage-based-routing-v2 selection and ``enforce_model_rate_limits`` read. + +The router charges one request before the routing pre-call checks run, hands it back when a check +refuses the attempt, and charges the tokens of each logged success. Nothing else writes these keys. +""" + +from collections.abc import Callable, Iterable, Mapping +from datetime import datetime +from typing import TYPE_CHECKING, Any, Final + +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.litellm_core_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs +from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck +from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.utils import StandardLoggingPayload +from litellm.utils import get_utc_datetime + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + Span = _Span +else: + Span = Any + +USAGE_TTL: Final = 60 + + +def reads_deployment_usage(callbacks: Iterable[object]) -> bool: + """Whether one of ``callbacks`` reads the per-deployment usage counters: usage-based-routing-v2 + selection or ``enforce_model_rate_limits``. Without a reader there is nothing to count.""" + return any(isinstance(callback, LowestTPMLoggingHandler_v2 | ModelRateLimitingCheck) for callback in callbacks) + + +def usage_counter_keys(deployment: Mapping[str, object], minute: str) -> tuple[str, str]: + """The ``(tpm_key, rpm_key)`` a deployment's usage is counted under this minute.""" + model_info: Final = deployment.get("model_info") + params: Final = deployment.get("litellm_params") + model_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None + model: Final = params.get("model") if isinstance(params, Mapping) else None + return f"{model_id}:{model}:tpm:{minute}", f"{model_id}:{model}:rpm:{minute}" + + +class DeploymentUsageCounter(CustomLogger): + """Counts requests and tokens against a deployment in the router's dual cache, through one Redis + pipeline per write so every worker then holds the shared total. Requests are charged by the + pre-call checks that run a reader; ``charges`` says whether a registered callback reads the + counters, so a logged success counts its tokens only then.""" + + def __init__( + self, + dual_cache: DualCache, + charges: Callable[[], bool] = lambda: True, + clock: Callable[[], datetime] = get_utc_datetime, + ) -> None: + super().__init__() + self.dual_cache: Final = dual_cache + self.charges: Final = charges + self.clock: Final = clock + + def _minute(self) -> str: + return self.clock().strftime("%H-%M") + + def rpm_key(self, deployment: Mapping[str, object]) -> str: + return usage_counter_keys(deployment, self._minute())[1] + + @with_service_target(ROUTER_USAGE_TARGET) + async def charge_request(self, deployment: Mapping[str, object], parent_otel_span: Span | None = None) -> int | None: + """One more request this minute; returns the count, or None when the cache could not say.""" + return await self._async_increment(self.rpm_key(deployment), 1, parent_otel_span) + + @with_service_target(ROUTER_USAGE_TARGET) + async def refund_request(self, deployment: Mapping[str, object], parent_otel_span: Span | None = None) -> None: + """A refused attempt does not count.""" + _ = await self._async_increment(self.rpm_key(deployment), -1, parent_otel_span) + + @with_service_target(ROUTER_USAGE_TARGET) + def charge_request_sync(self, deployment: Mapping[str, object]) -> int | None: + return self._increment(self.rpm_key(deployment), 1) + + @with_service_target(ROUTER_USAGE_TARGET) + def refund_request_sync(self, deployment: Mapping[str, object]) -> None: + _ = self._increment(self.rpm_key(deployment), -1) + + async def _async_increment(self, key: str, value: int, parent_otel_span: Span | None) -> int | None: + operation: Final = RedisPipelineIncrementOperation(key=key, increment_value=value, ttl=USAGE_TTL) + try: + values: Final = await self.dual_cache.async_increment_cache_pipeline( + increment_list=[operation], parent_otel_span=parent_otel_span + ) + except Exception as e: # noqa: BLE001 # a cache that cannot count never fails the call + verbose_router_logger.debug("DeploymentUsageCounter: usage was not recorded. Got - %s", e) + return None + if values is None or len(values) != 1: + return None + count: Final = int(values[0]) + if self.dual_cache.redis_cache is not None: + await self.dual_cache.async_set_cache(key, count, local_only=True, ttl=USAGE_TTL) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + return count + + def _increment(self, key: str, value: int) -> int | None: + try: + count: Final = int(self.dual_cache.increment_cache(key=key, value=value, ttl=USAGE_TTL)) + except Exception as e: # noqa: BLE001 # see _async_increment + verbose_router_logger.debug("DeploymentUsageCounter: usage was not recorded. Got - %s", e) + return None + if self.dual_cache.redis_cache is not None: + self.dual_cache.set_cache(key, count, local_only=True, ttl=USAGE_TTL) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + return count + + def _logged_tokens(self, kwargs: Mapping[str, object]) -> tuple[str, int] | None: + """The ``(tpm_key prefix, total_tokens)`` of a logged success, or None when it counts nothing.""" + if is_batch_retrieve_call_type(kwargs.get("call_type")) or not self.charges(): + return None + logged: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") # pyright: ignore[reportAssignmentType] # the logging payload as CustomLogger receives it + if logged is None: + return None + model_id: Final = logged.get("model_id") + model: Final = logged.get("hidden_params", {}).get("litellm_model_name") + tokens: Final = logged.get("total_tokens") + if model_id is None or not model or not tokens: + return None + return f"{model_id}:{model}", int(tokens) + + @with_service_target(ROUTER_USAGE_TARGET) + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): # pyright: ignore[reportMissingParameterType] # CustomLogger's untyped hook + counted: Final = self._logged_tokens(kwargs) # pyright: ignore[reportUnknownArgumentType] # see above + if counted is None: + return + prefix, tokens = counted + _ = await self._async_increment(f"{prefix}:tpm:{self._minute()}", tokens, get_parent_otel_span_from_kwargs(kwargs)) # pyright: ignore[reportUnknownArgumentType] # see above + + @with_service_target(ROUTER_USAGE_TARGET) + def log_success_event(self, kwargs, response_obj, start_time, end_time): # pyright: ignore[reportMissingParameterType] # CustomLogger's untyped hook + counted: Final = self._logged_tokens(kwargs) # pyright: ignore[reportUnknownArgumentType] # see above + if counted is None: + return + prefix, tokens = counted + _ = self._increment(f"{prefix}:tpm:{self._minute()}", tokens) + + +__all__ = ["USAGE_TTL", "DeploymentUsageCounter", "reads_deployment_usage", "usage_counter_keys"] 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 a3da0ae1674..fcbab1c1a36 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 @@ -316,10 +316,10 @@ class ModelRateLimitingCheck(CustomLogger): io_reservation_made = True dt: Final = self.clock() + current_minute: Final = dt.strftime("%H-%M") daily: Final = _DailyLimitDeployment.model_validate(deployment) tpm_limit, rpm_limit = self._get_deployment_limits(deployment) - current_minute: Final = dt.strftime("%H-%M") tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) model_id: Final = deployment.get("model_info", {}).get("id") @@ -345,11 +345,9 @@ class ModelRateLimitingCheck(CustomLogger): ), ) - # Check RPM limit (atomic increment-first to avoid race conditions) if rpm_limit is not None: - current_rpm: Final = self._increment(rpm_key, 1, RoutingArgs.ttl) + current_rpm: Final = self.dual_cache.get_cache(key=rpm_key) if current_rpm is not None and current_rpm > rpm_limit: - self._increment(rpm_key, -1, RoutingArgs.ttl) raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", llm_provider="", @@ -365,12 +363,7 @@ class ModelRateLimitingCheck(CustomLogger): ), ) - try: - self._check_daily_limits(daily, dt) - except litellm.RateLimitError: - if rpm_limit is not None: - self._increment(rpm_key, -1, RoutingArgs.ttl) - raise + self._check_daily_limits(daily, dt) return deployment except litellm.RateLimitError: @@ -401,11 +394,11 @@ class ModelRateLimitingCheck(CustomLogger): io_reservation_made = True dt: Final = self.clock() + current_minute: Final = dt.strftime("%H-%M") daily: Final = _DailyLimitDeployment.model_validate(deployment) tpd_limit, rpd_limit = daily.limits() tpm_limit, rpm_limit = self._get_deployment_limits(deployment) - current_minute: Final = dt.strftime("%H-%M") tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) tpd_key, rpd_key = daily.cache_keys(dt) @@ -438,32 +431,30 @@ class ModelRateLimitingCheck(CustomLogger): if current_tpd is not None and current_tpd >= tpd_limit: raise daily.limit_error("tpd", tpd_limit, current_tpd, dt) - charged_rpm_key: Final = rpm_key if rpm_limit is not None else None - charged_rpd_key: Final = rpd_key if rpd_limit is not None else None - current_rpm, current_rpd = await self._async_increment_request_counts( - charged_rpm_key, charged_rpd_key, parent_otel_span - ) - if rpm_limit is not None and current_rpm is not None and current_rpm > rpm_limit: - await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) - raise litellm.RateLimitError( - message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", - llm_provider="", - model=model_name, - response=httpx.Response( - status_code=429, - content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", - headers={"retry-after": str(60)}, - request=httpx.Request( - method="model_rate_limit_check", - url="https://github.com/BerriAI/litellm", + if rpm_limit is not None: + current_rpm: Final = await self._async_get_current_tpm(rpm_key, rpm_limit, parent_otel_span) + if current_rpm is not None and current_rpm > rpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), ), - ), - num_retries=0, # Don't retry - return 429 immediately - ) + num_retries=0, # Don't retry - return 429 immediately + ) - if rpd_limit is not None and current_rpd is not None and current_rpd > rpd_limit: - await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) - raise daily.limit_error("rpd", rpd_limit, current_rpd, dt) + if rpd_limit is not None: + _, current_rpd = await self._async_increment_request_counts(None, rpd_key, parent_otel_span) + if current_rpd is not None and current_rpd > rpd_limit: + await self._async_refund_request_counts(None, rpd_key, parent_otel_span) + raise daily.limit_error("rpd", rpd_limit, current_rpd, dt) return deployment @@ -509,15 +500,10 @@ class ModelRateLimitingCheck(CustomLogger): return dt: Final = self.clock() - current_minute: Final = dt.strftime("%H-%M") - tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" reset_offset: Final = _reset_offset_from_kwargs(kwargs) # pyright: ignore[reportUnknownArgumentType] # CustomLogger passes an untyped kwargs dict tpd_key: Final = f"{model_id}:{model}:tpd:{daily_window_day(dt, reset_offset)}" - verbose_router_logger.debug("[TPM TRACKING] Incrementing %s by %s", tpm_key, total_tokens) - tokens: Final = int(total_tokens) - await self._async_increment(tpm_key, tokens, RoutingArgs.ttl) await self._async_increment(tpd_key, tokens, RoutingArgs.daily_ttl) except Exception as e: @@ -560,13 +546,10 @@ class ModelRateLimitingCheck(CustomLogger): return dt: Final = self.clock() - current_minute: Final = dt.strftime("%H-%M") - tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" reset_offset: Final = _reset_offset_from_kwargs(kwargs) # pyright: ignore[reportUnknownArgumentType] # CustomLogger passes an untyped kwargs dict tpd_key: Final = f"{model_id}:{model}:tpd:{daily_window_day(dt, reset_offset)}" tokens: Final = int(total_tokens) - self._increment(tpm_key, tokens, RoutingArgs.ttl) self._increment(tpd_key, tokens, RoutingArgs.daily_ttl) except Exception as e: diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py index f04ade1ce2d..bda43a87fbd 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py @@ -8,12 +8,8 @@ import io import os import time import traceback -from typing import Dict, Final, Iterator - -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map -from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome - +from collections.abc import Iterator +from typing import Final from unittest.mock import patch import pytest @@ -23,12 +19,16 @@ import litellm from litellm import Router from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.router_rust.steps.selection import routing_pre_call_checks_sync from litellm.router_strategy.lowest_tpm_rpm_v2 import ( LowestTPMLoggingHandler_v2 as LowestTPMLoggingHandler, ) +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter from litellm.types.router import DeploymentTypedDict, RouterRateLimitError from litellm.types.utils import StandardLoggingPayload -from litellm.utils import get_utc_datetime +from litellm.utils import _invalidate_model_cost_lowercase_map, get_utc_datetime +from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome ### UNIT TESTS FOR TPM/RPM ROUTING ### @@ -45,6 +45,17 @@ def restore_verbose_logger_level() -> Iterator[None]: @pytest.mark.usefixtures("restore_verbose_logger_level") + +def _router_usage_counter(router: Router) -> DeploymentUsageCounter: + return next(callback for callback in router.optional_callbacks or () if isinstance(callback, DeploymentUsageCounter)) + + +def _log_success(handler: LowestTPMLoggingHandler, counter: DeploymentUsageCounter, **kwargs) -> None: + """A logged success as the router fans it out: the counter's tokens, then the strategy's hook.""" + counter.log_success_event(**kwargs) + handler.log_success_event(**kwargs) + + def test_tpm_rpm_updated(): from tests.unit.create_mock_standard_logging_payload import create_standard_logging_payload @@ -80,8 +91,10 @@ def test_tpm_rpm_updated(): start_time = time.time() response_obj = {"usage": {"total_tokens": total_tokens}} end_time = time.time() - lowest_tpm_logger.pre_call_check(deployment=litellm_deployment_dict) - lowest_tpm_logger.log_success_event( + routing_pre_call_checks_sync((lowest_tpm_logger,), litellm_deployment_dict, DeploymentUsageCounter(test_cache)) + _log_success( + lowest_tpm_logger, + DeploymentUsageCounter(test_cache), response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -140,7 +153,9 @@ def test_get_available_deployments(): start_time = time.time() response_obj = {"usage": {"total_tokens": total_tokens}} end_time = time.time() - lowest_tpm_logger.log_success_event( + _log_success( + lowest_tpm_logger, + DeploymentUsageCounter(test_cache), response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -167,7 +182,9 @@ def test_get_available_deployments(): start_time = time.time() response_obj = {"usage": {"total_tokens": total_tokens}} end_time = time.time() - lowest_tpm_logger.log_success_event( + _log_success( + lowest_tpm_logger, + DeploymentUsageCounter(test_cache), response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -245,7 +262,9 @@ def test_router_get_available_deployments(): start_time = time.time() response_obj = {"usage": {"total_tokens": total_tokens}} end_time = time.time() - router.lowesttpm_logger_v2.log_success_event( + _log_success( + router.lowesttpm_logger_v2, + _router_usage_counter(router), response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -269,7 +288,9 @@ def test_router_get_available_deployments(): start_time = time.time() response_obj = {"usage": {"total_tokens": 20}} end_time = time.time() - router.lowesttpm_logger_v2.log_success_event( + _log_success( + router.lowesttpm_logger_v2, + _router_usage_counter(router), response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -331,7 +352,9 @@ def test_router_skip_rate_limited_deployments(): start_time = time.time() response_obj = {"usage": {"total_tokens": total_tokens}} end_time = time.time() - router.lowesttpm_logger_v2.log_success_event( + _log_success( + router.lowesttpm_logger_v2, + _router_usage_counter(router), response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -345,9 +368,9 @@ def test_router_skip_rate_limited_deployments(): model="azure-model", messages=[{"role": "user", "content": "Hey, how's it going?"}], ) - pytest.fail(f"Should have raised No Models Available error") + pytest.fail("Should have raised No Models Available error") except Exception as e: - print(f"An exception occurred! {str(e)}") + print(f"An exception occurred! {e!s}") def _embedding_image_uri() -> str: @@ -375,7 +398,7 @@ def _embedding_image_uri() -> str: ids=["image-model", "image-base-model", "text-model"], ) async def test_deployment_tpm_counts_embedding_image_as_media_only_on_image_models( - routing_strategy: str, litellm_params: Dict[str, str], model_info: Dict[str, str], admitted: bool + routing_strategy: str, litellm_params: dict[str, str], model_info: dict[str, str], admitted: bool ) -> None: uri: Final = _embedding_image_uri() assert len(uri) // 4 > 75_000 @@ -485,9 +508,9 @@ def test_single_deployment_tpm_zero(): model=model, messages=[{"role": "user", "content": "Hey, how's it going?"}], ) - pytest.fail(f"Should have raised No Models Available error") + pytest.fail("Should have raised No Models Available error") except Exception as e: - print(f"it worked - {str(e)}! \n{traceback.format_exc()}") + print(f"it worked - {e!s}! \n{traceback.format_exc()}") @pytest.mark.asyncio @@ -573,7 +596,7 @@ def test_return_potential_deployments(): test_cache = DualCache() lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) - args: Dict = { + args: dict = { "healthy_deployments": [ { "model_name": "model-test", 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 ff3ffc0db0c..9a89d25ba4a 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -17,6 +17,8 @@ import litellm from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.router_rust.steps.selection import routing_pre_call_checks, routing_pre_call_checks_sync +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) @@ -39,6 +41,32 @@ def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> return dual_cache +def _counter(check: ModelRateLimitingCheck) -> DeploymentUsageCounter: + return DeploymentUsageCounter(check.dual_cache, clock=check.clock) + + +def _pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: + """The router's pre-call step: the request counted once, then the check.""" + routing_pre_call_checks_sync((check,), deployment, _counter(check)) + return deployment + + +async def _async_pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: + await routing_pre_call_checks((check,), deployment, _counter(check)) + return deployment + + +def _log_success(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: + """A logged success as the router fans it out: the counter's tokens, then the check's daily tokens.""" + _counter(check).log_success_event(kwargs, None, None, None) + check.log_success_event(kwargs, None, None, None) + + +async def _async_log_success(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: + await _counter(check).async_log_success_event(kwargs, None, None, None) + await check.async_log_success_event(kwargs, None, None, None) + + class TestModelRateLimitingCheck: """Test the ModelRateLimitingCheck class directly.""" @@ -105,16 +133,12 @@ class TestModelRateLimitingCheck: "model_info": {"id": "test-id"}, } - result = check.pre_call_check(deployment) + result = _pre_call(check, deployment) assert result == deployment def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): - """Test that RateLimitError is raised when RPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 11 # Over limit after increment - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - + """The eleventh request of the minute is the first one over rpm 10.""" + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) deployment = { "rpm": 10, "litellm_params": {"model": "gpt-4"}, @@ -122,8 +146,10 @@ class TestModelRateLimitingCheck: "model_name": "test-model", } + for _ in range(10): + assert _pre_call(check, deployment) == deployment with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) + _pre_call(check, deployment) assert "RPM limit=10" in str(exc_info.value) assert "current usage=11" in str(exc_info.value) @@ -142,7 +168,7 @@ class TestModelRateLimitingCheck: "model_name": "test-model", } - result = check.pre_call_check(deployment) + result = _pre_call(check, deployment) assert result == deployment def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): @@ -160,7 +186,7 @@ class TestModelRateLimitingCheck: } with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) + _pre_call(check, deployment) assert "TPM limit=1000" in str(exc_info.value) assert "current usage=1000" in str(exc_info.value) @@ -171,7 +197,7 @@ class TestModelRateLimitingCheck: check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(TPM_DEPLOYMENT) + _pre_call(check, TPM_DEPLOYMENT) assert "current usage=1000" in str(exc_info.value) @@ -184,7 +210,7 @@ class TestModelRateLimitingCheck: check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(TPM_DEPLOYMENT) + _pre_call(check, TPM_DEPLOYMENT) assert "current usage=1000" in str(exc_info.value) @@ -195,7 +221,7 @@ class TestModelRateLimitingCheck: check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + _pre_call(check, {**TPM_DEPLOYMENT, "rpm": 1}) assert "RPM limit=1" in str(exc_info.value) @@ -203,16 +229,16 @@ class TestModelRateLimitingCheck: check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) deployment = {**TPM_DEPLOYMENT, "rpm": 1} - assert check.pre_call_check(deployment) == deployment + assert _pre_call(check, deployment) == deployment with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) + _pre_call(check, deployment) assert "RPM limit=1" in str(exc_info.value) - def test_log_success_event_increments_cache(self): - """Test that log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) + def test_log_success_event_counts_the_tokens_once(self): + """A logged success adds its tokens to the minute once, although the counter and the check both see it.""" + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache, clock=Clock(DAY_ONE)) kwargs = { "standard_logging_object": { @@ -222,11 +248,9 @@ class TestModelRateLimitingCheck: } } - check.log_success_event(kwargs, None, None, None) + _log_success(check, kwargs) - keys = tuple(call.kwargs["key"] for call in mock_cache.increment_cache.call_args_list) - assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys - assert all(call.kwargs["value"] == 50 for call in mock_cache.increment_cache.call_args_list) + assert dual_cache.get_cache("test-id:gpt-4:tpm:23-59") == 50 class TestModelRateLimitingCheckAsync: @@ -237,6 +261,7 @@ class TestModelRateLimitingCheckAsync: """Test that requests are allowed when no limits are set (async).""" mock_cache = MagicMock() mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() check = ModelRateLimitingCheck(dual_cache=mock_cache) @@ -245,18 +270,13 @@ class TestModelRateLimitingCheckAsync: "model_info": {"id": "test-id"}, } - result = await check.async_pre_call_check(deployment) + result = await _async_pre_call(check, deployment) assert result == deployment @pytest.mark.asyncio async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): - """Test that RateLimitError is raised when RPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[11.0]) # Over limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - + """The eleventh request of the minute is the first one over rpm 10 (async).""" + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) deployment = { "rpm": 10, "litellm_params": {"model": "gpt-4"}, @@ -264,20 +284,18 @@ class TestModelRateLimitingCheckAsync: "model_name": "test-model", } + for _ in range(10): + assert await _async_pre_call(check, deployment) == deployment with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) + await _async_pre_call(check, deployment) assert "RPM limit=10" in str(exc_info.value) + assert "current usage=11" in str(exc_info.value) @pytest.mark.asyncio async def test_async_pre_call_check_allows_request_under_limit(self): - """Test that requests are allowed when under the limit (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[6.0]) - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - + """Every request up to rpm 10 is allowed (async).""" + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) deployment = { "rpm": 10, "litellm_params": {"model": "gpt-4"}, @@ -285,8 +303,8 @@ class TestModelRateLimitingCheckAsync: "model_name": "test-model", } - result = await check.async_pre_call_check(deployment) - assert result == deployment + for _ in range(10): + assert await _async_pre_call(check, deployment) == deployment @pytest.mark.asyncio async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): @@ -304,7 +322,7 @@ class TestModelRateLimitingCheckAsync: } with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) + await _async_pre_call(check, deployment) assert "TPM limit=1000" in str(exc_info.value) @@ -315,7 +333,7 @@ class TestModelRateLimitingCheckAsync: check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(TPM_DEPLOYMENT) + await _async_pre_call(check, TPM_DEPLOYMENT) assert "current usage=1000" in str(exc_info.value) @@ -329,7 +347,7 @@ class TestModelRateLimitingCheckAsync: check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(TPM_DEPLOYMENT) + await _async_pre_call(check, TPM_DEPLOYMENT) assert "current usage=1000" in str(exc_info.value) @@ -343,7 +361,7 @@ class TestModelRateLimitingCheckAsync: check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + await _async_pre_call(check, {**TPM_DEPLOYMENT, "rpm": 1}) assert "RPM limit=1" in str(exc_info.value) @@ -352,18 +370,17 @@ class TestModelRateLimitingCheckAsync: check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) deployment = {**TPM_DEPLOYMENT, "rpm": 1} - assert await check.async_pre_call_check(deployment) == deployment + assert await _async_pre_call(check, deployment) == deployment with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) + await _async_pre_call(check, deployment) assert "RPM limit=1" in str(exc_info.value) @pytest.mark.asyncio - async def test_async_log_success_event_increments_cache(self): - """Test that async_log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - mock_cache.async_increment_cache = AsyncMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) + async def test_async_log_success_event_counts_the_tokens_once(self): + """A logged success adds its tokens to the minute once, although the counter and the check both see it.""" + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache, clock=Clock(DAY_ONE)) kwargs = { "standard_logging_object": { @@ -373,11 +390,9 @@ class TestModelRateLimitingCheckAsync: } } - await check.async_log_success_event(kwargs, None, None, None) + await _async_log_success(check, kwargs) - keys = tuple(call.kwargs["key"] for call in mock_cache.async_increment_cache.call_args_list) - assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys - assert all(call.kwargs["value"] == 50 for call in mock_cache.async_increment_cache.call_args_list) + assert await dual_cache.async_get_cache("test-id:gpt-4:tpm:23-59") == 50 class TestRouterWithEnforceModelRateLimits: @@ -400,8 +415,7 @@ class TestRouterWithEnforceModelRateLimits: # Check that the callback was added assert router.optional_callbacks is not None - assert len(router.optional_callbacks) == 1 - assert isinstance(router.optional_callbacks[0], ModelRateLimitingCheck) + assert any(isinstance(callback, ModelRateLimitingCheck) for callback in router.optional_callbacks) def test_router_optional_callbacks_contains_model_rate_limiting(self): """Test that ModelRateLimitingCheck is in the callbacks list.""" @@ -452,7 +466,7 @@ class TestModelRateLimitConcurrency: } async def attempt_request(): - return await check.async_pre_call_check(deployment) + return await _async_pre_call(check, deployment) results = await asyncio.gather( *[attempt_request() for _ in range(4)], @@ -483,16 +497,6 @@ class Clock: return self.now -def _pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: - return check.pre_call_check(deployment) - - -async def _async_pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: - return await check.async_pre_call_check(deployment) - - -async def _async_log_success(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: - await check.async_log_success_event(kwargs, None, None, None) class TestDailyLimits: @@ -606,36 +610,40 @@ class TestDailyLimits: await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) @pytest.mark.asyncio - async def test_async_pre_call_check_refunds_rpm_and_rpd_in_one_cache_trip(self): + async def test_async_pre_call_check_rejected_over_rpd_hands_back_rpd_then_rpm(self): mock_cache = MagicMock() mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(side_effect=[[1.0, 3.0], [0.0, 2.0]]) + mock_cache.async_set_cache = AsyncMock() + mock_cache.async_increment_cache_pipeline = AsyncMock(side_effect=[[1.0], [3.0], [2.0], [0.0]]) + mock_cache.redis_cache = None check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) - assert mock_cache.async_increment_cache_pipeline.await_count == 2 - refund = mock_cache.async_increment_cache_pipeline.call_args_list[1].kwargs["increment_list"] - assert [(op["key"], op["increment_value"], op["ttl"]) for op in refund] == [ - ("daily-test-id:gpt-4.1-mini:rpm:23-59", -1, 60), - ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", -1, 86400), + trips = [call.kwargs["increment_list"] for call in mock_cache.async_increment_cache_pipeline.call_args_list] + assert [[(op["key"], op["increment_value"], op["ttl"]) for op in trip] for trip in trips] == [ + [("daily-test-id:gpt-4.1-mini:rpm:23-59", 1, 60)], + [("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", 1, 86400)], + [("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", -1, 86400)], + [("daily-test-id:gpt-4.1-mini:rpm:23-59", -1, 60)], ] @pytest.mark.asyncio - async def test_async_pre_call_check_increments_rpm_and_rpd_in_one_cache_trip(self): + async def test_async_pre_call_check_charges_rpm_then_rpd(self): mock_cache = MagicMock() mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[1.0, 1.0]) + mock_cache.async_set_cache = AsyncMock() + mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[1.0]) + mock_cache.redis_cache = None check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT - assert mock_cache.async_increment_cache_pipeline.await_count == 1 - operations = mock_cache.async_increment_cache_pipeline.call_args.kwargs["increment_list"] - assert [(op["key"], op["ttl"]) for op in operations] == [ - ("daily-test-id:gpt-4.1-mini:rpm:23-59", 60), - ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", 86400), + trips = [call.kwargs["increment_list"] for call in mock_cache.async_increment_cache_pipeline.call_args_list] + assert [[(op["key"], op["ttl"]) for op in trip] for trip in trips] == [ + [("daily-test-id:gpt-4.1-mini:rpm:23-59", 60)], + [("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", 86400)], ] assert not mock_cache.async_increment_cache.called From 915a27b96575ff3e228e242b58ac15473fd8c2c5 Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 01:51:50 +0000 Subject: [PATCH 02/12] test(router): ticket regression, counter tests, v2 selection admits the rpm-th request Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/lowest_tpm_rpm_v2.py | 2 +- .../test_router_routing_groups.py | 61 +++++++++++- .../test_deployment_usage_counter.py | 92 +++++++++++++++++++ 3 files changed, 153 insertions(+), 2 deletions(-) create mode 100644 tests/unit/router_utils/test_deployment_usage_counter.py diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 978daef72c8..ce92506bdd2 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -344,7 +344,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if item_tpm + input_tokens > _deployment_tpm or ( (rpm_dict is not None and item in rpm_dict) and rpm_dict[item] is not None - and (rpm_dict[item] + 1 >= _deployment_rpm) + and (rpm_dict[item] + 1 > _deployment_rpm) ): continue elif item_tpm == lowest_tpm: diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 8a0d5b9a1d6..1cbbe78573d 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -10,7 +10,6 @@ import datetime import uuid from collections.abc import Mapping from typing import Final, Literal -from unittest.mock import patch import pytest from pydantic import ValidationError @@ -21,6 +20,7 @@ from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.router_api import RoutingContext from litellm.router_strategy.least_busy import picks_reserve +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup from litellm.utils import Rules, function_setup @@ -1150,3 +1150,62 @@ async def test_update_settings_changed_routing_strategy_args_flushes_replaced_se await t except asyncio.CancelledError: pass + + +def _rate_limited_deployment(model_name: str, deployment_id: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test", "mock_response": "ok", "rpm": 10}, + "model_info": {"id": deployment_id}, + } + + +def _ticket_router() -> Router: + """``enforce_model_rate_limits`` next to a usage-based-routing-v2 group, as the customer ran it.""" + return Router( + model_list=[_rate_limited_deployment("lat-model", "lat-a"), _rate_limited_deployment("ubr-model", "ubr-a")], + routing_groups=[ + {"group_name": "latency-group", "models": ["lat-model"], "routing_strategy": "latency-based-routing"}, + {"group_name": "ubr-group", "models": ["ubr-model"], "routing_strategy": "usage-based-routing-v2"}, + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + num_retries=0, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["lat-model", "ubr-model"]) +async def test_a_deployment_with_rpm_10_serves_ten_requests_a_minute_next_to_a_usage_based_v2_group(model: str): + router = _ticket_router() + counter = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentUsageCounter)) + deployment = router.get_deployment(model_id=f"{model[:3]}-a") + assert deployment is not None + + outcomes = [ + await _outcome(router, model) for _ in range(12) + ] + + assert outcomes == [200] * 10 + [429] * 2, outcomes + assert await router.cache.async_get_cache(counter.rpm_key(deployment.model_dump())) == 10 + + +async def _outcome(router: Router, model: str) -> int: + try: + await router.acompletion(model=model, messages=[{"role": "user", "content": "hi"}]) + except litellm.RateLimitError: + return 429 + return 200 + + +@pytest.mark.asyncio +async def test_a_success_adds_its_tokens_to_the_deployment_once_next_to_a_usage_based_v2_group(): + router = _ticket_router() + counter = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentUsageCounter)) + deployment = router.get_deployment(model_id="lat-a") + assert deployment is not None + + response = await router.acompletion(model="lat-model", messages=[{"role": "user", "content": "hi"}]) + await asyncio.sleep(0.2) + + tpm_key: Final = counter.rpm_key(deployment.model_dump()).replace(":rpm:", ":tpm:") + assert await router.cache.async_get_cache(tpm_key) == response.usage.total_tokens diff --git a/tests/unit/router_utils/test_deployment_usage_counter.py b/tests/unit/router_utils/test_deployment_usage_counter.py new file mode 100644 index 00000000000..d906908a49e --- /dev/null +++ b/tests/unit/router_utils/test_deployment_usage_counter.py @@ -0,0 +1,92 @@ +from datetime import datetime, timezone + +import pytest + +from litellm.caching.dual_cache import DualCache +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, reads_deployment_usage +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + +NOON = datetime(2026, 10, 28, 12, 0, 30, tzinfo=timezone.utc) +DEPLOYMENT = {"litellm_params": {"model": "openai/gpt-4o"}, "model_info": {"id": "deploy-1"}} +RPM_KEY = "deploy-1:openai/gpt-4o:rpm:12-00" +TPM_KEY = "deploy-1:openai/gpt-4o:tpm:12-00" +SUCCESS = { + "standard_logging_object": { + "model_id": "deploy-1", + "total_tokens": 40, + "hidden_params": {"litellm_model_name": "openai/gpt-4o"}, + } +} + + +def _counter(cache: DualCache, charges: bool = True) -> DeploymentUsageCounter: + return DeploymentUsageCounter(cache, charges=lambda: charges, clock=lambda: NOON) + + +@pytest.mark.asyncio +async def test_each_charged_request_adds_one_to_the_minute(): + cache = DualCache() + counter = _counter(cache) + + counts = [await counter.charge_request(DEPLOYMENT) for _ in range(3)] + + assert counts == [1, 2, 3] + assert await cache.async_get_cache(RPM_KEY) == 3 + + +@pytest.mark.asyncio +async def test_a_refunded_request_leaves_the_minute_as_it_was(): + cache = DualCache() + counter = _counter(cache) + await counter.charge_request(DEPLOYMENT) + await counter.charge_request(DEPLOYMENT) + + await counter.refund_request(DEPLOYMENT) + + assert await cache.async_get_cache(RPM_KEY) == 1 + + +def test_sync_charges_and_refunds_count_like_the_async_ones(): + cache = DualCache() + counter = _counter(cache) + + assert [counter.charge_request_sync(DEPLOYMENT) for _ in range(2)] == [1, 2] + counter.refund_request_sync(DEPLOYMENT) + + assert cache.get_cache(RPM_KEY) == 1 + + +@pytest.mark.asyncio +async def test_a_logged_success_adds_its_tokens_to_the_minute_once(): + cache = DualCache() + counter = _counter(cache) + + await counter.async_log_success_event(SUCCESS, None, None, None) + counter.log_success_event(SUCCESS, None, None, None) + + assert await cache.async_get_cache(TPM_KEY) == 80 + + +@pytest.mark.asyncio +async def test_a_logged_success_counts_nothing_when_no_callback_reads_the_counters(): + cache = DualCache() + counter = _counter(cache, charges=False) + + await counter.async_log_success_event(SUCCESS, None, None, None) + counter.log_success_event(SUCCESS, None, None, None) + + assert await cache.async_get_cache(TPM_KEY) is None + + +@pytest.mark.asyncio +async def test_a_batch_retrieve_is_not_usage(): + cache = DualCache() + + await _counter(cache).async_log_success_event({**SUCCESS, "call_type": "aretrieve_batch"}, None, None, None) + + assert await cache.async_get_cache(TPM_KEY) is None + + +def test_only_the_rate_limit_check_and_usage_based_v2_read_the_counters(): + assert reads_deployment_usage((object(), ModelRateLimitingCheck(dual_cache=DualCache()))) + assert not reads_deployment_usage((object(), DeploymentUsageCounter(DualCache()))) From a9d6ea94a23810e3f777ecb4a80f9dd6a47e319c Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 02:05:52 +0000 Subject: [PATCH 03/12] test(integration): a deployment with rpm 10 serves ten requests a minute next to a usage-based-routing-v2 group Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...nforce_model_rate_limits_routing_groups.py | 103 ++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 tests/integration/routing/test_enforce_model_rate_limits_routing_groups.py diff --git a/tests/integration/routing/test_enforce_model_rate_limits_routing_groups.py b/tests/integration/routing/test_enforce_model_rate_limits_routing_groups.py new file mode 100644 index 00000000000..0f8e7aba680 --- /dev/null +++ b/tests/integration/routing/test_enforce_model_rate_limits_routing_groups.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import gateway_from_environment +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import owned_config +from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 180) + +RPM: Final = 10 +EXTRA_REQUESTS: Final = 2 +LATENCY_GROUP: Final = "lat-model" +USAGE_GROUP: Final = "ubr-model" + + +@dataclass(frozen=True, slots=True) +class Rig: + owned: OwnedProxy + latency: Wire + usage: Wire + + +def _reply(request: Request) -> Reply: + if is_model_info_probe(request): + return model_list_reply() + return chat_reply("chatcmpl-rpm", OPENAI_MODEL, "served", stream=False) + + +def _deployment(model_name: str, model_id: str, wire: Wire) -> dict[str, JsonValue]: + return { + "model_name": model_name, + "litellm_params": { + "model": OPENAI_MODEL, + "api_base": wire.url + "/v1", + "api_key": "synthetic-openai-key", + "rpm": RPM, + }, + "model_info": {"id": model_id}, + } + + +def _config(directory: Path, latency: Wire, usage: Wire) -> Path: + """``enforce_model_rate_limits`` next to a usage-based-routing-v2 routing group, as the customer ran it.""" + return owned_config( + directory, + [_deployment(LATENCY_GROUP, "lat-a", latency), _deployment(USAGE_GROUP, "ubr-a", usage)], + litellm_settings={"cache": False}, + router_settings={ + "optional_pre_call_checks": ["enforce_model_rate_limits"], + "num_retries": 0, + "routing_groups": [ + {"group_name": "latency-group", "models": [LATENCY_GROUP], "routing_strategy": "latency-based-routing"}, + {"group_name": "ubr-group", "models": [USAGE_GROUP], "routing_strategy": "usage-based-routing-v2"}, + ], + }, + ) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("rpm-routing-groups") + with ( + gateway_from_environment() as gateway, + wire_server(_reply) as latency, + wire_server(_reply) as usage, + owned_proxy_process(gateway, directory, {}, config=_config(directory, latency, usage)) as owned, + ): + yield Rig(owned, latency, usage) + + +def _chat(rig: Rig, model: str) -> httpx.Response: + return rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"rpm {uuid.uuid4().hex}"}]}, + ) + + +def _forwarded(wire: Wire) -> int: + return sum(1 for request in wire.drain() if not is_model_info_probe(request)) + + +@pytest.mark.parametrize( + ("model", "wire"), [(LATENCY_GROUP, "latency"), (USAGE_GROUP, "usage")], ids=["other-group", "v2-group"] +) +def test_a_deployment_with_rpm_10_serves_ten_requests_a_minute(rig: Rig, model: str, wire: str) -> None: + responses: Final = tuple(_chat(rig, model) for _ in range(RPM + EXTRA_REQUESTS)) + + assert [response.status_code for response in responses] == [200] * RPM + [429] * EXTRA_REQUESTS, [ + response.text for response in responses + ] + assert _forwarded(getattr(rig, wire)) == RPM From b25a0a3335d3af6885370a6ecac1279f1cdccc16 Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 02:10:51 +0000 Subject: [PATCH 04/12] chore(router): type the usage counter hooks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/lowest_tpm_rpm_v2.py | 2 +- .../router_utils/deployment_usage_counter.py | 30 +++++++++++-------- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index ce92506bdd2..21cc28ca626 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -227,7 +227,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): num_retries=deployment.get("num_retries"), ) else: - result = await self.router_cache.async_get_cache(key=rpm_key) + result: Final = await self.router_cache.async_get_cache(key=rpm_key) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", diff --git a/litellm/router_utils/deployment_usage_counter.py b/litellm/router_utils/deployment_usage_counter.py index b550e971dab..50728526c14 100644 --- a/litellm/router_utils/deployment_usage_counter.py +++ b/litellm/router_utils/deployment_usage_counter.py @@ -5,9 +5,11 @@ The router charges one request before the routing pre-call checks run, hands it refuses the attempt, and charges the tokens of each logged success. Nothing else writes these keys. """ +from __future__ import annotations + from collections.abc import Callable, Iterable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger @@ -23,11 +25,7 @@ from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_utc_datetime if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span - - Span = _Span -else: - Span = Any + from opentelemetry.trace import Span USAGE_TTL: Final = 60 @@ -71,7 +69,9 @@ class DeploymentUsageCounter(CustomLogger): return usage_counter_keys(deployment, self._minute())[1] @with_service_target(ROUTER_USAGE_TARGET) - async def charge_request(self, deployment: Mapping[str, object], parent_otel_span: Span | None = None) -> int | None: + async def charge_request( + self, deployment: Mapping[str, object], parent_otel_span: Span | None = None + ) -> int | None: """One more request this minute; returns the count, or None when the cache could not say.""" return await self._async_increment(self.rpm_key(deployment), 1, parent_otel_span) @@ -129,16 +129,22 @@ class DeploymentUsageCounter(CustomLogger): return f"{model_id}:{model}", int(tokens) @with_service_target(ROUTER_USAGE_TARGET) - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): # pyright: ignore[reportMissingParameterType] # CustomLogger's untyped hook - counted: Final = self._logged_tokens(kwargs) # pyright: ignore[reportUnknownArgumentType] # see above + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + counted: Final = self._logged_tokens(kwargs) if counted is None: return prefix, tokens = counted - _ = await self._async_increment(f"{prefix}:tpm:{self._minute()}", tokens, get_parent_otel_span_from_kwargs(kwargs)) # pyright: ignore[reportUnknownArgumentType] # see above + _ = await self._async_increment( + f"{prefix}:tpm:{self._minute()}", tokens, get_parent_otel_span_from_kwargs(kwargs) + ) @with_service_target(ROUTER_USAGE_TARGET) - def log_success_event(self, kwargs, response_obj, start_time, end_time): # pyright: ignore[reportMissingParameterType] # CustomLogger's untyped hook - counted: Final = self._logged_tokens(kwargs) # pyright: ignore[reportUnknownArgumentType] # see above + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + counted: Final = self._logged_tokens(kwargs) if counted is None: return prefix, tokens = counted From f27b39db9fb401dd0ba5cae7923ae0f6f1425974 Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 02:20:56 +0000 Subject: [PATCH 05/12] chore(router): typed usage payload for the deployment usage counter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/lowest_tpm_rpm_v2.py | 4 +- .../router_utils/deployment_usage_counter.py | 43 +++++++++++++------ .../pre_call_checks/model_rate_limit_check.py | 2 +- 3 files changed, 32 insertions(+), 17 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 21cc28ca626..7d27083219c 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -153,7 +153,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ), ) else: - result: Final = self.router_cache.get_cache(key=rpm_key) + result: Final[int | None] = self.router_cache.get_cache(key=rpm_key) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", @@ -227,7 +227,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): num_retries=deployment.get("num_retries"), ) else: - result: Final = await self.router_cache.async_get_cache(key=rpm_key) + result: Final[int | None] = await self.router_cache.async_get_cache(key=rpm_key) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", diff --git a/litellm/router_utils/deployment_usage_counter.py b/litellm/router_utils/deployment_usage_counter.py index 50728526c14..8726ff8f44f 100644 --- a/litellm/router_utils/deployment_usage_counter.py +++ b/litellm/router_utils/deployment_usage_counter.py @@ -11,6 +11,8 @@ from collections.abc import Callable, Iterable, Mapping from datetime import datetime from typing import TYPE_CHECKING, Final +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache @@ -21,13 +23,27 @@ from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET from litellm.types.caching import RedisPipelineIncrementOperation -from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_utc_datetime if TYPE_CHECKING: from opentelemetry.trace import Span USAGE_TTL: Final = 60 +_STR_MAP: Final = TypeAdapter(Mapping[str, object]) + + +class _HiddenParams(BaseModel): + model_config = ConfigDict(frozen=True) + litellm_model_name: str | None = None + + +class _LoggedUsage(BaseModel): + """The slice of a ``standard_logging_object`` that says which deployment used how many tokens.""" + + model_config = ConfigDict(frozen=True, protected_namespaces=()) + model_id: str | None = None + total_tokens: int = 0 + hidden_params: _HiddenParams = _HiddenParams() def reads_deployment_usage(callbacks: Iterable[object]) -> bool: @@ -38,10 +54,8 @@ def reads_deployment_usage(callbacks: Iterable[object]) -> bool: def usage_counter_keys(deployment: Mapping[str, object], minute: str) -> tuple[str, str]: """The ``(tpm_key, rpm_key)`` a deployment's usage is counted under this minute.""" - model_info: Final = deployment.get("model_info") - params: Final = deployment.get("litellm_params") - model_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None - model: Final = params.get("model") if isinstance(params, Mapping) else None + model_id: Final = _STR_MAP.validate_python(deployment.get("model_info") or {}).get("id") + model: Final = _STR_MAP.validate_python(deployment.get("litellm_params") or {}).get("model") return f"{model_id}:{model}:tpm:{minute}", f"{model_id}:{model}:rpm:{minute}" @@ -57,7 +71,7 @@ class DeploymentUsageCounter(CustomLogger): charges: Callable[[], bool] = lambda: True, clock: Callable[[], datetime] = get_utc_datetime, ) -> None: - super().__init__() + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger.__init__ is untyped self.dual_cache: Final = dual_cache self.charges: Final = charges self.clock: Final = clock @@ -106,7 +120,9 @@ class DeploymentUsageCounter(CustomLogger): def _increment(self, key: str, value: int) -> int | None: try: - count: Final = int(self.dual_cache.increment_cache(key=key, value=value, ttl=USAGE_TTL)) + count: Final = int( + self.dual_cache.increment_cache(key=key, value=value, ttl=USAGE_TTL) # pyright: ignore[reportUnknownMemberType] # DualCache.increment_cache is untyped + ) except Exception as e: # noqa: BLE001 # see _async_increment verbose_router_logger.debug("DeploymentUsageCounter: usage was not recorded. Got - %s", e) return None @@ -118,15 +134,14 @@ class DeploymentUsageCounter(CustomLogger): """The ``(tpm_key prefix, total_tokens)`` of a logged success, or None when it counts nothing.""" if is_batch_retrieve_call_type(kwargs.get("call_type")) or not self.charges(): return None - logged: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") # pyright: ignore[reportAssignmentType] # the logging payload as CustomLogger receives it - if logged is None: + try: + logged: Final = _LoggedUsage.model_validate(kwargs.get("standard_logging_object")) + except ValidationError: return None - model_id: Final = logged.get("model_id") - model: Final = logged.get("hidden_params", {}).get("litellm_model_name") - tokens: Final = logged.get("total_tokens") - if model_id is None or not model or not tokens: + model: Final = logged.hidden_params.litellm_model_name + if logged.model_id is None or not model or not logged.total_tokens: return None - return f"{model_id}:{model}", int(tokens) + return f"{logged.model_id}:{model}", logged.total_tokens @with_service_target(ROUTER_USAGE_TARGET) async def async_log_success_event( 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 fcbab1c1a36..34e5bc95025 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 @@ -346,7 +346,7 @@ class ModelRateLimitingCheck(CustomLogger): ) if rpm_limit is not None: - current_rpm: Final = self.dual_cache.get_cache(key=rpm_key) + current_rpm: Final[int | None] = self.dual_cache.get_cache(key=rpm_key) if current_rpm is not None and current_rpm > rpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", From 735a468d49fc75c36b1b17da58841d6e81ed4be4 Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 02:57:17 +0000 Subject: [PATCH 06/12] fix(router): the usage counter is router state, tokens ride the usage write, checks read the attempt's own charge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_rust/core.py | 5 +- litellm/router_rust/steps/callbacks.py | 9 +- litellm/router_rust/steps/selection.py | 24 +-- litellm/router_rust/surface/legacy_members.py | 14 +- litellm/router_strategy/lowest_tpm_rpm_v2.py | 15 +- .../router_utils/deployment_usage_counter.py | 173 +++++++++--------- .../pre_call_checks/model_rate_limit_check.py | 17 +- .../test_request_redis_batch_post_call.py | 31 +--- .../router_strategy/test_lowest_tpm_rpm_v2.py | 39 ++-- .../test_router_routing_groups.py | 11 +- .../test_deployment_usage_counter.py | 103 +++++++---- .../test_enforce_model_rate_limits.py | 29 ++- 12 files changed, 273 insertions(+), 197 deletions(-) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 647e183ffac..1a3b08b7e1a 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -207,7 +207,7 @@ from litellm.router_rust.surface.settings import ( from litellm.router_strategy.least_busy import CALL_IDENTITY, call_id_of, owning_picks, pick_reserves, picks_reserve from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET from litellm.router_utils.cooldown_handlers import get_cooldown_deployment_info, mark_deployment_selection_failure -from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, usage_twins from litellm.router_utils.endpoint_eligibility import endpoint_for_call, selected_endpoint from litellm.router_utils.kubernetes_pod_discovery import ( KubernetesPodDiscovery, @@ -3533,6 +3533,7 @@ class RustCore( ) for key, amount in zip(keys, amounts) ] + operations.extend(usage_twins(operations, self._checks.counts_deployment_usage())) try: values: Final = await self._cache.async_increment_cache_pipeline(increment_list=operations) except Exception as e: # noqa: BLE001 # usage that cannot be written never fails the call, as on the Python router @@ -5050,7 +5051,7 @@ class _Driver: if deferred and response is None: return False model: Final = call.kwargs.get("model") - return self._core.deployment_tracks_usage( + return self._core._checks.counts_deployment_usage() or self._core.deployment_tracks_usage( # pyright: ignore[reportPrivateUsage] # the driver serves its own core (as_str_mapping(call.kwargs.get("model_info")) or {}).get("id"), model if isinstance(model, str) else call.group, ) diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index b5e597ce42b..44591a34e81 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -116,7 +116,7 @@ class OptionalChecks: self._cache: Final = cache self._hosted: Final = hosted self._ctx: Final = ctx - self.usage_counter: Final = DeploymentUsageCounter(cache, charges=self.reads_deployment_usage) + self.usage_counter: Final = DeploymentUsageCounter(cache) self.optional_callbacks: list[object] | None = None # mutable-ok: the router's live list, as Python keeps it self.router_budget_logger: RouterBudgetLimiting | None = None self.provider_budget_config: object = arguments.get("provider_budget_config") @@ -146,9 +146,7 @@ class OptionalChecks: def build(self, checks: object, model_list: object) -> None: """The constructor's checks: a budget appends ``router_budget_limiting`` to the list given (in place, as Python does), the list is added, a per-group affinity config adds the affinity - check, and ``apply_default_settings`` adds the empty default list. The usage counter the - checks and selectors read is owned first.""" - self._own(self.usage_counter) + check, and ``apply_default_settings`` adds the empty default list.""" listed: list[object] | None = checks if isinstance(checks, list) else None # pyright: ignore[reportUnknownVariableType] # mutable-ok: appended in place, as Python does # rebind-ok: the argument as passed if RouterBudgetLimiting.should_init_router_budget_limiter( # pyright: ignore[reportUnknownMemberType] # untyped model_list model_list=model_list if isinstance(model_list, list) else None, # pyright: ignore[reportUnknownArgumentType] # the argument as passed @@ -216,7 +214,8 @@ class OptionalChecks: session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, # pyright: ignore[reportArgumentType] # same ) - def reads_deployment_usage(self) -> bool: + def counts_deployment_usage(self) -> bool: + """Whether a registered routing callback reads the per-deployment usage counters.""" return reads_deployment_usage(self.routing_callbacks()) def _own(self, callback: CustomLogger) -> None: diff --git a/litellm/router_rust/steps/selection.py b/litellm/router_rust/steps/selection.py index d634d764b08..a141f99d1ad 100644 --- a/litellm/router_rust/steps/selection.py +++ b/litellm/router_rust/steps/selection.py @@ -30,7 +30,7 @@ from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 -from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, reads_deployment_usage +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, end_attempt, reads_deployment_usage from litellm.types.router import RoutingGroup, RoutingStrategy if TYPE_CHECKING: @@ -356,16 +356,17 @@ async def routing_pre_call_checks( (``_routing_callbacks``) checks it before the call. A check's error hands the request back and fails the attempt; the core records that failure like any other.""" checks: Final = tuple(callback for callback in callbacks if isinstance(callback, CustomLogger)) - counted: Final = reads_deployment_usage(checks) - if counted: - _ = await counter.charge_request(deployment) + charge: Final = await counter.charge_request(deployment) if reads_deployment_usage(checks) else None try: for callback in checks: await callback.async_pre_call_check(dict(deployment), None) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict except Exception: - if counted: - await counter.refund_request(deployment) + if charge is not None: + await counter.refund_request(charge) raise + finally: + if charge is not None: + end_attempt(charge) def routing_pre_call_checks_sync( @@ -373,16 +374,17 @@ def routing_pre_call_checks_sync( ) -> None: """``routing_strategy_pre_call_checks``: the sync twin, with the sync check of each ``CustomLogger``.""" checks: Final = tuple(callback for callback in callbacks if isinstance(callback, CustomLogger)) - counted: Final = reads_deployment_usage(checks) - if counted: - _ = counter.charge_request_sync(deployment) + charge: Final = counter.charge_request_sync(deployment) if reads_deployment_usage(checks) else None try: for callback in checks: callback.pre_call_check(dict(deployment)) # pyright: ignore[reportUnknownMemberType] # see above except Exception: - if counted: - counter.refund_request_sync(deployment) + if charge is not None: + counter.refund_request_sync(charge) raise + finally: + if charge is not None: + end_attempt(charge) def picked_id(picked: object) -> str | None: diff --git a/litellm/router_rust/surface/legacy_members.py b/litellm/router_rust/surface/legacy_members.py index f5773720b1c..c66c7903ab2 100644 --- a/litellm/router_rust/surface/legacy_members.py +++ b/litellm/router_rust/surface/legacy_members.py @@ -44,6 +44,7 @@ from litellm.router_strategy.auto_router.auto_router_model_naming import ( count_capability_routers, ) from litellm.router_utils.cooldown_handlers import set_cooldown_deployments +from litellm.router_utils.deployment_usage_counter import end_attempt, usage_twins from litellm.router_utils.endpoint_eligibility import endpoint_selection from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_io_token_refund_failure from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata @@ -347,13 +348,13 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): from litellm import exceptions as litellm_exceptions counter: Final = self._opt.usage_counter - _ = await counter.charge_request(deployment, parent_otel_span) + charge: Final = await counter.charge_request(deployment, parent_otel_span) for callback in self._opt.routing_callbacks(): if isinstance(callback, CustomLogger): try: await callback.async_pre_call_check(deployment, parent_otel_span) # pyright: ignore[reportUnknownMemberType] # see above except litellm_exceptions.RateLimitError as e: - await counter.refund_request(deployment, parent_otel_span) + await counter.refund_request(charge, parent_otel_span) _dispatch_failure(logging_obj, e) _ = set_cooldown_deployments( ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the core's RoutingContext @@ -364,9 +365,10 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): ) raise except Exception as e: - await counter.refund_request(deployment, parent_otel_span) + await counter.refund_request(charge, parent_otel_span) _dispatch_failure(logging_obj, e) raise + end_attempt(charge) async def async_callback_filter_deployments( self, @@ -478,7 +480,8 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): then holds what Redis answered.""" from litellm.types.caching import RedisPipelineIncrementOperation - if not self._registry_state.cached_deployment_tracks_usage(deployment_id, model_group): + counted: Final = self._opt.counts_deployment_usage() + if not counted and not self._registry_state.cached_deployment_tracks_usage(deployment_id, model_group): return current_minute: Final = get_utc_datetime().strftime("%H-%M") tpm_key: Final = RouterCacheEnum.TPM.value.format( @@ -487,10 +490,11 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): rpm_key: Final = RouterCacheEnum.RPM.value.format( id=deployment_id, current_minute=current_minute, model=deployment_name ) - operations: Final[Sequence[RedisPipelineIncrementOperation]] = [ + router_operations: Final[Sequence[RedisPipelineIncrementOperation]] = [ RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=_USAGE_TTL) for key, increment_value in ((tpm_key, total_tokens), (rpm_key, 1)) ] + operations: Final = [*router_operations, *usage_twins(router_operations, counted)] values: Final = await self._cache.async_increment_cache_pipeline( increment_list=list(operations), parent_otel_span=parent_otel_span ) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 7d27083219c..627592db691 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -6,7 +6,7 @@ from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass from functools import partial -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, ClassVar, Final import httpx @@ -18,6 +18,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models from litellm.litellm_core_utils.token_counter import embedding_aware_token_count +from litellm.router_utils.deployment_usage_counter import charged_rpm from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose @@ -85,6 +86,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): Increments tpm/rpm limit using redis.incr """ + reads_deployment_usage: ClassVar[bool] = True + test_flag: bool = False logged_success: int = 0 logged_failure: int = 0 @@ -153,7 +156,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ), ) else: - result: Final[int | None] = self.router_cache.get_cache(key=rpm_key) + admitted: Final = charged_rpm(rpm_key) + result: Final[int | None] = ( + admitted if admitted is not None else self.router_cache.get_cache(key=rpm_key) + ) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", @@ -227,7 +233,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): num_retries=deployment.get("num_retries"), ) else: - result: Final[int | None] = await self.router_cache.async_get_cache(key=rpm_key) + admitted: Final = charged_rpm(rpm_key) + result: Final[int | None] = ( + admitted if admitted is not None else await self.router_cache.async_get_cache(key=rpm_key) + ) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", diff --git a/litellm/router_utils/deployment_usage_counter.py b/litellm/router_utils/deployment_usage_counter.py index 8726ff8f44f..55865e068de 100644 --- a/litellm/router_utils/deployment_usage_counter.py +++ b/litellm/router_utils/deployment_usage_counter.py @@ -1,55 +1,65 @@ """The one writer of the per-deployment ``{id}:{model}:rpm|tpm:{HH-MM}`` usage counters that usage-based-routing-v2 selection and ``enforce_model_rate_limits`` read. -The router charges one request before the routing pre-call checks run, hands it back when a check -refuses the attempt, and charges the tokens of each logged success. Nothing else writes these keys. +The router charges one request before the routing pre-call checks run and hands it back when a check +refuses the attempt. The tokens of a served call ride the router's own usage write, as the twin of its +``global_router:`` tpm increment (``enforcement_tpm_operation``). Nothing else writes these keys. """ from __future__ import annotations from collections.abc import Callable, Iterable, Mapping +from contextvars import ContextVar +from dataclasses import dataclass from datetime import datetime from typing import TYPE_CHECKING, Final -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import TypeAdapter 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.litellm_core_utils.batch_utils import is_batch_retrieve_call_type -from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs -from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 -from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck -from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.router import RouterCacheEnum from litellm.utils import get_utc_datetime if TYPE_CHECKING: from opentelemetry.trace import Span USAGE_TTL: Final = 60 +ROUTER_USAGE_TARGET: Final = "router_usage" _STR_MAP: Final = TypeAdapter(Mapping[str, object]) +_ROUTER_USAGE_PREFIX: Final = RouterCacheEnum.TPM.value.split("{", 1)[0] -class _HiddenParams(BaseModel): - model_config = ConfigDict(frozen=True) - litellm_model_name: str | None = None +@dataclass(frozen=True, slots=True) +class UsageCharge: + """One request charged to a deployment's rpm counter: the key it went to, and what the counter + read right after, None when the cache could not say.""" + + key: str + count: int | None -class _LoggedUsage(BaseModel): - """The slice of a ``standard_logging_object`` that says which deployment used how many tokens.""" +_ADMITTED: Final[ContextVar[UsageCharge | None]] = ContextVar("deployment_usage_charge", default=None) - model_config = ConfigDict(frozen=True, protected_namespaces=()) - model_id: str | None = None - total_tokens: int = 0 - hidden_params: _HiddenParams = _HiddenParams() + +def charged_rpm(rpm_key: str) -> int | None: + """What this attempt's own charge counted on ``rpm_key``, so a check compares the atomic + increment it was admitted with rather than a cache read another attempt may have overtaken.""" + charge: Final = _ADMITTED.get() + return charge.count if charge is not None and charge.key == rpm_key else None + + +def end_attempt(charge: UsageCharge) -> None: + if _ADMITTED.get() is charge: + _ADMITTED.set(None) def reads_deployment_usage(callbacks: Iterable[object]) -> bool: - """Whether one of ``callbacks`` reads the per-deployment usage counters: usage-based-routing-v2 - selection or ``enforce_model_rate_limits``. Without a reader there is nothing to count.""" - return any(isinstance(callback, LowestTPMLoggingHandler_v2 | ModelRateLimitingCheck) for callback in callbacks) + """Whether one of ``callbacks`` reads the per-deployment usage counters (it says so with a + ``reads_deployment_usage`` attribute). Without a reader there is nothing to count.""" + return any(getattr(callback, "reads_deployment_usage", False) is True for callback in callbacks) def usage_counter_keys(deployment: Mapping[str, object], minute: str) -> tuple[str, str]: @@ -59,21 +69,32 @@ def usage_counter_keys(deployment: Mapping[str, object], minute: str) -> tuple[s return f"{model_id}:{model}:tpm:{minute}", f"{model_id}:{model}:rpm:{minute}" -class DeploymentUsageCounter(CustomLogger): - """Counts requests and tokens against a deployment in the router's dual cache, through one Redis - pipeline per write so every worker then holds the shared total. Requests are charged by the - pre-call checks that run a reader; ``charges`` says whether a registered callback reads the - counters, so a logged success counts its tokens only then.""" +def enforcement_tpm_operation(operation: RedisPipelineIncrementOperation) -> RedisPipelineIncrementOperation | None: + """The ``{id}:{model}:tpm:{HH-MM}`` increment that mirrors a ``global_router:`` tpm increment of the + router's usage write, None for any other key.""" + key: Final = operation["key"] + if not key.startswith(_ROUTER_USAGE_PREFIX) or ":tpm:" not in key: + return None + return RedisPipelineIncrementOperation( + key=key.removeprefix(_ROUTER_USAGE_PREFIX), increment_value=operation["increment_value"], ttl=USAGE_TTL + ) - def __init__( - self, - dual_cache: DualCache, - charges: Callable[[], bool] = lambda: True, - clock: Callable[[], datetime] = get_utc_datetime, - ) -> None: - super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger.__init__ is untyped + +def usage_twins( + operations: Iterable[RedisPipelineIncrementOperation], counted: bool +) -> tuple[RedisPipelineIncrementOperation, ...]: + """The enforcement increments that ride the router's usage write, none when nothing reads them.""" + if not counted: + return () + return tuple(twin for twin in map(enforcement_tpm_operation, operations) if twin is not None) + + +class DeploymentUsageCounter: + """Counts requests against a deployment in the router's dual cache, through one Redis pipeline per + write so every worker then holds the shared total.""" + + def __init__(self, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime) -> None: self.dual_cache: Final = dual_cache - self.charges: Final = charges self.clock: Final = clock def _minute(self) -> str: @@ -85,22 +106,35 @@ class DeploymentUsageCounter(CustomLogger): @with_service_target(ROUTER_USAGE_TARGET) async def charge_request( self, deployment: Mapping[str, object], parent_otel_span: Span | None = None - ) -> int | None: - """One more request this minute; returns the count, or None when the cache could not say.""" - return await self._async_increment(self.rpm_key(deployment), 1, parent_otel_span) + ) -> UsageCharge: + """One more request this minute, remembered for the checks of this attempt (``charged_rpm``).""" + key: Final = self.rpm_key(deployment) + charge: Final = UsageCharge(key, await self._async_increment(key, 1, parent_otel_span)) + _ADMITTED.set(charge) + return charge @with_service_target(ROUTER_USAGE_TARGET) - async def refund_request(self, deployment: Mapping[str, object], parent_otel_span: Span | None = None) -> None: - """A refused attempt does not count.""" - _ = await self._async_increment(self.rpm_key(deployment), -1, parent_otel_span) + async def refund_request(self, charge: UsageCharge, parent_otel_span: Span | None = None) -> None: + """A refused attempt does not count, on the minute it was charged to.""" + _ = await self._async_increment(charge.key, -1, parent_otel_span) + end_attempt(charge) @with_service_target(ROUTER_USAGE_TARGET) - def charge_request_sync(self, deployment: Mapping[str, object]) -> int | None: - return self._increment(self.rpm_key(deployment), 1) + def charge_request_sync(self, deployment: Mapping[str, object]) -> UsageCharge: + key: Final = self.rpm_key(deployment) + charge: Final = UsageCharge(key, self._increment(key, 1)) + _ADMITTED.set(charge) + return charge @with_service_target(ROUTER_USAGE_TARGET) - def refund_request_sync(self, deployment: Mapping[str, object]) -> None: - _ = self._increment(self.rpm_key(deployment), -1) + def refund_request_sync(self, charge: UsageCharge) -> None: + _ = self._increment(charge.key, -1) + end_attempt(charge) + + def _held(self, key: str, value: int, count: int) -> int: + """What this worker's memory keeps: a charge never lowers what a later charge already wrote.""" + held: Final = self.dual_cache.in_memory_cache.get_cache(key) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # InMemoryCache reads are untyped + return max(count, int(held)) if value > 0 and isinstance(held, int | float) else count async def _async_increment(self, key: str, value: int, parent_otel_span: Span | None) -> int | None: operation: Final = RedisPipelineIncrementOperation(key=key, increment_value=value, ttl=USAGE_TTL) @@ -115,7 +149,7 @@ class DeploymentUsageCounter(CustomLogger): return None count: Final = int(values[0]) if self.dual_cache.redis_cache is not None: - await self.dual_cache.async_set_cache(key, count, local_only=True, ttl=USAGE_TTL) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + await self.dual_cache.async_set_cache(key, self._held(key, value, count), local_only=True, ttl=USAGE_TTL) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped return count def _increment(self, key: str, value: int) -> int | None: @@ -127,43 +161,18 @@ class DeploymentUsageCounter(CustomLogger): verbose_router_logger.debug("DeploymentUsageCounter: usage was not recorded. Got - %s", e) return None if self.dual_cache.redis_cache is not None: - self.dual_cache.set_cache(key, count, local_only=True, ttl=USAGE_TTL) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + self.dual_cache.set_cache(key, self._held(key, value, count), local_only=True, ttl=USAGE_TTL) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped return count - def _logged_tokens(self, kwargs: Mapping[str, object]) -> tuple[str, int] | None: - """The ``(tpm_key prefix, total_tokens)`` of a logged success, or None when it counts nothing.""" - if is_batch_retrieve_call_type(kwargs.get("call_type")) or not self.charges(): - return None - try: - logged: Final = _LoggedUsage.model_validate(kwargs.get("standard_logging_object")) - except ValidationError: - return None - model: Final = logged.hidden_params.litellm_model_name - if logged.model_id is None or not model or not logged.total_tokens: - return None - return f"{logged.model_id}:{model}", logged.total_tokens - @with_service_target(ROUTER_USAGE_TARGET) - async def async_log_success_event( - self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime - ) -> None: - counted: Final = self._logged_tokens(kwargs) - if counted is None: - return - prefix, tokens = counted - _ = await self._async_increment( - f"{prefix}:tpm:{self._minute()}", tokens, get_parent_otel_span_from_kwargs(kwargs) - ) - - @with_service_target(ROUTER_USAGE_TARGET) - def log_success_event( - self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime - ) -> None: - counted: Final = self._logged_tokens(kwargs) - if counted is None: - return - prefix, tokens = counted - _ = self._increment(f"{prefix}:tpm:{self._minute()}", tokens) - - -__all__ = ["USAGE_TTL", "DeploymentUsageCounter", "reads_deployment_usage", "usage_counter_keys"] +__all__ = [ + "USAGE_TTL", + "DeploymentUsageCounter", + "UsageCharge", + "charged_rpm", + "end_attempt", + "enforcement_tpm_operation", + "reads_deployment_usage", + "usage_counter_keys", + "usage_twins", +] 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 34e5bc95025..a18b61ae17d 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 @@ -14,7 +14,7 @@ is logged the first time such a deployment is seen. import contextlib from collections.abc import Callable, Mapping from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal import httpx from pydantic import BaseModel, ConfigDict @@ -25,6 +25,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.deployment_usage_counter import charged_rpm from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( async_io_token_pre_call_check, async_io_token_reconcile_success, @@ -134,6 +135,8 @@ class ModelRateLimitingCheck(CustomLogger): this check actively enforces those limits across ALL routing strategies. """ + reads_deployment_usage: ClassVar[bool] = True + def __init__(self, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime) -> None: self.dual_cache = dual_cache self.clock = clock @@ -346,7 +349,10 @@ class ModelRateLimitingCheck(CustomLogger): ) if rpm_limit is not None: - current_rpm: Final[int | None] = self.dual_cache.get_cache(key=rpm_key) + admitted: Final = charged_rpm(rpm_key) + current_rpm: Final[int | None] = ( + admitted if admitted is not None else self.dual_cache.get_cache(key=rpm_key) + ) if current_rpm is not None and current_rpm > rpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", @@ -432,7 +438,12 @@ class ModelRateLimitingCheck(CustomLogger): raise daily.limit_error("tpd", tpd_limit, current_tpd, dt) if rpm_limit is not None: - current_rpm: Final = await self._async_get_current_tpm(rpm_key, rpm_limit, parent_otel_span) + admitted: Final = charged_rpm(rpm_key) + current_rpm: Final = ( + admitted + if admitted is not None + else await self._async_get_current_tpm(rpm_key, rpm_limit, parent_otel_span) + ) if current_rpm is not None and current_rpm > rpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 4b73bb6dd34..a601b68102f 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -9,7 +9,7 @@ import asyncio import datetime import hashlib import json -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Sequence import pytest @@ -34,7 +34,6 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( ) from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement from litellm.proxy.utils import InternalUsageCache -from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ModelResponse @@ -118,33 +117,13 @@ def _response_cache(redis_cache: FakeRedisCache) -> Cache: return cache -def _tpm_router(redis_cache: FakeRedisCache) -> tuple[LowestTPMLoggingHandler_v2, DualCache]: - router_cache = DualCache() - router_cache.attach_redis_cache(redis_cache) - return LowestTPMLoggingHandler_v2(router_cache=router_cache, routing_args={"ttl": 60}), router_cache - - -def _tpm_kwargs() -> Mapping[str, object]: - return { - "standard_logging_object": { - "model_group": "gpt", - "model_id": "dep-a", - "hidden_params": {"litellm_model_name": "openai/gpt-4o-mini"}, - "total_tokens": 42, - }, - "litellm_params": {"metadata": {}}, - } - - @pytest.mark.asyncio async def test_every_post_call_owner_rides_one_pipeline_that_goes_out_when_the_callbacks_are_done(): client = FakeClient(_ok_replies) redis_cache = PostCallFakeRedisCache(client) limiter = _limiter(redis_cache) - tpm, router_cache = _tpm_router(redis_cache) with request_redis_batch_scope(): - await tpm.async_log_success_event(_tpm_kwargs(), None, None, None) await limiter.async_increment_tokens_with_ttl_preservation(_token_ops("{api_key:k1}:tokens")) await limiter._release_stashed_parallel_slot( _slot_stash("slot-1", "{api_key:k1}:parallel"), None, in_logging_callback=True @@ -153,16 +132,10 @@ async def test_every_post_call_owner_rides_one_pipeline_that_goes_out_when_the_c await flush_post_call_redis_batches() assert len(client.pipelines) == 1 - assert _names(client) == ["INCRBYFLOAT", "EXPIRE", "EVALSHA", "EVALSHA"] + assert _names(client) == ["EVALSHA", "EVALSHA"] evalshas = [c for c in client.pipelines[0].commands if c[0] == "EVALSHA"] assert [c[1] for c in evalshas] == [sha_of(TOKEN_INCREMENT_SCRIPT), sha_of(PARALLEL_RELEASE_SCRIPT)] assert redis_cache.alone == [] - assert ( - await router_cache.in_memory_cache.async_get_cache( - next(k for k in router_cache.in_memory_cache.cache_dict if ":tpm:" in k) - ) - == 42 - ) @pytest.mark.asyncio diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py index bda43a87fbd..e1aa58f7f9b 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py @@ -9,6 +9,7 @@ import os import time import traceback from collections.abc import Iterator +from datetime import datetime from typing import Final from unittest.mock import patch @@ -24,8 +25,9 @@ from litellm.router_rust.steps.selection import routing_pre_call_checks_sync from litellm.router_strategy.lowest_tpm_rpm_v2 import ( LowestTPMLoggingHandler_v2 as LowestTPMLoggingHandler, ) -from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter -from litellm.types.router import DeploymentTypedDict, RouterRateLimitError +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, enforcement_tpm_operation +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.router import DeploymentTypedDict, RouterCacheEnum, RouterRateLimitError from litellm.types.utils import StandardLoggingPayload from litellm.utils import _invalidate_model_cost_lowercase_map, get_utc_datetime from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -46,13 +48,24 @@ def restore_verbose_logger_level() -> Iterator[None]: @pytest.mark.usefixtures("restore_verbose_logger_level") -def _router_usage_counter(router: Router) -> DeploymentUsageCounter: - return next(callback for callback in router.optional_callbacks or () if isinstance(callback, DeploymentUsageCounter)) + +def _count_tokens(cache: DualCache, model_id: str, model: str, tokens: int, now: datetime) -> None: + """A served call's tokens as the router's usage write counts them: the enforcement twin of its + ``global_router:`` tpm increment.""" + minute: Final = now.strftime("%H-%M") + twin: Final = enforcement_tpm_operation( + RedisPipelineIncrementOperation( + key=RouterCacheEnum.TPM.value.format(id=model_id, model=model, current_minute=minute), increment_value=tokens, ttl=60 + ) + ) + assert twin is not None + cache.increment_cache(key=twin["key"], value=twin["increment_value"], ttl=twin["ttl"]) -def _log_success(handler: LowestTPMLoggingHandler, counter: DeploymentUsageCounter, **kwargs) -> None: - """A logged success as the router fans it out: the counter's tokens, then the strategy's hook.""" - counter.log_success_event(**kwargs) +def _log_success(handler: LowestTPMLoggingHandler, cache: DualCache, **kwargs) -> None: + """A served call as the router counts it: its tokens on the deployment's minute, then the strategy's hook.""" + logged = kwargs["kwargs"]["standard_logging_object"] + _count_tokens(cache, logged["model_id"], logged["hidden_params"]["litellm_model_name"], logged["total_tokens"], get_utc_datetime()) handler.log_success_event(**kwargs) @@ -94,7 +107,7 @@ def test_tpm_rpm_updated(): routing_pre_call_checks_sync((lowest_tpm_logger,), litellm_deployment_dict, DeploymentUsageCounter(test_cache)) _log_success( lowest_tpm_logger, - DeploymentUsageCounter(test_cache), + test_cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -155,7 +168,7 @@ def test_get_available_deployments(): end_time = time.time() _log_success( lowest_tpm_logger, - DeploymentUsageCounter(test_cache), + test_cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -184,7 +197,7 @@ def test_get_available_deployments(): end_time = time.time() _log_success( lowest_tpm_logger, - DeploymentUsageCounter(test_cache), + test_cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -264,7 +277,7 @@ def test_router_get_available_deployments(): end_time = time.time() _log_success( router.lowesttpm_logger_v2, - _router_usage_counter(router), + router.cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -290,7 +303,7 @@ def test_router_get_available_deployments(): end_time = time.time() _log_success( router.lowesttpm_logger_v2, - _router_usage_counter(router), + router.cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -354,7 +367,7 @@ def test_router_skip_rate_limited_deployments(): end_time = time.time() _log_success( router.lowesttpm_logger_v2, - _router_usage_counter(router), + router.cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 1cbbe78573d..50944be0f88 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -20,9 +20,9 @@ from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.router_api import RoutingContext from litellm.router_strategy.least_busy import picks_reserve -from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter +from litellm.router_utils.deployment_usage_counter import usage_counter_keys from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup -from litellm.utils import Rules, function_setup +from litellm.utils import Rules, function_setup, get_utc_datetime def _model_list(): @@ -1177,7 +1177,6 @@ def _ticket_router() -> Router: @pytest.mark.parametrize("model", ["lat-model", "ubr-model"]) async def test_a_deployment_with_rpm_10_serves_ten_requests_a_minute_next_to_a_usage_based_v2_group(model: str): router = _ticket_router() - counter = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentUsageCounter)) deployment = router.get_deployment(model_id=f"{model[:3]}-a") assert deployment is not None @@ -1186,7 +1185,8 @@ async def test_a_deployment_with_rpm_10_serves_ten_requests_a_minute_next_to_a_u ] assert outcomes == [200] * 10 + [429] * 2, outcomes - assert await router.cache.async_get_cache(counter.rpm_key(deployment.model_dump())) == 10 + _, rpm_key = usage_counter_keys(deployment.model_dump(), get_utc_datetime().strftime("%H-%M")) + assert await router.cache.async_get_cache(rpm_key) == 10 async def _outcome(router: Router, model: str) -> int: @@ -1200,12 +1200,11 @@ async def _outcome(router: Router, model: str) -> int: @pytest.mark.asyncio async def test_a_success_adds_its_tokens_to_the_deployment_once_next_to_a_usage_based_v2_group(): router = _ticket_router() - counter = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentUsageCounter)) deployment = router.get_deployment(model_id="lat-a") assert deployment is not None response = await router.acompletion(model="lat-model", messages=[{"role": "user", "content": "hi"}]) await asyncio.sleep(0.2) - tpm_key: Final = counter.rpm_key(deployment.model_dump()).replace(":rpm:", ":tpm:") + tpm_key, _ = usage_counter_keys(deployment.model_dump(), get_utc_datetime().strftime("%H-%M")) assert await router.cache.async_get_cache(tpm_key) == response.usage.total_tokens diff --git a/tests/unit/router_utils/test_deployment_usage_counter.py b/tests/unit/router_utils/test_deployment_usage_counter.py index d906908a49e..4f8cb3336cc 100644 --- a/tests/unit/router_utils/test_deployment_usage_counter.py +++ b/tests/unit/router_utils/test_deployment_usage_counter.py @@ -1,26 +1,35 @@ -from datetime import datetime, timezone +import asyncio +from datetime import datetime, timedelta, timezone import pytest from litellm.caching.dual_cache import DualCache -from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, reads_deployment_usage +from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.router_utils.deployment_usage_counter import ( + DeploymentUsageCounter, + charged_rpm, + end_attempt, + enforcement_tpm_operation, + reads_deployment_usage, + usage_twins, +) from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck NOON = datetime(2026, 10, 28, 12, 0, 30, tzinfo=timezone.utc) DEPLOYMENT = {"litellm_params": {"model": "openai/gpt-4o"}, "model_info": {"id": "deploy-1"}} RPM_KEY = "deploy-1:openai/gpt-4o:rpm:12-00" -TPM_KEY = "deploy-1:openai/gpt-4o:tpm:12-00" -SUCCESS = { - "standard_logging_object": { - "model_id": "deploy-1", - "total_tokens": 40, - "hidden_params": {"litellm_model_name": "openai/gpt-4o"}, - } -} -def _counter(cache: DualCache, charges: bool = True) -> DeploymentUsageCounter: - return DeploymentUsageCounter(cache, charges=lambda: charges, clock=lambda: NOON) +class _Clock: + def __init__(self) -> None: + self.now = NOON + + def __call__(self) -> datetime: + return self.now + + +def _counter(cache: DualCache, clock: _Clock | None = None) -> DeploymentUsageCounter: + return DeploymentUsageCounter(cache, clock=clock or _Clock()) @pytest.mark.asyncio @@ -28,9 +37,10 @@ async def test_each_charged_request_adds_one_to_the_minute(): cache = DualCache() counter = _counter(cache) - counts = [await counter.charge_request(DEPLOYMENT) for _ in range(3)] + charges = [await counter.charge_request(DEPLOYMENT) for _ in range(3)] - assert counts == [1, 2, 3] + assert [charge.count for charge in charges] == [1, 2, 3] + assert {charge.key for charge in charges} == {RPM_KEY} assert await cache.async_get_cache(RPM_KEY) == 3 @@ -39,54 +49,77 @@ async def test_a_refunded_request_leaves_the_minute_as_it_was(): cache = DualCache() counter = _counter(cache) await counter.charge_request(DEPLOYMENT) - await counter.charge_request(DEPLOYMENT) + second = await counter.charge_request(DEPLOYMENT) - await counter.refund_request(DEPLOYMENT) + await counter.refund_request(second) assert await cache.async_get_cache(RPM_KEY) == 1 +@pytest.mark.asyncio +async def test_a_refund_after_the_minute_rolled_over_goes_back_to_the_charged_minute(): + cache = DualCache() + clock = _Clock() + counter = _counter(cache, clock) + charge = await counter.charge_request(DEPLOYMENT) + clock.now = NOON + timedelta(minutes=1) + + await counter.refund_request(charge) + + assert await cache.async_get_cache(RPM_KEY) == 0 + assert await cache.async_get_cache("deploy-1:openai/gpt-4o:rpm:12-01") is None + + def test_sync_charges_and_refunds_count_like_the_async_ones(): cache = DualCache() counter = _counter(cache) - assert [counter.charge_request_sync(DEPLOYMENT) for _ in range(2)] == [1, 2] - counter.refund_request_sync(DEPLOYMENT) + charges = [counter.charge_request_sync(DEPLOYMENT) for _ in range(2)] + counter.refund_request_sync(charges[1]) + assert [charge.count for charge in charges] == [1, 2] assert cache.get_cache(RPM_KEY) == 1 @pytest.mark.asyncio -async def test_a_logged_success_adds_its_tokens_to_the_minute_once(): +async def test_a_check_reads_the_count_its_own_charge_was_admitted_with(): cache = DualCache() counter = _counter(cache) - await counter.async_log_success_event(SUCCESS, None, None, None) - counter.log_success_event(SUCCESS, None, None, None) + async def attempt() -> int | None: + charge = await counter.charge_request(DEPLOYMENT) + await asyncio.sleep(0) + seen = charged_rpm(RPM_KEY) + end_attempt(charge) + return seen - assert await cache.async_get_cache(TPM_KEY) == 80 + assert sorted(await asyncio.gather(attempt(), attempt(), attempt())) == [1, 2, 3] + assert charged_rpm(RPM_KEY) is None @pytest.mark.asyncio -async def test_a_logged_success_counts_nothing_when_no_callback_reads_the_counters(): - cache = DualCache() - counter = _counter(cache, charges=False) +async def test_a_charge_answers_only_for_its_own_key(): + counter = _counter(DualCache()) - await counter.async_log_success_event(SUCCESS, None, None, None) - counter.log_success_event(SUCCESS, None, None, None) + charge = await counter.charge_request(DEPLOYMENT) - assert await cache.async_get_cache(TPM_KEY) is None + assert charged_rpm(RPM_KEY) == 1 + assert charged_rpm("deploy-2:openai/gpt-4o:rpm:12-00") is None + end_attempt(charge) -@pytest.mark.asyncio -async def test_a_batch_retrieve_is_not_usage(): - cache = DualCache() +def test_the_router_usage_tpm_write_has_an_enforcement_twin_and_the_rpm_write_has_none(): + tpm = RedisPipelineIncrementOperation(key="global_router:deploy-1:openai/gpt-4o:tpm:12-00", increment_value=40, ttl=60) + rpm = RedisPipelineIncrementOperation(key="global_router:deploy-1:openai/gpt-4o:rpm:12-00", increment_value=1, ttl=60) - await _counter(cache).async_log_success_event({**SUCCESS, "call_type": "aretrieve_batch"}, None, None, None) - - assert await cache.async_get_cache(TPM_KEY) is None + assert enforcement_tpm_operation(tpm) == RedisPipelineIncrementOperation( + key="deploy-1:openai/gpt-4o:tpm:12-00", increment_value=40, ttl=60 + ) + assert enforcement_tpm_operation(rpm) is None + assert usage_twins((tpm, rpm), counted=True) == (enforcement_tpm_operation(tpm),) + assert usage_twins((tpm, rpm), counted=False) == () -def test_only_the_rate_limit_check_and_usage_based_v2_read_the_counters(): +def test_only_a_callback_that_says_so_reads_the_counters(): assert reads_deployment_usage((object(), ModelRateLimitingCheck(dual_cache=DualCache()))) assert not reads_deployment_usage((object(), DeploymentUsageCounter(DualCache()))) 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 9a89d25ba4a..ac3c6b1ea58 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -18,10 +18,12 @@ from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.router_rust.steps.selection import routing_pre_call_checks, routing_pre_call_checks_sync -from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, enforcement_tpm_operation from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.router import RouterCacheEnum TPM_DEPLOYMENT = { "tpm": 1000, @@ -56,14 +58,35 @@ async def _async_pre_call(check: ModelRateLimitingCheck, deployment: dict[str, o return deployment + +def _count_tokens(cache: DualCache, model_id: str, model: str, tokens: int, now: datetime) -> None: + """A served call's tokens as the router's usage write counts them: the enforcement twin of its + ``global_router:`` tpm increment.""" + minute: Final = now.strftime("%H-%M") + twin: Final = enforcement_tpm_operation( + RedisPipelineIncrementOperation( + key=RouterCacheEnum.TPM.value.format(id=model_id, model=model, current_minute=minute), increment_value=tokens, ttl=60 + ) + ) + assert twin is not None + cache.increment_cache(key=twin["key"], value=twin["increment_value"], ttl=twin["ttl"]) + + +def _count_logged_tokens(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: + logged: Final = kwargs["standard_logging_object"] + assert isinstance(logged, dict) + _count_tokens( + check.dual_cache, logged["model_id"], logged["hidden_params"]["litellm_model_name"], logged["total_tokens"], check.clock() + ) + def _log_success(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: """A logged success as the router fans it out: the counter's tokens, then the check's daily tokens.""" - _counter(check).log_success_event(kwargs, None, None, None) + _count_logged_tokens(check, kwargs) check.log_success_event(kwargs, None, None, None) async def _async_log_success(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: - await _counter(check).async_log_success_event(kwargs, None, None, None) + _count_logged_tokens(check, kwargs) await check.async_log_success_event(kwargs, None, None, None) From 0fb04cfd60f5660db7c9b2ea52c0b3757991773f Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 03:09:47 +0000 Subject: [PATCH 07/12] test(router): two attempts charged before either reply lands admit exactly the rpm Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_enforce_model_rate_limits.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) 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 ac3c6b1ea58..fc10db49b2e 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -24,6 +24,7 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ) from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.router import RouterCacheEnum +from tests.unit.caching.test_redis_batch import FakeClient, FakeRedisCache TPM_DEPLOYMENT = { "tpm": 1000, @@ -913,3 +914,54 @@ class TestDailyResetTime: router.cache.in_memory_cache.get_cache("daily-test-id:openai/gpt-4.1-mini:tpd:2026-10-27") # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # InMemoryCache.get_cache is untyped ) assert day_before is None + + +class _RepliesAfterBothCharged(FakeRedisCache): + """Redis that answers the first increment only once the second was sent, and the second only once + the first attempt has checked its limit: the first attempt's local write-back lands on a count the + second attempt already raised.""" + + def __init__(self) -> None: + super().__init__(FakeClient(lambda command: None)) + self.total = 0 + self.sent = 0 + self.both_sent = asyncio.Event() + self.first_checked = asyncio.Event() + + async def async_increment_pipeline( # pyright: ignore[reportIncompatibleMethodOverride] # fake, no server + self, increment_list: list[RedisPipelineIncrementOperation], **kwargs: object + ) -> list[float]: + self.total += int(increment_list[0]["increment_value"]) + answered: Final = float(self.total) + self.sent += 1 + order: Final = self.sent + if order == 1: + await self.both_sent.wait() + if order == 2: + self.both_sent.set() + await self.first_checked.wait() + return [answered] + + +@pytest.mark.asyncio +async def test_two_attempts_charged_before_either_reply_lands_admit_exactly_the_rpm(): + redis_cache = _RepliesAfterBothCharged() + dual_cache = DualCache() + dual_cache.attach_redis_cache(redis_cache) + check = ModelRateLimitingCheck(dual_cache=dual_cache, clock=Clock(DAY_ONE)) + deployment = {"rpm": 1, "litellm_params": {"model": "gpt-4"}, "model_info": {"id": "race-id"}} + + async def attempt() -> int: + try: + await _async_pre_call(check, deployment) + except litellm.RateLimitError: + return 429 + return 200 + + attempts = {asyncio.ensure_future(attempt()) for _ in range(2)} + checked_first, still_charged = await asyncio.wait(attempts, return_when=asyncio.FIRST_COMPLETED) + redis_cache.first_checked.set() + + assert [task.result() for task in checked_first] == [200] + assert await asyncio.gather(*still_charged) == [429] + assert redis_cache.total == 1 From b2b5a84d8c08aa039d4c1fc838c01d5e82d0b1b8 Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 03:38:14 +0000 Subject: [PATCH 08/12] fix(router): usage keys from pydantic model info, checks get the deployment dict itself, cache hits count no tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_rust/core.py | 5 +++- litellm/router_rust/steps/selection.py | 6 +++-- .../router_utils/deployment_usage_counter.py | 26 +++++++++++++++---- 3 files changed, 29 insertions(+), 8 deletions(-) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 1a3b08b7e1a..cfddcd5989b 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -207,7 +207,7 @@ from litellm.router_rust.surface.settings import ( from litellm.router_strategy.least_busy import CALL_IDENTITY, call_id_of, owning_picks, pick_reserves, picks_reserve from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET from litellm.router_utils.cooldown_handlers import get_cooldown_deployment_info, mark_deployment_selection_failure -from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, usage_twins +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, note_served_from_cache, usage_twins from litellm.router_utils.endpoint_eligibility import endpoint_for_call, selected_endpoint from litellm.router_utils.kubernetes_pod_discovery import ( KubernetesPodDiscovery, @@ -5030,6 +5030,9 @@ class _Driver: for _ in range(self._fallback_depth): verbose_router_logger.info("Successful fallback b/w models.") tokens: Final = total_tokens(response) + note_served_from_cache( + (as_str_mapping(getattr(response, "_hidden_params", None)) or {}).get("cache_hit") is True + ) # A wrapped stream reports its success or failure when it ends; the attempt counts the # request and the tokens read so far, which the end's report leaves out. # A realtime attempt counts its session's success itself (`attempt`), when it ended cleanly. diff --git a/litellm/router_rust/steps/selection.py b/litellm/router_rust/steps/selection.py index a141f99d1ad..cdd15164bea 100644 --- a/litellm/router_rust/steps/selection.py +++ b/litellm/router_rust/steps/selection.py @@ -357,9 +357,10 @@ async def routing_pre_call_checks( fails the attempt; the core records that failure like any other.""" checks: Final = tuple(callback for callback in callbacks if isinstance(callback, CustomLogger)) charge: Final = await counter.charge_request(deployment) if reads_deployment_usage(checks) else None + handed: Final = deployment if isinstance(deployment, dict) else dict(deployment) try: for callback in checks: - await callback.async_pre_call_check(dict(deployment), None) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict + await callback.async_pre_call_check(handed, None) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict except Exception: if charge is not None: await counter.refund_request(charge) @@ -375,9 +376,10 @@ def routing_pre_call_checks_sync( """``routing_strategy_pre_call_checks``: the sync twin, with the sync check of each ``CustomLogger``.""" checks: Final = tuple(callback for callback in callbacks if isinstance(callback, CustomLogger)) charge: Final = counter.charge_request_sync(deployment) if reads_deployment_usage(checks) else None + handed: Final = deployment if isinstance(deployment, dict) else dict(deployment) try: for callback in checks: - callback.pre_call_check(dict(deployment)) # pyright: ignore[reportUnknownMemberType] # see above + callback.pre_call_check(handed) # pyright: ignore[reportUnknownMemberType] # see above except Exception: if charge is not None: counter.refund_request_sync(charge) diff --git a/litellm/router_utils/deployment_usage_counter.py b/litellm/router_utils/deployment_usage_counter.py index 55865e068de..d2afda69508 100644 --- a/litellm/router_utils/deployment_usage_counter.py +++ b/litellm/router_utils/deployment_usage_counter.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from datetime import datetime from typing import TYPE_CHECKING, Final -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger @@ -42,6 +42,13 @@ class UsageCharge: _ADMITTED: Final[ContextVar[UsageCharge | None]] = ContextVar("deployment_usage_charge", default=None) +_SERVED_FROM_CACHE: Final[ContextVar[bool]] = ContextVar("deployment_usage_served_from_cache", default=False) + + +def note_served_from_cache(cache_hit: bool) -> None: + """A response the cache answered used no deployment tokens; the usage write that follows in this + attempt's context leaves the enforcement counters alone (``usage_twins``).""" + _SERVED_FROM_CACHE.set(cache_hit) def charged_rpm(rpm_key: str) -> int | None: @@ -62,10 +69,17 @@ def reads_deployment_usage(callbacks: Iterable[object]) -> bool: return any(getattr(callback, "reads_deployment_usage", False) is True for callback in callbacks) +def _fields(value: object) -> Mapping[str, object]: + """A deployment section as a mapping, whether it came as a dict or as its pydantic model.""" + if isinstance(value, BaseModel): + return value.model_dump() + return _STR_MAP.validate_python(value or {}) + + def usage_counter_keys(deployment: Mapping[str, object], minute: str) -> tuple[str, str]: """The ``(tpm_key, rpm_key)`` a deployment's usage is counted under this minute.""" - model_id: Final = _STR_MAP.validate_python(deployment.get("model_info") or {}).get("id") - model: Final = _STR_MAP.validate_python(deployment.get("litellm_params") or {}).get("model") + model_id: Final = _fields(deployment.get("model_info")).get("id") + model: Final = _fields(deployment.get("litellm_params")).get("model") return f"{model_id}:{model}:tpm:{minute}", f"{model_id}:{model}:rpm:{minute}" @@ -83,8 +97,9 @@ def enforcement_tpm_operation(operation: RedisPipelineIncrementOperation) -> Red def usage_twins( operations: Iterable[RedisPipelineIncrementOperation], counted: bool ) -> tuple[RedisPipelineIncrementOperation, ...]: - """The enforcement increments that ride the router's usage write, none when nothing reads them.""" - if not counted: + """The enforcement increments that ride the router's usage write, none when nothing reads them or + when the cache answered the call.""" + if not counted or _SERVED_FROM_CACHE.get(): return () return tuple(twin for twin in map(enforcement_tpm_operation, operations) if twin is not None) @@ -172,6 +187,7 @@ __all__ = [ "charged_rpm", "end_attempt", "enforcement_tpm_operation", + "note_served_from_cache", "reads_deployment_usage", "usage_counter_keys", "usage_twins", From db6e269ecaa2b63efa3f2d56213b909bd262bcaa Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 03:45:21 +0000 Subject: [PATCH 09/12] fix(router): charge the request before taking the deployment slot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_rust/core.py | 30 +++++++++++++----- litellm/router_rust/steps/selection.py | 43 ++++++++++++++++++++------ 2 files changed, 56 insertions(+), 17 deletions(-) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index cfddcd5989b..9f0d7de43fd 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -124,6 +124,7 @@ from litellm.router_rust.steps.prerouting import PreRouter, PreRoutingStrategies from litellm.router_rust.steps.selection import ( SELECTOR_ATTRIBUTES, StrategySelectors, + charge_for_checks, group_documents, group_settings, member_map, @@ -132,6 +133,7 @@ from litellm.router_rust.steps.selection import ( pick_async, pick_sync, picked_id, + pre_call_checks_with, routing_pre_call_checks, routing_pre_call_checks_sync, ) @@ -5174,15 +5176,27 @@ class _Driver: return {**self._failed(step, deployment, call, refused), "recorded": True} # The deployment's max_parallel_requests slot; an open chat stream takes it along. port = self._core._port # pyright: ignore[reportPrivateUsage] # the driver serves its own core # rebind-ok: the same port on every attempt - with enter_deployment_slot(port, self._plan, deployment, call.kwargs) as slot: - # A search tool is no deployment: nothing to check; Python's `_arerank` and - # `amoderation` run no routing pre-call checks either. - if self._plan.kind not in _NO_ROUTING_PRE_CALL_CHECKS: + # A search tool is no deployment: nothing to check; Python's `_arerank` and + # `amoderation` run no routing pre-call checks either. + checked: Final = self._plan.kind not in _NO_ROUTING_PRE_CALL_CHECKS + checks: Final = self._core._checks # pyright: ignore[reportPrivateUsage] # see above + picked_deployment: Final = as_str_mapping(step.get("deployment")) or {} + charge: Final = ( + await charge_for_checks(checks.routing_callbacks(), picked_deployment, checks.usage_counter) + if checked + else None + ) + try: + slot_taken: Final = enter_deployment_slot(port, self._plan, deployment, call.kwargs) + except Exception as refused: # noqa: BLE001 # a full deployment's 429, handed to the Rust call loop + if charge is not None: + await checks.usage_counter.refund_request(charge) + return await self._refused(step, deployment, call, refused) + with slot_taken as slot: + if checked: try: - await routing_pre_call_checks( - self._core._checks.routing_callbacks(), # pyright: ignore[reportPrivateUsage] # see above - as_str_mapping(step.get("deployment")) or {}, - self._core._checks.usage_counter, # pyright: ignore[reportPrivateUsage] # see above + await pre_call_checks_with( + checks.routing_callbacks(), picked_deployment, checks.usage_counter, charge ) except Exception as refused: # noqa: BLE001 # handed to the Rust call loop as the attempt's failure # Only `_acompletion` hands the check its logging object. diff --git a/litellm/router_rust/steps/selection.py b/litellm/router_rust/steps/selection.py index cdd15164bea..bed528ec3a3 100644 --- a/litellm/router_rust/steps/selection.py +++ b/litellm/router_rust/steps/selection.py @@ -30,7 +30,12 @@ from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 -from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, end_attempt, reads_deployment_usage +from litellm.router_utils.deployment_usage_counter import ( + DeploymentUsageCounter, + UsageCharge, + end_attempt, + reads_deployment_usage, +) from litellm.types.router import RoutingGroup, RoutingStrategy if TYPE_CHECKING: @@ -348,18 +353,31 @@ def pick_sync( return None -async def routing_pre_call_checks( +def _checks_among(callbacks: Iterable[object]) -> tuple[CustomLogger, ...]: + return tuple(callback for callback in callbacks if isinstance(callback, CustomLogger)) + + +async def charge_for_checks( callbacks: Iterable[object], deployment: Mapping[str, object], counter: DeploymentUsageCounter +) -> UsageCharge | None: + """The attempt's request counted against the picked deployment, once, when a routing callback + reads the usage counters. Charged before the deployment's slot is taken, so the cache round trip + never holds a slot another request is waiting for.""" + return await counter.charge_request(deployment) if reads_deployment_usage(_checks_among(callbacks)) else None + + +async def pre_call_checks_with( + callbacks: Iterable[object], + deployment: Mapping[str, object], + counter: DeploymentUsageCounter, + charge: UsageCharge | None, ) -> None: - """``async_routing_strategy_pre_call_checks``: the attempt's request is counted against the picked - deployment once, then every ``CustomLogger`` among the router's routing callbacks - (``_routing_callbacks``) checks it before the call. A check's error hands the request back and - fails the attempt; the core records that failure like any other.""" - checks: Final = tuple(callback for callback in callbacks if isinstance(callback, CustomLogger)) - charge: Final = await counter.charge_request(deployment) if reads_deployment_usage(checks) else None + """Every ``CustomLogger`` among the router's routing callbacks (``_routing_callbacks``) checks the + picked deployment before the call. A check's error hands the charged request back and fails the + attempt; the core records that failure like any other.""" handed: Final = deployment if isinstance(deployment, dict) else dict(deployment) try: - for callback in checks: + for callback in _checks_among(callbacks): await callback.async_pre_call_check(handed, None) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict except Exception: if charge is not None: @@ -370,6 +388,13 @@ async def routing_pre_call_checks( end_attempt(charge) +async def routing_pre_call_checks( + callbacks: Iterable[object], deployment: Mapping[str, object], counter: DeploymentUsageCounter +) -> None: + """``async_routing_strategy_pre_call_checks``: ``charge_for_checks`` then ``pre_call_checks_with``.""" + await pre_call_checks_with(callbacks, deployment, counter, await charge_for_checks(callbacks, deployment, counter)) + + def routing_pre_call_checks_sync( callbacks: Iterable[object], deployment: Mapping[str, object], counter: DeploymentUsageCounter ) -> None: From 4facf2bd56014e64ca1b39ad35896bcac3e8965c Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 03:46:02 +0000 Subject: [PATCH 10/12] test(router): the minute's tpm is the router's usage write to count, websocket turns use the full rpm Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/test_router/core/test_gaps_l8400_precall.py | 7 +++---- tests/unit/test_router/test_io_token_rate_limits.py | 10 +++++----- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/unit/test_router/core/test_gaps_l8400_precall.py b/tests/unit/test_router/core/test_gaps_l8400_precall.py index 1e64257bac8..be6d943deb5 100644 --- a/tests/unit/test_router/core/test_gaps_l8400_precall.py +++ b/tests/unit/test_router/core/test_gaps_l8400_precall.py @@ -163,16 +163,15 @@ class TestResponsesWebsocketTurnAdmission: async def test_usage_based_routing_admits_turns_until_the_deployment_rpm_is_used_up( self, make_core: MakeCore, host: EndpointHost ) -> None: - """Opening the socket already counts two requests (its pre-call check and its success), so an - rpm of 5 leaves room for three turns.""" + """Opening the socket counts one request, so an rpm of 5 leaves room for four turns.""" core: Final = make_core(model_list=[deployment("g", "a", rpm=5)], routing_strategy="usage-based-routing-v2") admit: Final = await self.open_socket(core, host) - admitted: Final = [dict(await admit()) for _ in range(3)] # pyright: ignore[reportGeneralTypeIssues] + admitted: Final = [dict(await admit()) for _ in range(4)] # pyright: ignore[reportGeneralTypeIssues] with pytest.raises(litellm.RateLimitError) as refused: await admit() # pyright: ignore[reportGeneralTypeIssues] - assert admitted == [{}, {}, {}] + assert admitted == [{}, {}, {}, {}] assert str(refused.value).startswith( "litellm.RateLimitError: No deployments available for selected model. Passed model=g. Deployments=" ) 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 f09ad5973e2..8e88e473c68 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -800,9 +800,9 @@ class TestModelRateLimitingCheckIOTokens: # The first deployment's ITPM counter is not driven negative... assert await dual_cache.async_get_cache(key=itpm_key) == 0 - # ...and the non-IO deployment's TPM usage is tracked normally. + # ...and the minute's TPM is the router's usage write to count, not the check's. tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" - assert await dual_cache.async_get_cache(key=tpm_key) == 12 + assert await dual_cache.async_get_cache(key=tpm_key) is None @pytest.mark.asyncio async def test_stale_reservation_refunded_before_retry_overwrites_it(self): @@ -1236,8 +1236,8 @@ class TestModelRateLimitingCheckIOTokens: # ITPM reconciled down from the 5-token reservation to actual usage (3). assert await dual_cache.async_get_cache(key=itpm_key) == 3 - # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. - assert await dual_cache.async_get_cache(key=tpm_key) == 7 + # The minute's TPM is the router's usage write to count, not the check's. + assert await dual_cache.async_get_cache(key=tpm_key) is None def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): from litellm.utils import get_utc_datetime @@ -1267,7 +1267,7 @@ class TestModelRateLimitingCheckIOTokens: check.log_success_event(kwargs, response, None, None) assert dual_cache.get_cache(key=itpm_key) == 3 - assert dual_cache.get_cache(key=tpm_key) == 7 + assert dual_cache.get_cache(key=tpm_key) is None @pytest.mark.asyncio async def test_failure_refunds_itpm_reservation(self): From 1b83778beac3c92e3b29c2b44140abbdb37f8be6 Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 03:53:39 +0000 Subject: [PATCH 11/12] fix(router): a full deployment's 429 raises as before, with the charged request handed back Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_rust/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 9f0d7de43fd..3278ad63508 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -5188,10 +5188,10 @@ class _Driver: ) try: slot_taken: Final = enter_deployment_slot(port, self._plan, deployment, call.kwargs) - except Exception as refused: # noqa: BLE001 # a full deployment's 429, handed to the Rust call loop + except Exception: # noqa: BLE001 # a full deployment's 429 raises as before, with the request handed back if charge is not None: await checks.usage_counter.refund_request(charge) - return await self._refused(step, deployment, call, refused) + raise with slot_taken as slot: if checked: try: From cfbadc54468b96d71b5b6b88af083b3a71ffcda6 Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 6 Nov 2026 06:40:41 +0000 Subject: [PATCH 12/12] fix(router): drop the unused _corrected_cooldowns helper left behind by the main merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_rust/core.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 533c01c0f1e..be9eab67989 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -3491,15 +3491,6 @@ class RustCore( read: Final[object] = self._cache.batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above - def _corrected_cooldowns(self, keys: Sequence[str], results: object) -> Sequence[object]: - """The cooldowns a batch read found, once ``active_cooldowns_from_results`` corrected the in-memory - TTL of each it promoted from Redis and dropped each that expired.""" - if not isinstance(results, list): - return [None] * len(keys) - ids: Final = [key.removeprefix("deployment:").removesuffix(":cooldown") for key in keys] - _ = self._cooldowns.active_cooldowns_from_results(ids, results) # pyright: ignore[reportUnknownArgumentType] # see above - return results # pyright: ignore[reportUnknownVariableType] # see above - async def _increment_usage(self, keys: Sequence[object], amounts: Sequence[object]) -> None: try: counted: Final = await self._stores.increment_usage(keys, amounts)