diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index eef454862f9..0c858828c7b 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -130,6 +130,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, @@ -138,7 +139,9 @@ 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, ) from litellm.router_rust.surface.alerting import initialize_alerting from litellm.router_rust.surface.fallbacks import ( @@ -211,6 +214,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_handlers import get_cooldown_deployment_info, mark_deployment_selection_failure +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter, note_served_from_cache from litellm.router_utils.endpoint_eligibility import endpoint_for_call, selected_endpoint from litellm.router_utils.kubernetes_pod_discovery import ( KubernetesPodDiscovery, @@ -1320,6 +1324,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 @@ -3375,7 +3382,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) @@ -3409,7 +3416,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) @@ -3436,7 +3443,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) @@ -3458,7 +3465,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 ----------------------------------------------------------------------------- @@ -3498,6 +3505,15 @@ class RustCore( return if not counted: raise _not_usage(keys) + from litellm.types.caching import RedisPipelineIncrementOperation + + await self._checks.usage_counter.charge_tokens( + tuple( + RedisPipelineIncrementOperation(key=str(key), increment_value=_amount(amount), ttl=None) + for key, amount in zip(keys, amounts) + ), + self._checks.counts_deployment_usage(), + ) async def _answer_cache(self, step: Mapping[str, object]) -> object: with service_target(_cache_target(str(step.get("key")))): @@ -5031,6 +5047,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. @@ -5052,7 +5071,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, ) @@ -5173,14 +5192,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: # 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) + raise + 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 {}, + 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/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index 7a5b52761e2..7f373f03a8f 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: @@ -2452,16 +2455,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 cd2d69d13a8..af4dd71d1c9 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, @@ -118,6 +119,7 @@ class OptionalChecks: self._cache: Final = cache self._hosted: Final = hosted self._ctx: Final = ctx + 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") @@ -219,6 +221,10 @@ class OptionalChecks: session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, # pyright: ignore[reportArgumentType] # same ) + 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: 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..bed528ec3a3 100644 --- a/litellm/router_rust/steps/selection.py +++ b/litellm/router_rust/steps/selection.py @@ -30,6 +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, + UsageCharge, + end_attempt, + reads_deployment_usage, +) from litellm.types.router import RoutingGroup, RoutingStrategy if TYPE_CHECKING: @@ -347,21 +353,65 @@ 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): - await callback.async_pre_call_check(dict(deployment), None) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict +def _checks_among(callbacks: Iterable[object]) -> tuple[CustomLogger, ...]: + return tuple(callback for callback in callbacks if isinstance(callback, CustomLogger)) -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): - callback.pre_call_check(dict(deployment)) # pyright: ignore[reportUnknownMemberType] # see above +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: + """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_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: + await counter.refund_request(charge) + raise + finally: + if charge is not None: + 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: + """``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(handed) # pyright: ignore[reportUnknownMemberType] # see above + except Exception: + 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 bef35673aaa..ea1234beba3 100644 --- a/litellm/router_rust/surface/legacy_members.py +++ b/litellm/router_rust/surface/legacy_members.py @@ -38,6 +38,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, @@ -45,9 +46,11 @@ 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 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 +from litellm.types.router import RouterCacheEnum from litellm.utils import get_utc_datetime if TYPE_CHECKING: @@ -332,10 +335,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) @with_service_target(ROUTER_COOLDOWNS_TARGET) async def async_routing_strategy_pre_call_checks( @@ -348,11 +350,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 + 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(charge, parent_otel_span) _dispatch_failure(logging_obj, e) _ = set_cooldown_deployments( ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the core's RoutingContext @@ -363,8 +368,10 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): ) raise except Exception as e: + await counter.refund_request(charge, parent_otel_span) _dispatch_failure(logging_obj, e) raise + end_attempt(charge) async def async_callback_filter_deployments( self, @@ -474,10 +481,21 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): parent_otel_span: "Span | None", ) -> None: """``_increment_deployment_usage`` with its one request.""" - if not self._registry_state.cached_deployment_tracks_usage(deployment_id, model_group): + from litellm.types.caching import RedisPipelineIncrementOperation + + 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") await self._usage.increment_deployment(deployment_id, deployment_name, current_minute, total_tokens, 1) + tpm_key: Final = RouterCacheEnum.TPM.value.format( + id=deployment_id, current_minute=current_minute, model=deployment_name + ) + await self._opt.usage_counter.charge_tokens( + (RedisPipelineIncrementOperation(key=tpm_key, increment_value=total_tokens, ttl=_USAGE_TTL),), + counted, + parent_otel_span, + ) @with_service_target(ROUTER_USAGE_TARGET) async def async_deployment_callback_on_failure( diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index d34d7edf9e4..3b05aafdcfd 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 @@ -16,10 +16,10 @@ 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.redis.stores.router_usage import ROUTER_USAGE_TARGET +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 @@ -87,6 +87,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 @@ -134,7 +136,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="", @@ -155,9 +157,10 @@ 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) + 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}", @@ -214,7 +217,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="", @@ -231,8 +234,10 @@ 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) + 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}", @@ -260,9 +265,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.") @@ -274,22 +277,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 @@ -310,32 +297,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: @@ -388,7 +354,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/litellm/router_utils/deployment_usage_counter.py b/litellm/router_utils/deployment_usage_counter.py new file mode 100644 index 00000000000..500e669f539 --- /dev/null +++ b/litellm/router_utils/deployment_usage_counter.py @@ -0,0 +1,212 @@ +"""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 and hands it back when a check +refuses the attempt. The tokens of a served call follow the router's own usage write, as the twin of its +``global_router:`` tpm increment (``charge_tokens``). 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, 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.redis.stores.router_usage 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 +_STR_MAP: Final = TypeAdapter(Mapping[str, object]) +_ROUTER_USAGE_PREFIX: Final = RouterCacheEnum.TPM.value.split("{", 1)[0] + + +@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 + + +_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: + """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 (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 _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 = _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}" + + +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 usage_twins( + operations: Iterable[RedisPipelineIncrementOperation], counted: bool +) -> tuple[RedisPipelineIncrementOperation, ...]: + """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) + + +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.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 + ) -> 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, 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]) -> 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, charge: UsageCharge) -> None: + _ = self._increment(charge.key, -1) + end_attempt(charge) + + @with_service_target(ROUTER_USAGE_TARGET) + async def charge_tokens( + self, + operations: Iterable[RedisPipelineIncrementOperation], + counted: bool, + parent_otel_span: Span | None = None, + ) -> None: + """The served call's tokens, as the twins of the router's ``global_router:`` tpm increments: counted + in this worker's memory now, in Redis on the request's post-call pipeline.""" + twins: Final = usage_twins(operations, counted) + if not twins: + return + try: + await self.dual_cache.async_increment_cache_pipeline_post_call(twins, parent_otel_span) + except Exception as e: # noqa: BLE001 # see _async_increment + verbose_router_logger.debug("DeploymentUsageCounter: tokens were not recorded. Got - %s", e) + + 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) + 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, 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: + try: + 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 + if self.dual_cache.redis_cache is not None: + 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 + + +__all__ = [ + "ROUTER_USAGE_TARGET", + "USAGE_TTL", + "DeploymentUsageCounter", + "UsageCharge", + "charged_rpm", + "end_attempt", + "enforcement_tpm_operation", + "note_served_from_cache", + "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 38ae7e0d951..22df2f3b646 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 @@ -26,6 +26,7 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.integrations.custom_logger import CustomLogger from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +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 @@ -316,10 +319,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 +348,12 @@ 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) + 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: - 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 +369,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 +400,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 +437,35 @@ 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", - ), - ), - num_retries=0, # Don't retry - return 429 immediately + if rpm_limit is not None: + 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}", + 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 + ) - 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 +511,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 +557,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/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 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 41d508a0eea..58eb7538d4e 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -9,13 +9,11 @@ import asyncio import datetime import hashlib import json -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Sequence from typing import Final import pytest -from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters - import litellm from litellm.caching.caching import Cache from litellm.caching.dual_cache import DualCache @@ -36,10 +34,12 @@ 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.redis.stores.rate_limits import CounterIncrement, RateLimitScope -from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 +from litellm.router_utils.deployment_usage_counter import DeploymentUsageCounter from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.router import RouterCacheEnum from litellm.types.utils import ModelResponse from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters from .test_redis_batch import FakeClient, FakeNatives, FakeRedisCache, fake_send @@ -132,22 +132,15 @@ def _response_cache(redis_cache: FakeRedisCache) -> Cache: return cache -def _tpm_router(redis_cache: FakeRedisCache) -> tuple[LowestTPMLoggingHandler_v2, DualCache]: +def _tpm_counter(redis_cache: FakeRedisCache) -> tuple[DeploymentUsageCounter, DualCache]: router_cache = DualCache() router_cache.attach_redis_cache(redis_cache) - return LowestTPMLoggingHandler_v2(router_cache=router_cache, routing_args={"ttl": 60}), router_cache + return DeploymentUsageCounter(router_cache), 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": {}}, - } +def _router_tpm_op(total_tokens: int) -> RedisPipelineIncrementOperation: + key = RouterCacheEnum.TPM.value.format(id="dep-a", model="openai/gpt-4o-mini", current_minute="00-00") + return RedisPipelineIncrementOperation(key=key, increment_value=total_tokens, ttl=60) @pytest.mark.asyncio @@ -157,10 +150,10 @@ async def test_post_call_owners_wait_for_the_callbacks_and_the_limiter_declares_ calls: list[tuple[str, list[str], list[object]]] = [] # mutable-ok: filled by the fake scripts fakes = FakeRateLimits(token_increment_script=_scripts(calls, [10.0]), parallel_release_script=_scripts(calls, [0])) limiter = _limiter(fakes) - tpm, router_cache = _tpm_router(redis_cache) + counter, router_cache = _tpm_counter(redis_cache) with request_redis_batch_scope(): - await tpm.async_log_success_event(_tpm_kwargs(), None, None, None) + await counter.charge_tokens((_router_tpm_op(42),), counted=True) await limiter.async_increment_tokens_with_ttl_preservation(_token_ops(_K1)) await limiter._release_stashed_parallel_slot(_slot_stash("slot-1", _K1), None, in_logging_callback=True) assert client.pipelines == [] # nothing goes out while the callbacks are still declaring @@ -437,7 +430,9 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_request_boundary(): natives = FakeNatives() - with request_redis_batch_scope(post_call_deadline=60, native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + with request_redis_batch_scope( + post_call_deadline=60, native_factory=natives, native_send_factory=fake_send + ) as request: # pyright: ignore[reportArgumentType] # test double request.native_post_call() await request.flush_all() at_the_boundary = list(natives.built[0].events) @@ -452,7 +447,9 @@ async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_requ async def test_native_post_call_ops_still_waiting_on_their_callbacks_are_drained_at_shutdown(): natives = FakeNatives() - with request_redis_batch_scope(post_call_deadline=60, native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + with request_redis_batch_scope( + post_call_deadline=60, native_factory=natives, native_send_factory=fake_send + ) as request: # pyright: ignore[reportArgumentType] # test double request.native_post_call() await request.flush_all() await drain_post_call_redis_batches() 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..e1aa58f7f9b 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,9 @@ 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 datetime import datetime +from typing import Final from unittest.mock import patch import pytest @@ -23,12 +20,17 @@ 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.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 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 +47,28 @@ def restore_verbose_logger_level() -> Iterator[None]: @pytest.mark.usefixtures("restore_verbose_logger_level") + + +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, 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) + + def test_tpm_rpm_updated(): from tests.unit.create_mock_standard_logging_payload import create_standard_logging_payload @@ -80,8 +104,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, + test_cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -140,7 +166,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, + test_cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -167,7 +195,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, + test_cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -245,7 +275,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.cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -269,7 +301,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.cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -331,7 +365,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.cache, response_obj=response_obj, kwargs=kwargs, start_time=start_time, @@ -345,9 +381,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 +411,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 +521,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 +609,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/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 8a0d5b9a1d6..50944be0f88 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,8 +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 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(): @@ -1150,3 +1150,61 @@ 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() + 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 + _, 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: + 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() + 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, _ = 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 new file mode 100644 index 00000000000..4f8cb3336cc --- /dev/null +++ b/tests/unit/router_utils/test_deployment_usage_counter.py @@ -0,0 +1,125 @@ +import asyncio +from datetime import datetime, timedelta, timezone + +import pytest + +from litellm.caching.dual_cache import DualCache +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" + + +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 +async def test_each_charged_request_adds_one_to_the_minute(): + cache = DualCache() + counter = _counter(cache) + + charges = [await counter.charge_request(DEPLOYMENT) for _ in range(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 + + +@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) + second = await counter.charge_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) + + 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_check_reads_the_count_its_own_charge_was_admitted_with(): + cache = DualCache() + counter = _counter(cache) + + 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 sorted(await asyncio.gather(attempt(), attempt(), attempt())) == [1, 2, 3] + assert charged_rpm(RPM_KEY) is None + + +@pytest.mark.asyncio +async def test_a_charge_answers_only_for_its_own_key(): + counter = _counter(DualCache()) + + charge = await counter.charge_request(DEPLOYMENT) + + assert charged_rpm(RPM_KEY) == 1 + assert charged_rpm("deploy-2:openai/gpt-4o:rpm:12-00") is None + end_attempt(charge) + + +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) + + 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_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/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_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index ff3ffc0db0c..fc10db49b2e 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -17,9 +17,14 @@ 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, 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 +from tests.unit.caching.test_redis_batch import FakeClient, FakeRedisCache TPM_DEPLOYMENT = { "tpm": 1000, @@ -39,6 +44,53 @@ 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 _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.""" + _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: + _count_logged_tokens(check, kwargs) + await check.async_log_success_event(kwargs, None, None, None) + + class TestModelRateLimitingCheck: """Test the ModelRateLimitingCheck class directly.""" @@ -105,16 +157,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 +170,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 +192,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 +210,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 +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) + _pre_call(check, TPM_DEPLOYMENT) assert "current usage=1000" in str(exc_info.value) @@ -184,7 +234,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 +245,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 +253,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 +272,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 +285,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 +294,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 +308,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 +327,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 +346,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 +357,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 +371,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 +385,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 +394,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 +414,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 +439,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 +490,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 +521,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 +634,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 @@ -882,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 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):