diff --git a/litellm/constants.py b/litellm/constants.py index e07d727973d..cfee777f36a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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) diff --git a/litellm/litellm_core_utils/env_utils.py b/litellm/litellm_core_utils/env_utils.py index 6883af98bbd..462b0b3df47 100644 --- a/litellm/litellm_core_utils/env_utils.py +++ b/litellm/litellm_core_utils/env_utils.py @@ -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( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 404aff84994..14b08a09116 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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 diff --git a/litellm/router.py b/litellm/router.py index 3a423b79e49..20af2e5649e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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 diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 0b3d48c9e1d..7580c7dadd3 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -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)", diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index f4b41211c6e..f41d19cd62c 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -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 diff --git a/tests/integration/routing/test_otpm_default_reservation_fraction.py b/tests/integration/routing/test_otpm_default_reservation_fraction.py new file mode 100644 index 00000000000..589a2535a03 --- /dev/null +++ b/tests/integration/routing/test_otpm_default_reservation_fraction.py @@ -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) diff --git a/tests/unit/litellm_core_utils/test_env_utils.py b/tests/unit/litellm_core_utils/test_env_utils.py index 1ddd172e274..89b091e47c2 100644 --- a/tests/unit/litellm_core_utils/test_env_utils.py +++ b/tests/unit/litellm_core_utils/test_env_utils.py @@ -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 diff --git a/tests/unit/test_constants.py b/tests/unit/test_constants.py index 99b049f1803..d9c34e95364 100644 --- a/tests/unit/test_constants.py +++ b/tests/unit/test_constants.py @@ -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", + } ) diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index 05921aa70b3..e5ff4531f84 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -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(