scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

fix(router): cap the default OTPM reservation and reconcile token counters after each call (#49054)

Browse source 
* fix(router): cap the default OTPM reservation for requests without max_tokens

* fix(router): reconcile ITPM and OTPM reservations after every successful call

* fix(router): read the default OTPM reservation from the served model and make its fraction configurable

The default reservation now looks up max output from litellm_params.model, as main does, so a pricing base_model never changes it. NO_MAX_TOKENS_TPM_RESERVATION_FRACTION sets the fraction for both the router check and the key and team TPM floor. A max_tokens that is not a number or string counts as no cap instead of raising inside the check.

* fix(router): refund a failed attempt's io token reservation in make_call

The proxy passes one logging object to every retry, and it runs failure
callbacks only for the first attempt, so the last attempt's ITPM and
OTPM reservation stayed held until the minute ended. Each attempt that
raises now gives its reservation back before the retry or fallback

* fix(router): count Anthropic cache writes once in ITPM reconcile

Native /v1/messages usage reports input_tokens without cache reads and writes, so the reconcile now maps it through AnthropicConfig.calculate_usage before subtracting cached tokens.

* fix(router): use a deployment max_tokens, charge late usage to its finish minute, and read the reservation fraction as a fraction

* fix(router): reserve with the deployment cap only through the param the request's API sends

* fix(router): never let a deployment's output cap raise the reservation above the default

* test(router): move the reservation fraction env var test to the integration suite

The unit test started a child interpreter, which the unit guide forbids. The integration test boots a real proxy with the env var set and checks admission against the held reservation.

* test(router): retry the reservation fraction scenario when the minute rolls over

The OTPM counter is keyed by minute, so a rollover between the held stream and the capped requests gave the capped request a fresh counter. The control request now leaves room for usage the first attempt left in the new minute.

* test(router): give the rollover retry its own deployment

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
paveln83 2026-10-29 20:32:04 -07:00 • committed by GitHub
parent 342ce6ae4a
commit b91fc92b6d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 831 additions and 138 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/constants.py
View file

@ -7,6 +7,7 @@ from typing import Final, Literal
from litellm.litellm_core_utils.env_utils import (
get_env_float,
get_env_float_in_range,
get_env_float_or_none,
get_env_int,
get_env_int_in_range,
@ -149,6 +150,9 @@ DEFAULT_FAILURE_THRESHOLD_PERCENT: Final = get_env_float(
"DEFAULT_FAILURE_THRESHOLD_PERCENT", 0.5
) # default cooldown a deployment if 50% of requests fail in a given minute
DEFAULT_MAX_TOKENS: Final = get_env_int("DEFAULT_MAX_TOKENS", 4096)
NO_MAX_TOKENS_TPM_RESERVATION_FRACTION: Final = get_env_float_in_range(
"NO_MAX_TOKENS_TPM_RESERVATION_FRACTION", 0.25, 0.01, 1.0
)
DEFAULT_ALLOWED_FAILS: Final = get_env_int("DEFAULT_ALLOWED_FAILS", 3)
DEFAULT_REDIS_SYNC_INTERVAL: Final = get_env_int("DEFAULT_REDIS_SYNC_INTERVAL", 1)
DEFAULT_COOLDOWN_TIME_SECONDS: Final = get_env_int("DEFAULT_COOLDOWN_TIME_SECONDS", 5)

10
litellm/litellm_core_utils/env_utils.py
View file

@ -71,7 +71,15 @@ def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int)
Values outside the range fall back to the default and warn, so a misconfigured knob can
neither crash the caller nor silently change the meaning of what it computes.
"""
value: Final = get_env_int(env_var, default)
return _within_range(env_var, get_env_int(env_var, default), default, minimum, maximum)
def get_env_float_in_range(env_var: str, default: float, minimum: float, maximum: float) -> float:
"""Float counterpart of `get_env_int_in_range`."""
return _within_range(env_var, get_env_float(env_var, default), default, minimum, maximum)
def _within_range(env_var: str, value: _NumberT, default: _NumberT, minimum: _NumberT, maximum: _NumberT) -> _NumberT:
if minimum <= value <= maximum:
return value
logging.getLogger("LiteLLM").warning(

17
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -42,7 +42,11 @@ from litellm.caching.redis_batch import (
active_request_redis_batch,
)
from litellm.caching.redis_cache import log_redis_failure
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.constants import (
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
NO_MAX_TOKENS_TPM_RESERVATION_FRACTION,
)
from litellm.cost_calculator import get_usage_object
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens
@ -508,11 +512,6 @@ REDIS_NODE_HASHTAG_NAME: Final = "all_keys"
# *some* output budget; these define that fallback estimate.
DEFAULT_MAX_TOKENS_ESTIMATE: Final = 4096
DEFAULT_CHARS_PER_TOKEN: Final = 4
# Fraction of the available output budget reserved as the upfront floor when
# the request omits max_tokens. Applied to both DEFAULT_MAX_TOKENS_ESTIMATE
# (baseline floor) and to the smallest configured TPM limit (capped floor for
# small per-tenant TPM caps).
_TPM_FLOOR_FRACTION: Final = 4
# Both embeddings and the Responses API put their prompt in data["input"],
# but only embeddings have no output tokens. Every "is this an embedding"
# check on data["input"] must exclude these call types, or a Responses call
@ -1003,10 +1002,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
per-tenant cap can't be tripped by the floor alone. Returns the
baseline floor when no limit is provided.
"""
baseline: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION
baseline: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)
if min_configured_tpm_limit is None:
return baseline
return min(baseline, max(1, min_configured_tpm_limit // _TPM_FLOOR_FRACTION))
return min(baseline, max(1, int(min_configured_tpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)))
@staticmethod
def _is_embedding_request(data: object, call_type: str | None) -> bool:
@ -1138,7 +1137,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if call_type in RESPONSES_API_CALL_TYPES
else base_capped_floor
)
baseline_floor: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION
baseline_floor: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)
is_embedding: Final = PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type)
if (
capped_floor >= baseline_floor

30
litellm/router.py
View file

@ -377,6 +377,7 @@ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check impor
EncryptedContentAffinityCheck,
)
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
async_io_token_refund_failure,
build_io_token_rate_limit_headers,
deployment_has_io_token_limits,
refund_stale_reservation_before_retry,
@ -9500,7 +9501,7 @@ class Router(DeprecatedRouterMembers):
raise original_exception
async def make_call(self, original_function: Any, *args, **kwargs):
async def make_call(self, original_function: Any, *args: object, **kwargs: object):
"""
Handler for making a call to the .completion()/.embeddings()/etc. functions.
"""
@ -9509,17 +9510,26 @@ class Router(DeprecatedRouterMembers):
endpoint_function: Final = (
generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function)
)
with endpoint_selection(
getattr(endpoint_function, "__name__", None),
getattr(endpoint_function, "__module__", None),
model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group),
):
response = original_function(*args, **kwargs)
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
response = await response
try:
with endpoint_selection(
getattr(endpoint_function, "__name__", None),
getattr(endpoint_function, "__module__", None),
model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group),
):
response = original_function(*args, **kwargs)
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
response = await response
except Exception:
with contextlib.suppress(Exception):
await async_io_token_refund_failure(
self.cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs)
)
raise
await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs)
## PROCESS RESPONSE HEADERS
response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs)
response = await self.set_response_headers(
response=response, model_group=model_group if isinstance(model_group, str) else None, request_kwargs=kwargs
)
return response

185
litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py
View file

@ -1,7 +1,8 @@
"""
Separate ITPM/OTPM (input/output tokens per minute) deployment rate limits.
- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM
- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM; a request
without max_tokens reserves the model's max output, capped at a fraction of the OTPM limit
- Post-call: reconcile ITPM to actual input tokens and OTPM to actual output tokens
- Cached prompt-read tokens are excluded from ITPM post-call accounting
@ -22,6 +23,7 @@ from litellm import token_counter
from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION
from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET
@ -127,10 +129,15 @@ def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] |
return itpm_key, otpm_key
def _tokenizer_model(deployment: object) -> str:
def _served_model(deployment: object) -> str:
params: Final = deployment.get("litellm_params") if is_str_keyed_mapping(deployment) else None
served: Final = params.get("model") if is_str_keyed_mapping(params) else None
return deployment_openai_model(deployment) or (served if isinstance(served, str) else "")
return served if isinstance(served, str) else ""
def _tokenizer_and_served_models(deployment: object) -> tuple[str, str]:
served: Final = _served_model(deployment)
return deployment_openai_model(deployment) or served, served
def _estimate_input_tokens(request_kwargs: dict[str, Any] | None, model: str = "") -> int:
@ -160,24 +167,61 @@ def _model_max_output_tokens(model_name: str) -> int | None:
return None
def _resolve_max_tokens(request_kwargs: dict[str, Any] | None, deployment: dict) -> int:
if request_kwargs:
# An explicit max_tokens=0 must be honored, not treated as absent and
# replaced by the model default.
explicit = request_kwargs.get("max_tokens")
if explicit is None:
explicit = request_kwargs.get("max_completion_tokens")
if explicit is None:
explicit = request_kwargs.get("max_output_tokens")
if explicit is not None:
return max(0, int(explicit))
_CHAT_OUTPUT_CAP_KEYS: Final = ("max_tokens", "max_completion_tokens")
_RESPONSES_OUTPUT_CAP_KEYS: Final = ("max_output_tokens",)
_OUTPUT_CAP_KEYS: Final = _CHAT_OUTPUT_CAP_KEYS + _RESPONSES_OUTPUT_CAP_KEYS
model_name: Final = (deployment.get("litellm_params") or {}).get("model")
if model_name:
model_max: Final = _model_max_output_tokens(model_name)
if model_max is not None:
return model_max
return 4096
def _explicit_max_tokens(
request_kwargs: Mapping[str, object] | None, cap_keys: tuple[str, ...] = _OUTPUT_CAP_KEYS
) -> int | None:
if not request_kwargs:
return None
explicit: Final = next(
(request_kwargs[key] for key in cap_keys if request_kwargs.get(key) is not None),
None,
)
match explicit:
case int() | float() | str():
return max(0, int(explicit))
case _:
return None
def _sent_deployment_cap_keys(request_kwargs: Mapping[str, object] | None) -> tuple[str, ...]:
is_responses_request: Final = (
request_kwargs is not None and "input" in request_kwargs and "messages" not in request_kwargs
)
return _RESPONSES_OUTPUT_CAP_KEYS if is_responses_request else _CHAT_OUTPUT_CAP_KEYS
def _deployment_output_cap(request_kwargs: Mapping[str, object] | None, deployment: Mapping[str, object]) -> int | None:
params: Final = deployment.get("litellm_params")
return (
_explicit_max_tokens(params, _sent_deployment_cap_keys(request_kwargs))
if is_str_keyed_mapping(params)
else None
)
def _default_otpm_reservation(model_name: str, otpm_limit: int) -> int:
model_max: Final = _model_max_output_tokens(model_name) if model_name else None
share_of_limit: Final = max(1, int(otpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION))
return min(4096 if model_max is None else model_max, share_of_limit)
def _otpm_reservation(
request_kwargs: Mapping[str, object] | None,
deployment: Mapping[str, object],
model_name: str,
otpm_limit: int,
) -> int:
from_request: Final = _explicit_max_tokens(request_kwargs)
if from_request is not None:
return from_request
default: Final = _default_otpm_reservation(model_name, otpm_limit)
from_deployment: Final = _deployment_output_cap(request_kwargs, deployment)
return default if from_deployment is None else min(from_deployment, default)
def _get_usage_tokens(usage: object) -> tuple[int, int, int]:
@ -210,6 +254,12 @@ def _extract_response_usage(response_obj: object) -> object:
return getattr(response_obj, "usage", None)
def _openai_shaped_usage(usage: object) -> object:
if isinstance(usage, dict) and litellm.AnthropicConfig.is_anthropic_usage_object(usage):
return litellm.AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None)
return usage
def _usage_is_present(usage: object) -> bool:
"""
True only if usage carries an actual input/output breakdown.
@ -238,7 +288,7 @@ def _resolve_reconcile_usage_tokens(
fields. When usage cannot be resolved, return ``usage_resolved=False`` so
callers keep the pre-call reservation instead of refunding it as zero usage.
"""
usage: Final = _extract_response_usage(response_obj)
usage: Final = _openai_shaped_usage(_extract_response_usage(response_obj))
if _usage_is_present(usage):
prompt_tokens, completion_tokens, cached_tokens = _get_usage_tokens(usage)
return max(0, prompt_tokens - cached_tokens), completion_tokens, True
@ -261,6 +311,21 @@ def _resolve_reconcile_usage_tokens(
return 0, 0, False
def _key_for_this_minute(reserved_key: str) -> str:
counter, _, _ = reserved_key.rpartition(":")
return f"{counter}:{get_utc_datetime().strftime('%H-%M')}"
def _reconcile_increments(
reservation: tuple[int, int, str | None, str | None], billable_input: int, completion_tokens: int
) -> tuple[tuple[str, int], ...]:
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
deltas: Final = ((itpm_key, billable_input - itpm_reserved), (otpm_key, completion_tokens - otpm_reserved))
return tuple(
(key if delta < 0 else _key_for_this_minute(key), delta) for key, delta in deltas if key is not None and delta
)
def _stash_reservation_in_metadata(
request_kwargs: dict[str, object] | None,
*,
@ -311,7 +376,9 @@ def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object,
return ()
top_level: Final = (kwargs.get("metadata"), kwargs.get("litellm_metadata"))
litellm_params: Final = _as_mutable_mapping(kwargs.get("litellm_params"))
from_params: Final = () if litellm_params is None else (litellm_params.get("metadata"),)
from_params: Final = (
() if litellm_params is None else (litellm_params.get("metadata"), litellm_params.get("litellm_metadata"))
)
standard_logging_object: Final = _as_mutable_mapping(kwargs.get("standard_logging_object"))
from_logging_object: Final = () if standard_logging_object is None else (standard_logging_object.get("metadata"),)
return top_level + from_params + from_logging_object
@ -433,9 +500,11 @@ def io_token_pre_call_check(
return deployment
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
_model: Final = _tokenizer_model(deployment)
_model, served_model = _tokenizer_and_served_models(deployment)
estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model)
max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
otpm_to_reserve: Final = (
_otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0
)
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M")
@ -458,7 +527,7 @@ def io_token_pre_call_check(
)
if otpm_limit is not None:
otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit)
otpm_reserved = otpm_to_reserve
try:
_sync_increment_with_rollback(
dual_cache,
@ -497,9 +566,11 @@ async def async_io_token_pre_call_check(
return deployment
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
_model: Final = _tokenizer_model(deployment)
_model, served_model = _tokenizer_and_served_models(deployment)
estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model)
max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
otpm_to_reserve: Final = (
_otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0
)
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M")
@ -523,7 +594,7 @@ async def async_io_token_pre_call_check(
)
if otpm_limit is not None:
otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit)
otpm_reserved = otpm_to_reserve
try:
await _increment_with_rollback(
dual_cache,
@ -564,7 +635,8 @@ def io_token_reconcile_success(
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
if itpm_key is None and otpm_key is None:
return
@ -572,23 +644,8 @@ def io_token_reconcile_success(
try:
if usage_resolved:
if itpm_key is not None:
itpm_delta: Final = billable_input - itpm_reserved
if itpm_delta != 0:
dual_cache.increment_cache(
key=itpm_key,
value=itpm_delta,
ttl=RoutingArgsTTL,
)
if otpm_key is not None:
otpm_delta: Final = completion_tokens - otpm_reserved
if otpm_delta != 0:
dual_cache.increment_cache(
key=otpm_key,
value=otpm_delta,
ttl=RoutingArgsTTL,
)
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
@ -618,39 +675,23 @@ async def async_io_token_reconcile_success(
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
if itpm_key is None and otpm_key is None:
return
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
# Reconcile against the exact key that held the reservation (which encodes
# the reservation's minute), not a key recomputed at response time. This
# tracks actual usage even when the pre-call estimate was 0, and avoids a
# minute-boundary mismatch for calls that span into the next minute. Always
# clear the stash afterwards (even if an increment throws) so a retry or a
# duplicate success event can't re-process it.
# A refund goes back to the minute that held the reservation, while usage above the reservation is charged
# to the minute the response finished in, so a call that runs past its minute still counts against a live
# window. Always clear the stash afterwards (even if an increment throws) so a retry or a duplicate success
# event can't re-process it.
try:
if usage_resolved:
if itpm_key is not None:
itpm_delta: Final = billable_input - itpm_reserved
if itpm_delta != 0:
await dual_cache.async_increment_cache(
key=itpm_key,
value=itpm_delta,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
if otpm_key is not None:
otpm_delta: Final = completion_tokens - otpm_reserved
if otpm_delta != 0:
await dual_cache.async_increment_cache(
key=otpm_key,
value=otpm_delta,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
await dual_cache.async_increment_cache(
key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span
)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",

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

@ -26,7 +26,6 @@ from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
ITPM_RESERVED_KEY,
async_io_token_pre_call_check,
async_io_token_reconcile_success,
async_io_token_refund_failure,
@ -485,24 +484,12 @@ class ModelRateLimitingCheck(CustomLogger):
try:
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object")
# IO token reconciliation works purely from the cache keys stashed in
# kwargs/metadata, so it must run before the model_id guard below
# (which only the TPM-tracking path needs). Otherwise a request whose
# standard_logging_object lacks model_id would never return its
# reservation, leaving the counter elevated until the TTL expires.
slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {}
kwargs_metadata: Final = kwargs.get("metadata") or {}
if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata:
await async_io_token_reconcile_success(
self.dual_cache,
kwargs,
response_obj,
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
)
# Fall through: a deployment can also configure tpm/rpm alongside
# itpm/otpm, and that path's pre-call check reads the tpm_key
# counter tracked below, so it must still be incremented here.
await async_io_token_reconcile_success(
self.dual_cache,
kwargs,
response_obj,
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
)
if standard_logging_object is None:
return
@ -557,17 +544,7 @@ class ModelRateLimitingCheck(CustomLogger):
"""
try:
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object")
slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {}
kwargs_metadata: Final = kwargs.get("metadata") or {}
if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata:
io_token_reconcile_success(
self.dual_cache,
kwargs,
response_obj,
)
# Fall through: a deployment can also configure tpm/rpm alongside
# itpm/otpm, and that path's pre-call check reads the tpm_key
# counter tracked below, so it must still be incremented here.
io_token_reconcile_success(self.dual_cache, kwargs, response_obj)
if standard_logging_object is None:
return

137
tests/integration/routing/test_otpm_default_reservation_fraction.py Normal file
View file

@ -0,0 +1,137 @@
from __future__ import annotations
import dataclasses
import threading
import uuid
from collections.abc import Callable, Iterator
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from queue import SimpleQueue
from typing import Final
import httpx
import pytest
from integration._support.client import JSON_OBJECT, 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
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
GROUPS: Final = ("otpm-fraction-first", "otpm-fraction-retry")
OTPM_LIMIT: Final = 1000
FRACTION: Final = 0.5
DEFAULT_RESERVATION: Final = int(OTPM_LIMIT * FRACTION)
@dataclass(frozen=True, slots=True)
class Rig:
owned: OwnedProxy
upstream: Wire
gates: SimpleQueue[threading.Event]
@dataclass(frozen=True, slots=True)
class Admissions:
held_forwarded_max_tokens: bool
statuses: tuple[int, int, int]
one_minute: bool
def _reply_for(gates: SimpleQueue[threading.Event]) -> Callable[[Request], Reply]:
def reply(request: Request) -> Reply:
if is_model_info_probe(request):
return model_list_reply()
if JSON_OBJECT.validate_json(request.body).get("stream") is True:
return dataclasses.replace(
chat_reply("chatcmpl-held", OPENAI_MODEL, "held", stream=True), gate_before_first=gates.get(timeout=60)
)
return chat_reply("chatcmpl-capped", OPENAI_MODEL, "capped", stream=False)
return reply
def _config(directory: Path, upstream: Wire) -> Path:
return owned_config(
directory,
[
{
"model_name": group,
"litellm_params": {
"model": OPENAI_MODEL,
"api_base": upstream.url + "/v1",
"api_key": "synthetic-openai-key",
"otpm": OTPM_LIMIT,
},
"model_info": {"id": group},
}
for group in GROUPS
],
litellm_settings={"cache": False},
router_settings={"optional_pre_call_checks": ["enforce_model_rate_limits"], "num_retries": 0},
)
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
directory: Final = tmp_path_factory.mktemp("otpm-fraction")
gates: Final[SimpleQueue[threading.Event]] = SimpleQueue()
with (
gateway_from_environment() as gateway,
wire_server(_reply_for(gates)) as upstream,
owned_proxy_process(
gateway,
directory,
{"NO_MAX_TOKENS_TPM_RESERVATION_FRACTION": str(FRACTION)},
config=_config(directory, upstream),
) as owned,
):
yield Rig(owned, upstream, gates)
def _chat(rig: Rig, group: str, **params: int | bool) -> httpx.Response:
return rig.owned.gateway.request(
"POST",
"/v1/chat/completions",
{"model": group, "messages": [{"role": "user", "content": f"fraction {uuid.uuid4().hex}"}], **params},
)
def _await_forwarded_chat(upstream: Wire) -> Request:
request: Final = upstream.received.get(timeout=60)
return _await_forwarded_chat(upstream) if is_model_info_probe(request) else request
def _utc_minute() -> str:
return datetime.now(UTC).strftime("%Y-%m-%d %H:%M")
def _admissions_beside_a_held_stream(rig: Rig, group: str) -> Admissions:
release: Final = threading.Event()
rig.gates.put(release)
started: Final = _utc_minute()
with ThreadPoolExecutor(max_workers=1) as pool:
held: Final = pool.submit(_chat, rig, group, stream=True)
try:
forwarded: Final = JSON_OBJECT.validate_json(_await_forwarded_chat(rig.upstream).body)
over: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION + 1)
fits: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION)
ended: Final = _utc_minute()
finally:
release.set()
return Admissions(
held_forwarded_max_tokens="max_tokens" in forwarded,
statuses=(over.status_code, fits.status_code, held.result(timeout=60).status_code),
one_minute=started == ended,
)
def test_reservation_fraction_env_var_sets_what_a_request_without_max_tokens_holds(rig: Rig) -> None:
first: Final = _admissions_beside_a_held_stream(rig, GROUPS[0])
admissions: Final = first if first.one_minute else _admissions_beside_a_held_stream(rig, GROUPS[1])
assert admissions == Admissions(held_forwarded_max_tokens=False, statuses=(429, 200, 200), one_minute=True)

10
tests/unit/litellm_core_utils/test_env_utils.py
View file

@ -6,6 +6,7 @@ import pytest
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.env_utils import (
get_env_float,
get_env_float_in_range,
get_env_float_or_none,
get_env_int,
get_env_int_or_none,
@ -94,3 +95,12 @@ def test_unparsable_stdout_log_limit_does_not_break_logging(
assert "first line" in messages
assert "second line" in messages
assert len([message for message in messages if raw in message]) == 1
@pytest.mark.parametrize(("raw", "expected"), [("0.5", 0.5), ("1", 1.0), ("4", 0.25), ("0", 0.25), ("half", 0.25)])
def test_float_in_range_keeps_values_inside_the_range_and_falls_back_otherwise(
monkeypatch: pytest.MonkeyPatch, raw: str, expected: float
) -> None:
monkeypatch.setenv("ENV_UTILS_TEST_FRACTION", raw)
assert get_env_float_in_range("ENV_UTILS_TEST_FRACTION", 0.25, 0.01, 1.0) == expected

9
tests/unit/test_constants.py
View file

@ -19,7 +19,14 @@ from unittest.mock import patch
_ENV_HELPERS = frozenset(
{"get_env_int", "get_env_float", "get_env_int_in_range", "get_env_int_or_none", "get_env_float_or_none"}
{
"get_env_int",
"get_env_float",
"get_env_int_in_range",
"get_env_float_in_range",
"get_env_int_or_none",
"get_env_float_or_none",
}
)

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

@ -3,6 +3,8 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits).
"""
import asyncio
from collections.abc import Awaitable, Callable, Mapping
from datetime import datetime
from typing import Final
import pytest
@ -10,13 +12,16 @@ import pytest
import litellm
from litellm import Router
from litellm.caching.dual_cache import DualCache
from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
ITPM_CACHE_KEY,
ITPM_RESERVED_KEY,
OTPM_CACHE_KEY,
OTPM_RESERVED_KEY,
_explicit_max_tokens,
_reservation_value,
_resolve_max_tokens,
async_io_token_pre_call_check,
async_io_token_reconcile_success,
build_io_token_rate_limit_headers,
@ -30,7 +35,16 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
ModelRateLimitingCheck,
)
from litellm.types.router import RouterCacheEnum
from litellm.types.utils import ModelResponse, Usage
from litellm.utils import get_utc_datetime
MANTLE_OPUS = "bedrock_mantle/anthropic.claude-opus-4-7"
MANTLE_OPUS_MAX_OUTPUT = litellm.get_model_info(model=MANTLE_OPUS)["max_output_tokens"]
UNMAPPED_MODEL = "openai/litellm-io-unmapped-model"
UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT = 4096
AZURE_BASE_MODEL = "azure/gpt-5"
AZURE_BASE_MODEL_MAX_OUTPUT = litellm.get_model_info(model=AZURE_BASE_MODEL)["max_output_tokens"]
class TestIOTokenRateLimitHelpers:
@ -48,13 +62,23 @@ class TestIOTokenRateLimitHelpers:
# A real non-zero estimate is reserved as-is.
assert _reservation_value(42, 100) == 42
def test_resolve_max_tokens_respects_explicit_zero(self):
deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}}
# An explicit max_tokens=0 is honored, not replaced by the model default.
assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0
# max_completion_tokens is the fallback only when max_tokens is absent.
assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12
assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9
@pytest.mark.parametrize(
"request_kwargs, expected",
[
({"max_tokens": 0}, 0),
({"max_completion_tokens": 12}, 12),
({"max_output_tokens": 9}, 9),
({"max_tokens": 5, "max_completion_tokens": 12}, 5),
({"max_tokens": None, "max_output_tokens": 9}, 9),
({"max_tokens": [100]}, None),
({"messages": []}, None),
(None, None),
],
)
def test_explicit_max_tokens_reads_every_output_cap_key(
self, request_kwargs: dict[str, object] | None, expected: int | None
):
assert _explicit_max_tokens(request_kwargs) == expected
def test_build_io_token_rate_limit_headers(self):
headers = build_io_token_rate_limit_headers(
@ -386,6 +410,244 @@ class TestModelRateLimitingCheckIOTokens:
normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY]
assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
@pytest.mark.parametrize(
"model, otpm_limit",
[
(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT // 2),
(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT - 1),
(UNMAPPED_MODEL, 1000),
(UNMAPPED_MODEL, 3),
],
)
async def test_request_without_max_tokens_is_admitted_below_model_max_output(
self, model: str, otpm_limit: int, use_sync_check: bool
):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": model, "otpm": otpm_limit},
"model_info": {"id": f"io-no-max-tokens-{otpm_limit}"},
"model_name": "no-max-tokens",
}
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
reserved: Final = request_kwargs["metadata"][OTPM_RESERVED_KEY]
assert isinstance(reserved, int)
assert 0 < reserved < otpm_limit
assert await dual_cache.async_get_cache(key=request_kwargs["metadata"][OTPM_CACHE_KEY]) == reserved
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model, otpm_limit",
[(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT), (UNMAPPED_MODEL, UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT)],
)
async def test_concurrent_requests_without_max_tokens_share_otpm_without_overshoot(
self, model: str, otpm_limit: int
):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": model, "otpm": otpm_limit},
"model_info": {"id": "io-no-max-tokens-race-id"},
"model_name": "no-max-tokens",
}
async def _attempt() -> dict[str, object] | None:
request_kwargs: Final = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
try:
await check.async_pre_call_check(deployment)
except litellm.RateLimitError:
return None
return request_kwargs["metadata"]
admitted: Final = [m for m in await asyncio.gather(*[_attempt() for _ in range(8)]) if m is not None]
otpm_key: Final = admitted[0][OTPM_CACHE_KEY]
assert len(admitted) == round(1 / NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)
assert {m[OTPM_CACHE_KEY] for m in admitted} == {otpm_key}
assert await dual_cache.async_get_cache(key=otpm_key) == sum(m[OTPM_RESERVED_KEY] for m in admitted)
assert await dual_cache.async_get_cache(key=otpm_key) <= otpm_limit
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model, model_max_output",
[(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT), (UNMAPPED_MODEL, UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT)],
)
async def test_request_without_max_tokens_reserves_model_max_output_when_otpm_allows(
self, model: str, model_max_output: int
):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": model, "otpm": model_max_output * 10},
"model_info": {"id": "io-no-max-tokens-roomy-id"},
"model_name": "roomy",
}
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
await check.async_pre_call_check(deployment)
assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == model_max_output
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
async def test_base_model_does_not_change_the_default_otpm_reservation(self, use_sync_check: bool):
check: Final = ModelRateLimitingCheck(dual_cache=DualCache())
otpm_limit: Final = AZURE_BASE_MODEL_MAX_OUTPUT * 10
async def _reserved(model_info: dict[str, str]) -> int:
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
deployment: Final = {
"litellm_params": {"model": "azure/litellm-io-custom-deployment", "otpm": otpm_limit},
"model_info": model_info,
"model_name": "custom-name",
}
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
return request_kwargs["metadata"][OTPM_RESERVED_KEY]
with_base_model: Final = await _reserved({"id": "io-with-base-model", "base_model": AZURE_BASE_MODEL})
without_base_model: Final = await _reserved({"id": "io-without-base-model"})
assert with_base_model == without_base_model
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
async def test_unparseable_max_tokens_leaves_itpm_uncharged(self, use_sync_check: bool):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": MANTLE_OPUS, "itpm": 2000, "otpm": 2000},
"model_info": {"id": "io-bad-max-tokens-id"},
"model_name": "opus",
}
async def _pre_call(request_kwargs: dict[str, object]) -> None:
set_io_token_rate_limit_request_kwargs(request_kwargs)
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
await _pre_call({"messages": [{"role": "user", "content": "hello there"}], "max_tokens": "abc", "metadata": {}})
valid_kwargs: Final = {
"messages": [{"role": "user", "content": "hello there"}],
"max_tokens": 1,
"metadata": {},
}
await _pre_call(valid_kwargs)
itpm_counter: Final = await dual_cache.async_get_cache(key=valid_kwargs["metadata"][ITPM_CACHE_KEY])
assert itpm_counter == valid_kwargs["metadata"][ITPM_RESERVED_KEY]
@pytest.mark.asyncio
async def test_explicit_max_tokens_above_default_share_is_reserved_in_full(self):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": MANTLE_OPUS, "otpm": 100},
"model_info": {"id": "io-explicit-max-tokens-id"},
"model_name": "opus",
}
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
await check.async_pre_call_check(deployment)
assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == 100
set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": "hi"}], "metadata": {}})
with pytest.raises(litellm.RateLimitError):
await check.async_pre_call_check(deployment)
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
async def test_deployment_max_tokens_sets_the_reservation_unless_the_request_sets_one(self, use_sync_check: bool):
check: Final = ModelRateLimitingCheck(dual_cache=DualCache())
deployment: Final = {
"litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, "max_tokens": 100},
"model_info": {"id": "io-deployment-max-tokens-id"},
"model_name": "opus",
}
without_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
with_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 300, "metadata": {}}
for request_kwargs in (without_cap, with_cap):
set_io_token_rate_limit_request_kwargs(request_kwargs)
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
assert without_cap["metadata"][OTPM_RESERVED_KEY] == 100
assert with_cap["metadata"][OTPM_RESERVED_KEY] == 300
@pytest.mark.asyncio
async def test_deployment_cap_sets_the_reservation_only_through_the_param_its_api_sends(self):
check: Final = ModelRateLimitingCheck(dual_cache=DualCache())
chat: Final = {"messages": [{"role": "user", "content": "hello"}]}
responses: Final = {"input": "hello"}
async def reserved(deployment_cap: Mapping[str, int], request_shape: Mapping[str, object]) -> int:
request_kwargs: Final = {**request_shape, "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
await check.async_pre_call_check(
{
"litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, **deployment_cap},
"model_info": {"id": f"io-sent-cap-{sorted(deployment_cap)}-{sorted(request_shape)}"},
"model_name": "opus",
}
)
return request_kwargs["metadata"][OTPM_RESERVED_KEY]
chat_default: Final = await reserved({}, chat)
responses_default: Final = await reserved({}, responses)
assert chat_default == responses_default > 100
assert await reserved({"max_tokens": 100}, chat) == 100
assert await reserved({"max_completion_tokens": 100}, chat) == 100
assert await reserved({"max_output_tokens": 100}, chat) == chat_default
assert await reserved({"max_output_tokens": 100}, responses) == 100
assert await reserved({"max_tokens": 100}, responses) == responses_default
assert await reserved({"max_completion_tokens": 100}, responses) == responses_default
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
async def test_deployment_cap_above_the_default_reserves_the_default(self, use_sync_check: bool):
check: Final = ModelRateLimitingCheck(dual_cache=DualCache())
otpm_limit: Final = UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT * 5
async def reserved(deployment_cap: Mapping[str, int]) -> int:
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
deployment: Final = {
"litellm_params": {"model": UNMAPPED_MODEL, "otpm": otpm_limit, **deployment_cap},
"model_info": {"id": f"io-deployment-cap-above-default-{deployment_cap}"},
"model_name": "unmapped",
}
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
return request_kwargs["metadata"][OTPM_RESERVED_KEY]
default: Final = await reserved({})
assert await reserved({"max_tokens": default + 1}) == default
assert await reserved({"max_tokens": otpm_limit * 2}) == default
def test_sync_io_pre_call_reserves_and_reconciles(self):
from litellm.utils import get_utc_datetime
@ -415,9 +677,9 @@ class TestModelRateLimitingCheckIOTokens:
"standard_logging_object": {
"model_id": "io-sync-id",
"hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"},
"metadata": dict(request_kwargs["metadata"]),
"metadata": {},
},
"metadata": request_kwargs["metadata"],
"litellm_params": {"metadata": request_kwargs["metadata"]},
}
response = ModelResponse(
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
@ -709,6 +971,32 @@ class TestModelRateLimitingCheckIOTokens:
assert await dual_cache.async_get_cache(key=itpm_key) == 4
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_reconcile", [False, True])
async def test_reconcile_charges_usage_above_the_reservation_to_the_minute_it_finished_in(
self, use_sync_reconcile: bool
):
dual_cache: Final = DualCache()
counter: Final = "global_router:io-late-id:bedrock_mantle/test:otpm"
reserved_minute_key: Final = f"{counter}:99-99"
await dual_cache.async_increment_cache(key=reserved_minute_key, value=10, ttl=60)
kwargs: Final = {"metadata": {ITPM_RESERVED_KEY: 0, OTPM_RESERVED_KEY: 10, OTPM_CACHE_KEY: reserved_minute_key}}
response: Final = ModelResponse(
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
usage=Usage(prompt_tokens=4, completion_tokens=25, total_tokens=29),
)
minute_before: Final = get_utc_datetime().strftime("%H-%M")
if use_sync_reconcile:
io_token_reconcile_success(dual_cache, kwargs, response)
else:
await async_io_token_reconcile_success(dual_cache, kwargs, response)
finished_minutes: Final = {minute_before, get_utc_datetime().strftime("%H-%M")}
charged: Final = [await dual_cache.async_get_cache(key=f"{counter}:{minute}") for minute in finished_minutes]
assert await dual_cache.async_get_cache(key=reserved_minute_key) == 10
assert sum(value or 0 for value in charged) == 25 - 10
@pytest.mark.asyncio
async def test_reconcile_missing_usage_keeps_reservation(self):
dual_cache = DualCache()
@ -799,16 +1087,16 @@ class TestModelRateLimitingCheckIOTokens:
assert await dual_cache.async_get_cache(key=itpm_key) == 4
def test_sync_reconcile_anthropic_dict_usage(self):
def test_sync_reconcile_anthropic_dict_usage_counts_uncached_input_once(self):
dual_cache = DualCache()
itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00"
otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00"
dual_cache.set_cache(key=itpm_key, value=6, ttl=60)
dual_cache.set_cache(key=itpm_key, value=10, ttl=60)
dual_cache.set_cache(key=otpm_key, value=4, ttl=60)
kwargs = {
"metadata": {
ITPM_RESERVED_KEY: 6,
ITPM_RESERVED_KEY: 10,
OTPM_RESERVED_KEY: 4,
ITPM_CACHE_KEY: itpm_key,
OTPM_CACHE_KEY: otpm_key,
@ -816,12 +1104,17 @@ class TestModelRateLimitingCheckIOTokens:
}
response = {
"type": "message",
"usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1},
"usage": {
"input_tokens": 3,
"output_tokens": 2,
"cache_read_input_tokens": 1,
"cache_creation_input_tokens": 4,
},
}
io_token_reconcile_success(dual_cache, kwargs, response)
assert dual_cache.get_cache(key=itpm_key) == 2
assert dual_cache.get_cache(key=itpm_key) == 3 + 4
assert dual_cache.get_cache(key=otpm_key) == 2
@pytest.mark.asyncio
@ -958,7 +1251,214 @@ class TestModelRateLimitingCheckIOTokens:
assert current == 0
ROUTER_IO_MODEL: Final = "anthropic/claude-haiku-5-5"
RouterCall = Callable[[Router, str | Exception], Awaitable[tuple[int, int]]]
def _router_with_io_limits() -> Router:
return Router(
model_list=[
{
"model_name": "limited",
"litellm_params": {"model": ROUTER_IO_MODEL, "api_key": "sk-test", "itpm": 50000, "otpm": 50000},
}
],
optional_pre_call_checks=["enforce_model_rate_limits"],
num_retries=1,
)
def _current_minute() -> str:
return get_utc_datetime().strftime("%H-%M")
def _counter(router: Router, template: RouterCacheEnum, minutes: frozenset[str]) -> int:
deployment: Final = router.model_list[0]
return sum(
router.cache.get_cache(
key=template.value.format(
id=deployment["model_info"]["id"],
model=deployment["litellm_params"]["model"],
current_minute=minute,
)
)
or 0
for minute in minutes
)
async def _io_counters_since(router: Router, first_minute: str) -> tuple[int, int]:
await asyncio.sleep(0)
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
minutes: Final = frozenset({first_minute, _current_minute()})
return _counter(router, RouterCacheEnum.ITPM, minutes), _counter(router, RouterCacheEnum.OTPM, minutes)
async def _sdk_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]:
response: Final = await router.acompletion(
model="limited",
messages=[{"role": "user", "content": "hello there"}],
max_tokens=100,
mock_response=mock_response,
)
return response.usage.prompt_tokens, response.usage.completion_tokens
async def _proxy_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]:
response: Final = await router.acompletion(
model="limited",
messages=[{"role": "user", "content": "hello there"}],
max_tokens=100,
metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_response,
)
return response.usage.prompt_tokens, response.usage.completion_tokens
async def _proxy_responses(router: Router, mock_response: str | Exception) -> tuple[int, int]:
response: Final = await router.aresponses(
model="limited",
input="hello there",
max_output_tokens=100,
litellm_metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_response,
)
return response.usage.input_tokens, response.usage.output_tokens
async def _proxy_messages(router: Router, mock_response: str | Exception) -> tuple[int, int]:
response: Final = await router.aanthropic_messages(
model="limited",
messages=[{"role": "user", "content": "hello there"}],
max_tokens=100,
litellm_metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_response,
)
return response["usage"]["input_tokens"], response["usage"]["output_tokens"]
ProxyCallSharingLogging = Callable[[Router, str | Exception, LiteLLMLoggingObject], Awaitable[object]]
PROXY_MESSAGES: Final = [{"role": "user", "content": "hello there"}]
def _logging_obj_shared_across_attempts(original_function: str) -> LiteLLMLoggingObject:
logging_obj, _ = litellm.utils.function_setup(
original_function,
litellm.utils.Rules(),
datetime.now(),
model="limited",
messages=PROXY_MESSAGES,
litellm_call_id="proxy-call",
)
return logging_obj
async def _proxy_chat_sharing_logging(
router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject
) -> object:
return await router.acompletion(
model="limited",
messages=PROXY_MESSAGES,
max_tokens=100,
metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_error,
litellm_call_id="proxy-call",
litellm_logging_obj=logging_obj,
)
async def _proxy_responses_sharing_logging(
router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject
) -> object:
return await router.aresponses(
model="limited",
input="hello there",
max_output_tokens=100,
litellm_metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_error,
litellm_call_id="proxy-call",
litellm_logging_obj=logging_obj,
)
async def _proxy_messages_sharing_logging(
router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject
) -> object:
return await router.aanthropic_messages(
model="limited",
messages=PROXY_MESSAGES,
max_tokens=100,
litellm_metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_error,
litellm_call_id="proxy-call",
litellm_logging_obj=logging_obj,
)
class TestRouterIOTokenIntegration:
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call",
[_sdk_chat, _proxy_chat, _proxy_responses, _proxy_messages],
ids=["sdk_chat", "proxy_chat_metadata", "proxy_responses_litellm_metadata", "proxy_messages_litellm_metadata"],
)
async def test_success_reconciles_io_counters_to_real_usage(self, call: RouterCall):
router: Final = _router_with_io_limits()
first_minute: Final = _current_minute()
usage: Final = await call(router, "hi")
assert await _io_counters_since(router, first_minute) == usage
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call", "mock_error"),
[
(
_proxy_responses,
litellm.InternalServerError(message="mock 500", llm_provider="anthropic", model=ROUTER_IO_MODEL),
),
(_proxy_messages, "litellm.InternalServerError"),
],
ids=["proxy_responses_litellm_metadata", "proxy_messages_litellm_metadata"],
)
async def test_retried_failure_refunds_each_io_reservation_once(
self, call: RouterCall, mock_error: str | Exception
):
router: Final = _router_with_io_limits()
first_minute: Final = _current_minute()
with pytest.raises(litellm.InternalServerError):
await call(router, mock_error)
assert await _io_counters_since(router, first_minute) == (0, 0)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call", "original_function", "mock_error"),
[
(_proxy_chat_sharing_logging, "acompletion", "litellm.InternalServerError"),
(
_proxy_responses_sharing_logging,
"aresponses",
litellm.InternalServerError(message="mock 500", llm_provider="anthropic", model=ROUTER_IO_MODEL),
),
(_proxy_messages_sharing_logging, "anthropic_messages", "litellm.InternalServerError"),
],
ids=["proxy_chat", "proxy_responses", "proxy_messages"],
)
async def test_retried_failure_sharing_one_logging_object_refunds_every_attempt(
self, call: ProxyCallSharingLogging, original_function: str, mock_error: str | Exception
):
router: Final = _router_with_io_limits()
first_minute: Final = _current_minute()
with pytest.raises(litellm.InternalServerError):
await call(router, mock_error, _logging_obj_shared_across_attempts(original_function))
assert await _io_counters_since(router, first_minute) == (0, 0)
@pytest.mark.asyncio
async def test_model_group_info_aggregates_io_limits(self):
router = Router(