scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge cfbadc5446 into 5103d8a0ee

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-11-07 07:24:21 +00:00 • committed by GitHub
commit 8fa0b027d0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 943 additions and 263 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
litellm/router_rust/core.py
View file

@ -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("<routing_strategy>.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.

9
litellm/router_rust/invoke/endpoints.py
View file

@ -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})

6
litellm/router_rust/steps/callbacks.py
View file

@ -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 = []

76
litellm/router_rust/steps/selection.py
View file

@ -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:

28
litellm/router_rust/surface/legacy_members.py
View file

@ -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(

66
litellm/router_strategy/lowest_tpm_rpm_v2.py
View file

@ -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:

212
litellm/router_utils/deployment_usage_counter.py Normal file
View file

@ -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",
]

82
litellm/router_utils/pre_call_checks/model_rate_limit_check.py
View file

@ -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:

103
tests/integration/routing/test_enforce_model_rate_limits_routing_groups.py Normal file
View file

@ -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

37
tests/unit/caching/test_request_redis_batch_post_call.py
View file

@ -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()

78
tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py
View file

@ -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",

62
tests/unit/router_strategy/test_router_routing_groups.py
View file

@ -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

125
tests/unit/router_utils/test_deployment_usage_counter.py Normal file
View file

@ -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())))

7
tests/unit/test_router/core/test_gaps_l8400_precall.py
View file

@ -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="
)

249
tests/unit/test_router/test_enforce_model_rate_limits.py
View file

@ -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

10
tests/unit/test_router/test_io_token_rate_limits.py
View file

@ -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):