mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
feat(router): add separate ITPM/OTPM deployment rate limits (#31952)
Some checks are pending
GitHub Actions Security Analysis / zizmor (push) Waiting to run
Some checks are pending
GitHub Actions Security Analysis / zizmor (push) Waiting to run
* feat(router): add separate ITPM/OTPM deployment rate limits Support input/output tokens per minute on deployments via enforce_model_rate_limits, with reservation, reconciliation, refund on failure, and rate-limit headers. Co-authored-by: Cursor <cursoragent@cursor.com> * chore(router): keep ITPM/OTPM diff minimal in router.py Drop unrelated Black reformatting from router.py and types/router.py so the PR only contains functional ITPM/OTPM changes. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(router): make ITPM/OTPM limits separate and atomic Address Greptile review on separate ITPM/OTPM deployment rate limits. - OTPM is now reserved atomically pre-call with rollback, matching the ITPM path, so concurrent requests can no longer overshoot the configured output limit before reconciliation - ITPM counts input tokens only; it no longer accumulates completion tokens, so the input-token limit and x-ratelimit-limit-input-tokens header describe input usage as their names imply - _read_reservation_from_kwargs only falls back to litellm_params.metadata when the top-level metadata channel is absent, so production requests carrying a litellm_params.metadata dict still reconcile and refund their reservation Adds regression tests for OTPM atomicity under concurrency, input-only ITPM enforcement, and reservation lookup when litellm_params.metadata is present. * fix(router): subtract input tokens only from remaining-input-tokens header The in-flight replay for x-ratelimit-remaining-input-tokens subtracted total tokens (input + output) instead of input tokens only, so clients saw remaining input quota understated by the completion token count on every response. Now consistent with the input-only ITPM counter. * fix(router): make itpm/otpm vs tpm/rpm precedence explicit When a deployment configures itpm/otpm alongside tpm/rpm, the io-token path takes over and the tpm/rpm limits are not enforced. Log a warning the first time such a conflicting deployment is seen so the supersession is not silent, and document the mutual exclusivity. Post-call reconciliation now only trues up a counter that was actually reserved against, so the itpm/otpm keys are no longer incremented for deployments that never configured that limit. * fix(router): track actual io-token usage on the reservation-minute key Post-call reconciliation now keys off the exact cache key stashed at pre-call time rather than one recomputed from the response-time minute. This fixes two issues: a request whose pre-call estimate was 0 now still writes its actual billable input to the ITPM counter (previously it was skipped, leaving the limit unenforceable for that request), and a call that finishes in a later minute reconciles against the minute it reserved against instead of pushing a negative delta into the next minute. Counters are only touched when their limit is configured. * fix(router): run io-token reconciliation before the model_id guard async_log_success_event gated IO reconciliation behind the model_id guard that only the TPM tracking path needs. Since reconciliation works entirely from the cache keys stashed in kwargs, a success event whose standard_logging_object lacks model_id would skip reconciliation and leave the reservation on the counter until the TTL expired, wasting quota. Route the IO path first. * fix(router): don't replay in-flight delta for itpm/otpm headers For ITPM/OTPM model groups the counter is incremented at reservation time (pre-call), so the remaining values returned by get_remaining_model_group_usage already account for the current request. Replaying the in-flight delta on top double-counted it and understated x-ratelimit-remaining-input/output-tokens by up to max_tokens on every response. Skip the delta for io-token groups; the legacy TPM/RPM replay path is unchanged. * fix(router): clear io-token reservation after reconcile/refund async_io_token_refund_failure and async_io_token_reconcile_success now clear the stashed reservation keys from the request metadata once done. Otherwise, on a model group mixing IO-limited and non-IO deployments, a failed IO call that retries on a non-IO fallback left the stale sentinel in the shared request metadata; the fallback's success handler would divert into IO reconciliation against the already-refunded key, driving the ITPM counter negative and skipping the non-IO deployment's TPM tracking. * fix(router): tidy reservation channel lookup and header guard Consolidate the reservation channel lookup into a single ordered helper shared by read and clear, so top-level metadata always wins over litellm_params metadata without the tangled per-iteration fallback. Also stop gating the router rate-limit header block on the presence of x-ratelimit-remaining-input/output-tokens. That block only emits those headers for ITPM/OTPM groups; for a non-IO group backed by a provider that natively returns input/output token headers, the extra conditions suppressed the router's own remaining-tokens/requests headers. * fix(router): strip client-supplied io-token reservation keys The reservation sentinels (_litellm_itpm_reserved, _litellm_itpm_cache_key, and the otpm equivalents) are server-only, but metadata is caller-controlled on proxy requests. An authenticated caller could forge these fields with an arbitrary cache key so the post-call reconcile/refund path would decrement any deployment's ITPM/OTPM counter and let it exceed the configured limit. Strip the reserved keys from the request metadata in set_io_token_rate_limit_request_kwargs, which runs before the router stashes its own reservation, so only a genuine server-side reservation is ever read post-call. * fix(router): track TPM routing load for io-limited deployments deployment_callback_on_success early-returned for any deployment with itpm/otpm set, so its total-token usage never landed in the router's TPM routing counter. TPM-aware routing strategies then saw 0 load for IO deployments and over-routed to them in mixed model groups. Only skip tracking when neither tpm/rpm nor itpm/otpm are configured; itpm/otpm enforcement still runs separately in ModelRateLimitingCheck, so the routing counter and the enforcement counters stay independent. * fix(router): expose standard tpm/rpm headers for io-limited groups get_remaining_model_group_usage returned early for ITPM/OTPM groups, so a group that also set tpm/rpm never emitted x-ratelimit-remaining-tokens / -requests; clients and prometheus gauges reading those saw no data. Build both header sets instead of returning early. Also simplify the in-flight header replay: only the tpm/rpm counters are incremented post-response, so the delta now adjusts just those. The itpm/otpm counters are incremented at reservation time (pre-call), so the input/output token headers already reflect the request and are left untouched - which removes the need for the separate io-group special case. * fix(router): roll back ITPM on any OTPM reservation error; dedup warning per instance Two follow-ups from review. The pre-call OTPM reservation only rolled back the ITPM reservation on a RateLimitError, so a transient cache error while reserving OTPM left the ITPM counter inflated until the TTL expired; catch any exception, release the ITPM reservation, then re-raise. Replace the module-level lru_cache warn-once (caching a logging side effect, which never re-warns in a long-lived process) with an instance-scoped set of already-warned deployment ids on ModelRateLimitingCheck. * fix(router): always clear reservation stash on reconcile; don't collapse id-less warning dedup Clear the reservation in a finally block so a mid-reconciliation cache error still removes the stash and a duplicate success event can't re-process it. Dedup the itpm/otpm-vs-tpm/rpm conflict warning per real deployment id; a deployment with no id no longer collapses every id-less deployment onto the str(None) key (which would suppress all but the first warning). * fix(router): skip io reservation when deployment can't be keyed _get_cache_keys returned a shared 'global_router:None:None:...' key when a deployment was missing model_info.id or litellm_params.model, so misconfigured deployments could share one rate-limit bucket. Return None in that case and skip io reservation for the request. * fix(router): honor explicit max_tokens=0 in io reservation _resolve_max_tokens used 'max_tokens or max_completion_tokens', so an explicit max_tokens=0 fell through to the model default. Only fall back to max_completion_tokens when max_tokens is absent. * fix(ci): satisfy lint budget, router coverage, and dashboard schema sync - Modernize the new itpm/otpm module's type hints to PEP 585 lowercase generics (Dict/Tuple/List -> dict/tuple/list) to clear the added UP006 violations; ratchet ruff-strict-budget.json's UP006 ceiling down to match. - Replace three try/except Exception blocks that must stay broad by design (token_counter and litellm.get_model_info raise untyped exceptions, and an io-token refund failure must never break the logging pipeline) with contextlib.suppress(Exception), matching the codebase's existing resolution for this exact BLE001 pattern. - Add direct unit tests for get_model_group_io_token_usage (multi-deployment aggregation and the empty-model-list case) in test_router_helper_utils.py, satisfying the router function-coverage check. - Regenerate the dashboard's schema.d.ts so the new itpm/otpm fields on GenericLiteLLMParams and ModelGroupInfo are reflected in the OpenAPI types. * fix: enforce io token rate limits consistently * fix: honor zero max tokens in otpm reservation * fix(lint): fix UP007 violation and resync ruff-strict-budget.json to base Convert Union[_Span, Any] to _Span | Any (safe on this repo's Python >=3.10 floor) to clear the new UP007 violation from the TYPE_CHECKING-gated Span alias. The previously committed ruff-strict-budget.json ratcheted UP006 down from a stale base; litellm_internal_staging has since tightened that same ceiling further on its own. Reset the file to the current base's committed values and re-ratchet from there so the budget only ever moves down relative to the actual merge-base, never against a stale snapshot. * fix(router): attach ITPM/OTPM headers on dict responses and harden reservation Strip itpm/otpm from provider kwargs, ensure messages are available for ITPM estimation, honor max_output_tokens on /v1/responses, and propagate rate-limit headers through /v1/messages dict responses via _hidden_params. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(router): attach ITPM/OTPM headers to streaming /v1/messages responses Wrap bare async iterators in HiddenParamsAsyncIteratorWrapper so set_response_headers can attach rate-limit headers to streaming Anthropic messages responses that lack a _hidden_params slot. Co-authored-by: Cursor <cursoragent@cursor.com> * style: ruff format add_retry_fallback_headers.py Fix CI ruff format check failure on get_hidden_params_dict call site. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(router): extract set_response_headers helpers to fix C901 budget Move header-attachment logic into add_retry_fallback_headers helpers so set_response_headers stays under the strict complexity ceiling. Co-authored-by: Cursor <cursoragent@cursor.com> * fix: keep IO token reservation when response usage is missing Missing usage was reconciled as zero and fully refunded the pre-call reservation, allowing limit bypass on repeated successful calls. Only adjust counters when usage is resolved from the response or standard logging fields; otherwise keep the reservation until TTL expires. Co-authored-by: Cursor <cursoragent@cursor.com> * fix: enforce RPM/TPM alongside IO-token limits on mixed deployments Deployments with both itpm/otpm and tpm/rpm previously returned after the IO reservation and skipped RPM/TPM checks. Run both paths and refund the IO reservation only when RPM/TPM rejects after a successful reservation. Co-authored-by: Cursor <cursoragent@cursor.com> * fix: track TPM usage on success for mixed IO+TPM deployments The early return after IO-token reconciliation in log_success_event and async_log_success_event skipped the TPM counter increment, so the tpm_key the pre-call check reads was never written and tpm_limit was never actually enforced on deployments that also configure itpm/otpm. Co-authored-by: Cursor <cursoragent@cursor.com> * fix: treat total-only usage as unresolved in IO-token reconcile usage/standard_logging_object entries carrying only total_tokens (no prompt/completion or input/output breakdown) were treated as resolved usage, resolving to (0, 0) and refunding the full reservation. Both _usage_is_present and the standard_logging_object fallback now require an actual input/output breakdown before reconciling, keeping the reservation otherwise. Co-authored-by: Cursor <cursoragent@cursor.com> * fix: reserve minimal token when input/output estimation fails _reservation_value(0, limit) reserved the entire limit whenever token estimation failed (empty/unsupported input, tokenizer error), letting one such request claim the whole bucket and 429 every concurrent request to the deployment until it completed. Reserve 1 token instead so estimation failures no longer serialize traffic. Co-authored-by: Cursor <cursoragent@cursor.com> * fix: refund IO reservation synchronously before retry deployment pick On retry, set_io_token_rate_limit_request_kwargs clears reservation sentinels from the shared kwargs dict before a background failure handler can refund them, stranding the counter until TTL. Refund and clear any stale reservation in _update_kwargs_with_deployment before stripping sentinels for the next attempt. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(io_token_rate_limit_check): use model-specific tokenizer for ITPM estimate; document sync-refund Redis ceiling Pass the deployment litellm_params.model to token_counter so it uses the model's native tokenizer instead of the generic fallback, narrowing the reservation over/under-estimate window between pre-call and post-call reconcile. Add a ponytail: comment to refund_stale_reservation_before_retry explaining the known ceiling: the synchronous DualCache.increment_cache issues a blocking Redis INCR when a Redis backend is configured. This only fires on streaming mid-stream retries (non-streaming failures await their failure handler before the retry picks a new deployment, leaving no sentinels to refund). Upgrade path: make _update_kwargs_with_deployment async. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Sameer Kankute
GitHub
parent
2e076b110f
commit
5b93ba0ada
14 changed files with 2533 additions and 100 deletions
|
|
@ -36,6 +36,8 @@ OPTIONAL_KWARGS_KEYS = frozenset(
|
|||
"aws_bedrock_project_id",
|
||||
"tpm",
|
||||
"rpm",
|
||||
"itpm",
|
||||
"otpm",
|
||||
"use_xai_oauth",
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -52,6 +52,7 @@ from litellm.proxy.dd_span_tagger import DDSpanTagger
|
|||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.router import RouterRateLimitError
|
||||
from litellm.types.utils import ServerToolUse
|
||||
|
|
@ -600,7 +601,7 @@ def _override_openai_response_model(
|
|||
if not requested_model:
|
||||
return
|
||||
|
||||
hidden_params = getattr(response_obj, "_hidden_params", {}) or {}
|
||||
hidden_params = get_hidden_params_dict(response_obj)
|
||||
if isinstance(hidden_params, dict):
|
||||
# Check if a fallback occurred - if so, preserve the actual model used
|
||||
fallback_headers = hidden_params.get("additional_headers", {}) or {}
|
||||
|
|
@ -896,7 +897,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
(e.g. Google native :generateContent) instead of base_process_llm_request.
|
||||
"""
|
||||
if isinstance(response, dict):
|
||||
hidden_params = response.get("_hidden_params") or {}
|
||||
hidden_params = get_hidden_params_dict(response)
|
||||
else:
|
||||
hidden_params = getattr(response, "_hidden_params", None) or {}
|
||||
if not isinstance(hidden_params, dict):
|
||||
|
|
@ -1429,7 +1430,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
_exception_raised = False
|
||||
try:
|
||||
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
||||
hidden_params = get_hidden_params_dict(response)
|
||||
model_id = self._get_model_id_from_response(hidden_params, self.data)
|
||||
|
||||
cache_key, api_base, response_cost = (
|
||||
|
|
@ -1704,7 +1705,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
log_context=f"litellm_call_id={logging_obj.litellm_call_id}",
|
||||
)
|
||||
|
||||
hidden_params = getattr(response, "_hidden_params", {}) or {} # get any updated response headers
|
||||
hidden_params = get_hidden_params_dict(response) # get any updated response headers
|
||||
additional_headers = hidden_params.get("additional_headers", {}) or {}
|
||||
|
||||
recover_response_cost = not response_cost and hidden_params.get("response_cost") is None
|
||||
|
|
@ -1732,6 +1733,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
)
|
||||
)
|
||||
|
||||
if isinstance(response, dict):
|
||||
response.pop("_hidden_params", None)
|
||||
|
||||
# Call response headers hook for non-streaming success
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=self.data,
|
||||
|
|
|
|||
|
|
@ -90,7 +90,12 @@ from litellm.router_utils.add_retry_fallback_headers import (
|
|||
_HiddenParamsHost,
|
||||
add_fallback_headers_to_response,
|
||||
add_retry_headers_to_response,
|
||||
apply_quality_router_decision_headers,
|
||||
apply_remaining_usage_headers,
|
||||
ensure_response_additional_headers,
|
||||
get_hidden_params_dict,
|
||||
prepare_response_for_header_attachment,
|
||||
response_in_flight_token_count,
|
||||
)
|
||||
from litellm.router_utils.batch_utils import (
|
||||
_get_router_metadata_variable_name,
|
||||
|
|
@ -133,6 +138,12 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
|||
from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
|
||||
ModelRateLimitingCheck,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
|
||||
build_io_token_rate_limit_headers,
|
||||
deployment_has_io_token_limits,
|
||||
refund_stale_reservation_before_retry,
|
||||
set_io_token_rate_limit_request_kwargs,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import (
|
||||
PromptCachingDeploymentCheck,
|
||||
)
|
||||
|
|
@ -1649,6 +1660,7 @@ class Router:
|
|||
)
|
||||
thread.start()
|
||||
|
||||
kwargs.setdefault("messages", messages)
|
||||
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||||
kwargs.pop("silent_model", None) # Ensure it's not in kwargs either
|
||||
model_name = litellm_params["model"]
|
||||
|
|
@ -2672,6 +2684,7 @@ class Router:
|
|||
)
|
||||
)
|
||||
|
||||
kwargs.setdefault("messages", messages)
|
||||
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||||
kwargs.pop("silent_model", None) # Ensure it's not in kwargs either
|
||||
|
||||
|
|
@ -2946,6 +2959,13 @@ class Router:
|
|||
}
|
||||
)
|
||||
|
||||
# A retry/fallback reuses this same kwargs dict for the next deployment.
|
||||
# Refund and clear any reservation the previous deployment attempt left
|
||||
# here before it's wiped below, instead of relying on that attempt's
|
||||
# (possibly still-pending) failure event to do it.
|
||||
refund_stale_reservation_before_retry(self.cache, kwargs)
|
||||
set_io_token_rate_limit_request_kwargs(kwargs)
|
||||
|
||||
## DEPLOYMENT-LEVEL TAGS
|
||||
deployment_tags = deployment.get("litellm_params", {}).get("tags")
|
||||
if deployment_tags:
|
||||
|
|
@ -6753,7 +6773,13 @@ class Router:
|
|||
deployment_id=id,
|
||||
)
|
||||
|
||||
## if all are none, return - no need to track current tpm/rpm usage for models with no tpm/rpm set
|
||||
deployment_dict = deployment_info if isinstance(deployment_info, dict) else deployment_info.model_dump()
|
||||
has_io_token_limits = deployment_has_io_token_limits(deployment_dict)
|
||||
|
||||
## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are
|
||||
## set. IO deployments still record TPM/RPM usage here so TPM-aware
|
||||
## routing strategies see their real load in mixed model groups; their
|
||||
## itpm/otpm enforcement runs separately in ModelRateLimitingCheck.
|
||||
if (
|
||||
tpm is None
|
||||
and rpm is None
|
||||
|
|
@ -6761,6 +6787,7 @@ class Router:
|
|||
and rpm_litellm_params is None
|
||||
and tpm_model_info is None
|
||||
and rpm_model_info is None
|
||||
and not has_io_token_limits
|
||||
):
|
||||
return
|
||||
|
||||
|
|
@ -8597,6 +8624,8 @@ class Router:
|
|||
|
||||
total_tpm: Optional[int] = None
|
||||
total_rpm: Optional[int] = None
|
||||
total_itpm: Optional[int] = None
|
||||
total_otpm: Optional[int] = None
|
||||
configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None
|
||||
model_list = self.get_model_list(model_name=model_group)
|
||||
if model_list is None:
|
||||
|
|
@ -8637,6 +8666,18 @@ class Router:
|
|||
if _deployment_rpm is None:
|
||||
_deployment_rpm = model_info_dict.get("rpm", None) # type: ignore
|
||||
|
||||
_deployment_itpm: Optional[int] = model.get("itpm") # type: ignore
|
||||
if _deployment_itpm is None:
|
||||
_deployment_itpm = model_litellm_params.get("itpm", None) # type: ignore
|
||||
if _deployment_itpm is None:
|
||||
_deployment_itpm = model_info_dict.get("itpm", None) # type: ignore
|
||||
|
||||
_deployment_otpm: Optional[int] = model.get("otpm") # type: ignore
|
||||
if _deployment_otpm is None:
|
||||
_deployment_otpm = model_litellm_params.get("otpm", None) # type: ignore
|
||||
if _deployment_otpm is None:
|
||||
_deployment_otpm = model_info_dict.get("otpm", None) # type: ignore
|
||||
|
||||
# get model info
|
||||
try:
|
||||
model_id = model_info_dict.get("id", None)
|
||||
|
|
@ -8777,6 +8818,16 @@ class Router:
|
|||
if total_rpm is None:
|
||||
total_rpm = 0
|
||||
total_rpm += _deployment_rpm # type: ignore
|
||||
|
||||
if _deployment_itpm is not None:
|
||||
if total_itpm is None:
|
||||
total_itpm = 0
|
||||
total_itpm += _deployment_itpm # type: ignore
|
||||
|
||||
if _deployment_otpm is not None:
|
||||
if total_otpm is None:
|
||||
total_otpm = 0
|
||||
total_otpm += _deployment_otpm # type: ignore
|
||||
if model_group_info is not None:
|
||||
## UPDATE WITH TOTAL TPM/RPM FOR MODEL GROUP
|
||||
if total_tpm is not None:
|
||||
|
|
@ -8785,6 +8836,12 @@ class Router:
|
|||
if total_rpm is not None:
|
||||
model_group_info.rpm = total_rpm
|
||||
|
||||
if total_itpm is not None:
|
||||
model_group_info.itpm = total_itpm
|
||||
|
||||
if total_otpm is not None:
|
||||
model_group_info.otpm = total_otpm
|
||||
|
||||
## UPDATE WITH CONFIGURABLE CLIENTSIDE AUTH PARAMS FOR MODEL GROUP
|
||||
if configurable_clientside_auth_params is not None:
|
||||
model_group_info.configurable_clientside_auth_params = configurable_clientside_auth_params
|
||||
|
|
@ -8887,6 +8944,58 @@ class Router:
|
|||
rpm_usage += t
|
||||
return tpm_usage, rpm_usage
|
||||
|
||||
async def get_model_group_io_token_usage(self, model_group: str) -> tuple[Optional[int], Optional[int]]:
|
||||
"""
|
||||
Returns current ITPM/OTPM usage for a model group (sum across deployments).
|
||||
"""
|
||||
dt = get_utc_datetime()
|
||||
current_minute = dt.strftime("%H-%M")
|
||||
itpm_keys: list[str] = []
|
||||
otpm_keys: list[str] = []
|
||||
|
||||
model_list = self.get_model_list(model_name=model_group)
|
||||
if model_list is None:
|
||||
return None, None
|
||||
|
||||
for model in model_list:
|
||||
model_id: Optional[str] = model.get("model_info", {}).get("id") # type: ignore
|
||||
litellm_model: Optional[str] = model["litellm_params"].get("model")
|
||||
if model_id is None or litellm_model is None:
|
||||
continue
|
||||
itpm_keys.append(
|
||||
RouterCacheEnum.ITPM.value.format(
|
||||
id=model_id,
|
||||
model=litellm_model,
|
||||
current_minute=current_minute,
|
||||
)
|
||||
)
|
||||
otpm_keys.append(
|
||||
RouterCacheEnum.OTPM.value.format(
|
||||
id=model_id,
|
||||
model=litellm_model,
|
||||
current_minute=current_minute,
|
||||
)
|
||||
)
|
||||
|
||||
combined_values = await self.cache.async_batch_get_cache(keys=itpm_keys + otpm_keys)
|
||||
if combined_values is None:
|
||||
return None, None
|
||||
|
||||
itpm_values = combined_values[: len(itpm_keys)]
|
||||
otpm_values = combined_values[len(itpm_keys) :]
|
||||
|
||||
total_itpm: Optional[int] = None
|
||||
for value in itpm_values:
|
||||
if isinstance(value, int):
|
||||
total_itpm = (total_itpm or 0) + value
|
||||
|
||||
total_otpm: Optional[int] = None
|
||||
for value in otpm_values:
|
||||
if isinstance(value, int):
|
||||
total_otpm = (total_otpm or 0) + value
|
||||
|
||||
return total_itpm, total_otpm
|
||||
|
||||
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
def _cached_get_model_group_info(self, model_group: str) -> Optional[ModelGroupInfo]:
|
||||
"""
|
||||
|
|
@ -8896,25 +9005,33 @@ class Router:
|
|||
"""
|
||||
return self.get_model_group_info(model_group)
|
||||
|
||||
async def get_remaining_model_group_usage(self, model_group: str) -> Dict[str, int]:
|
||||
async def get_remaining_model_group_usage(self, model_group: str) -> dict[str, int]:
|
||||
model_group_info = self._cached_get_model_group_info(model_group)
|
||||
|
||||
if model_group_info is not None and model_group_info.tpm is not None:
|
||||
tpm_limit = model_group_info.tpm
|
||||
else:
|
||||
tpm_limit = None
|
||||
returned_dict: dict[str, int] = {}
|
||||
|
||||
if model_group_info is not None and model_group_info.rpm is not None:
|
||||
rpm_limit = model_group_info.rpm
|
||||
else:
|
||||
rpm_limit = None
|
||||
# ITPM/OTPM groups emit input/output token headers, but they may also set
|
||||
# tpm/rpm, so build both sets rather than returning early - clients and
|
||||
# prometheus gauges that read the standard headers still get data.
|
||||
if model_group_info is not None and (model_group_info.itpm is not None or model_group_info.otpm is not None):
|
||||
current_itpm, current_otpm = await self.get_model_group_io_token_usage(model_group)
|
||||
returned_dict.update(
|
||||
build_io_token_rate_limit_headers(
|
||||
itpm_limit=model_group_info.itpm,
|
||||
otpm_limit=model_group_info.otpm,
|
||||
current_itpm=current_itpm,
|
||||
current_otpm=current_otpm,
|
||||
)
|
||||
)
|
||||
|
||||
tpm_limit = model_group_info.tpm if model_group_info is not None else None
|
||||
rpm_limit = model_group_info.rpm if model_group_info is not None else None
|
||||
|
||||
if tpm_limit is None and rpm_limit is None:
|
||||
return {}
|
||||
return returned_dict
|
||||
|
||||
current_tpm, current_rpm = await self.get_model_group_usage(model_group)
|
||||
|
||||
returned_dict = {}
|
||||
if tpm_limit is not None:
|
||||
returned_dict["x-ratelimit-remaining-tokens"] = tpm_limit - (current_tpm or 0)
|
||||
returned_dict["x-ratelimit-limit-tokens"] = tpm_limit
|
||||
|
|
@ -8937,69 +9054,25 @@ class Router:
|
|||
# - if healthy_deployments > 1, return model group rate limit headers
|
||||
# - else return the model's rate limit headers
|
||||
"""
|
||||
if response is not None and hasattr(response, "_hidden_params"):
|
||||
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
||||
if hasattr(hidden_params, "model_dump"):
|
||||
hidden_params = hidden_params.model_dump()
|
||||
if not isinstance(hidden_params, dict):
|
||||
return response
|
||||
response._hidden_params = hidden_params
|
||||
response = prepare_response_for_header_attachment(response)
|
||||
if response is None:
|
||||
return response
|
||||
|
||||
additional_headers = hidden_params.get("additional_headers")
|
||||
if not isinstance(additional_headers, dict):
|
||||
additional_headers = {}
|
||||
hidden_params["additional_headers"] = additional_headers
|
||||
additional_headers["x-litellm-model-group"] = model_group
|
||||
additional_headers = ensure_response_additional_headers(response)
|
||||
additional_headers["x-litellm-model-group"] = model_group
|
||||
apply_quality_router_decision_headers(additional_headers, request_kwargs)
|
||||
|
||||
# Lift QualityRouter routing decision into response headers for
|
||||
# transparency. The decision is stashed in request_kwargs.metadata
|
||||
# by QualityRouter.async_pre_routing_hook.
|
||||
metadata = (request_kwargs.get("metadata") or {}) if isinstance(request_kwargs, dict) else {}
|
||||
decision = metadata.get("quality_router_decision") if isinstance(metadata, dict) else None
|
||||
if isinstance(decision, dict):
|
||||
# Only emit headers for fields that have a meaningful value.
|
||||
# `complexity_tier` and `matched_keyword` are mutually exclusive
|
||||
# (the keyword path short-circuits classification), so each
|
||||
# request emits one or the other but not both.
|
||||
if decision.get("routed_model") is not None:
|
||||
additional_headers["x-litellm-quality-router-model"] = str(decision["routed_model"])
|
||||
if decision.get("quality_tier") is not None:
|
||||
additional_headers["x-litellm-quality-router-tier"] = str(decision["quality_tier"])
|
||||
if decision.get("routed_via") is not None:
|
||||
additional_headers["x-litellm-quality-router-via"] = str(decision["routed_via"])
|
||||
if decision.get("matched_keyword") is not None:
|
||||
additional_headers["x-litellm-quality-router-keyword"] = str(decision["matched_keyword"])
|
||||
if decision.get("complexity_tier") is not None:
|
||||
additional_headers["x-litellm-quality-router-complexity"] = str(decision["complexity_tier"])
|
||||
|
||||
if (
|
||||
"x-ratelimit-remaining-tokens" not in additional_headers
|
||||
and "x-ratelimit-remaining-requests" not in additional_headers
|
||||
and model_group is not None
|
||||
):
|
||||
remaining_usage = await self.get_remaining_model_group_usage(model_group)
|
||||
|
||||
# get_remaining_model_group_usage reads the router's TPM/RPM
|
||||
# counter, which is incremented post-response by
|
||||
# deployment_callback_on_success. So the values returned here
|
||||
# are pre-decrement for the current request, while vendor
|
||||
# headers (OpenAI/Anthropic/Azure) are post-decrement. Replay
|
||||
# the in-flight increment so router-derived headers match
|
||||
# vendor-derived semantics — for both the HTTP response sent
|
||||
# to the client and the prometheus gauges that read these
|
||||
# headers downstream (LIT-2719).
|
||||
in_flight_tokens = 0
|
||||
usage = getattr(response, "usage", None)
|
||||
if usage is not None:
|
||||
in_flight_tokens = getattr(usage, "total_tokens", 0) or 0
|
||||
in_flight_delta = {
|
||||
"x-ratelimit-remaining-tokens": in_flight_tokens,
|
||||
"x-ratelimit-remaining-requests": 1,
|
||||
}
|
||||
|
||||
for header, value in remaining_usage.items():
|
||||
if value is not None:
|
||||
additional_headers[header] = value - in_flight_delta.get(header, 0)
|
||||
if model_group is not None:
|
||||
remaining_usage = await self.get_remaining_model_group_usage(model_group)
|
||||
# get_remaining_model_group_usage reads the router's TPM/RPM counter,
|
||||
# which is incremented post-response by deployment_callback_on_success.
|
||||
# Replay the in-flight increment for TPM/RPM only (LIT-2719); ITPM/OTPM
|
||||
# counters are incremented at reservation time and must not be adjusted.
|
||||
apply_remaining_usage_headers(
|
||||
additional_headers,
|
||||
remaining_usage,
|
||||
response_in_flight_token_count(response),
|
||||
)
|
||||
return response
|
||||
|
||||
def _build_model_name_index(self, model_list: list) -> None:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import json
|
||||
from typing import Protocol, TypedDict, cast
|
||||
from typing import Any, Protocol, TypedDict, cast
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -15,8 +15,98 @@ class _HiddenParamsHost(Protocol):
|
|||
_hidden_params: dict[str, object]
|
||||
|
||||
|
||||
def get_hidden_params_dict(response: object) -> dict[str, object]:
|
||||
hidden_params: object = cast(object, getattr(response, "_hidden_params", None))
|
||||
class HiddenParamsAsyncIteratorWrapper:
|
||||
"""
|
||||
Wraps a bare async generator/iterator (e.g. a provider's raw SSE
|
||||
streaming response) that cannot itself hold a ``_hidden_params``
|
||||
attribute, so router-derived headers (ITPM/OTPM, model-group, retry,
|
||||
fallback) can attach to a streaming response the same way they attach
|
||||
to object-based responses (e.g. ``CustomStreamWrapper``).
|
||||
"""
|
||||
|
||||
def __init__(self, inner: object) -> None:
|
||||
self._inner = inner
|
||||
self._hidden_params: dict[str, object] = {}
|
||||
|
||||
def __aiter__(self) -> "HiddenParamsAsyncIteratorWrapper":
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> object:
|
||||
return await cast(Any, self._inner).__anext__()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
aclose = getattr(self._inner, "aclose", None)
|
||||
if callable(aclose):
|
||||
await aclose()
|
||||
|
||||
|
||||
def prepare_response_for_header_attachment(response: object) -> object | None:
|
||||
if response is None:
|
||||
return None
|
||||
if isinstance(response, dict) or hasattr(response, "_hidden_params"):
|
||||
return response
|
||||
if hasattr(response, "__anext__"):
|
||||
return HiddenParamsAsyncIteratorWrapper(response)
|
||||
return response
|
||||
|
||||
|
||||
def ensure_response_additional_headers(response: object) -> dict[str, object]:
|
||||
hidden_params = get_hidden_params_dict(response, create=isinstance(response, dict))
|
||||
_write_hidden_params(response, hidden_params)
|
||||
additional_headers = hidden_params.get("additional_headers")
|
||||
if not isinstance(additional_headers, dict):
|
||||
additional_headers = {}
|
||||
hidden_params["additional_headers"] = additional_headers
|
||||
return additional_headers
|
||||
|
||||
|
||||
def apply_quality_router_decision_headers(
|
||||
additional_headers: dict[str, object],
|
||||
request_kwargs: object,
|
||||
) -> None:
|
||||
metadata = (request_kwargs.get("metadata") or {}) if isinstance(request_kwargs, dict) else {}
|
||||
decision = metadata.get("quality_router_decision") if isinstance(metadata, dict) else None
|
||||
if not isinstance(decision, dict):
|
||||
return
|
||||
quality_header_fields = (
|
||||
("routed_model", "x-litellm-quality-router-model"),
|
||||
("quality_tier", "x-litellm-quality-router-tier"),
|
||||
("routed_via", "x-litellm-quality-router-via"),
|
||||
("matched_keyword", "x-litellm-quality-router-keyword"),
|
||||
("complexity_tier", "x-litellm-quality-router-complexity"),
|
||||
)
|
||||
for field, header in quality_header_fields:
|
||||
if decision.get(field) is not None:
|
||||
additional_headers[header] = str(decision[field])
|
||||
|
||||
|
||||
def response_in_flight_token_count(response: object) -> int:
|
||||
usage = response.get("usage") if isinstance(response, dict) else getattr(response, "usage", None)
|
||||
if usage is None:
|
||||
return 0
|
||||
if isinstance(usage, dict):
|
||||
total = int(usage.get("total_tokens") or 0)
|
||||
if total:
|
||||
return total
|
||||
return int(usage.get("input_tokens") or 0) + int(usage.get("output_tokens") or 0)
|
||||
return int(getattr(usage, "total_tokens", 0) or 0)
|
||||
|
||||
|
||||
def apply_remaining_usage_headers(
|
||||
additional_headers: dict[str, object],
|
||||
remaining_usage: dict[str, int],
|
||||
in_flight_tokens: int,
|
||||
) -> None:
|
||||
in_flight_delta = {
|
||||
"x-ratelimit-remaining-tokens": in_flight_tokens,
|
||||
"x-ratelimit-remaining-requests": 1,
|
||||
}
|
||||
for header, value in remaining_usage.items():
|
||||
if value is not None and header not in additional_headers:
|
||||
additional_headers[header] = value - in_flight_delta.get(header, 0)
|
||||
|
||||
|
||||
def _normalize_hidden_params(hidden_params: object) -> dict[str, object]:
|
||||
if isinstance(hidden_params, BaseModel):
|
||||
return cast("dict[str, object]", hidden_params.model_dump())
|
||||
if isinstance(hidden_params, dict):
|
||||
|
|
@ -24,6 +114,29 @@ def get_hidden_params_dict(response: object) -> dict[str, object]:
|
|||
return {}
|
||||
|
||||
|
||||
def get_hidden_params_dict(
|
||||
response: object,
|
||||
*,
|
||||
create: bool = False,
|
||||
) -> dict[str, object]:
|
||||
if isinstance(response, dict):
|
||||
hidden_params = _normalize_hidden_params(response.get("_hidden_params"))
|
||||
if not hidden_params and create:
|
||||
hidden_params = {}
|
||||
response["_hidden_params"] = hidden_params
|
||||
return hidden_params
|
||||
|
||||
hidden_params = _normalize_hidden_params(cast(object, getattr(response, "_hidden_params", None)))
|
||||
return hidden_params
|
||||
|
||||
|
||||
def _write_hidden_params(response: object, hidden_params: dict[str, object]) -> None:
|
||||
if isinstance(response, dict):
|
||||
response["_hidden_params"] = hidden_params
|
||||
elif hasattr(response, "_hidden_params"):
|
||||
cast(_HiddenParamsHost, response)._hidden_params = hidden_params
|
||||
|
||||
|
||||
def _ensure_additional_headers_dict(
|
||||
hidden_params: dict[str, object],
|
||||
) -> dict[str, object]:
|
||||
|
|
@ -73,15 +186,19 @@ def _add_headers_to_response(response: object, headers: dict[str, object]) -> ob
|
|||
if response is None:
|
||||
return response
|
||||
|
||||
if not isinstance(response, BaseModel) and not hasattr(response, "_hidden_params"):
|
||||
if (
|
||||
not isinstance(response, BaseModel)
|
||||
and not isinstance(response, dict)
|
||||
and not hasattr(response, "_hidden_params")
|
||||
):
|
||||
return response
|
||||
|
||||
hidden_params = get_hidden_params_dict(response)
|
||||
hidden_params = get_hidden_params_dict(response, create=isinstance(response, dict))
|
||||
additional_headers = _ensure_additional_headers_dict(hidden_params)
|
||||
additional_headers.update(headers)
|
||||
hidden_params["additional_headers"] = additional_headers
|
||||
|
||||
cast(_HiddenParamsHost, response)._hidden_params = hidden_params
|
||||
_write_hidden_params(response, hidden_params)
|
||||
return response
|
||||
|
||||
|
||||
|
|
@ -127,12 +244,12 @@ def add_fallback_headers_to_response(
|
|||
if fallback_errors is None or response is None:
|
||||
return response
|
||||
|
||||
hidden_params = get_hidden_params_dict(response)
|
||||
hidden_params = get_hidden_params_dict(response, create=isinstance(response, dict))
|
||||
additional_headers = _ensure_additional_headers_dict(hidden_params)
|
||||
merged_errors = get_fallback_errors_from_headers(additional_headers) + [
|
||||
cast("dict[str, object]", error) for error in fallback_errors
|
||||
]
|
||||
additional_headers["x-litellm-fallback-errors"] = json.dumps(merged_errors)
|
||||
hidden_params["additional_headers"] = additional_headers
|
||||
cast(_HiddenParamsHost, response)._hidden_params = hidden_params
|
||||
_write_hidden_params(response, hidden_params)
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -0,0 +1,711 @@
|
|||
"""
|
||||
Separate ITPM/OTPM (input/output tokens per minute) deployment rate limits.
|
||||
|
||||
- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM
|
||||
- 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
|
||||
|
||||
Used by ModelRateLimitingCheck when a deployment sets itpm/otpm.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import contextvars
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm import token_counter
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.types.router import RouterCacheEnum, RouterErrors
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
Span = _Span | Any
|
||||
else:
|
||||
Span = Any
|
||||
|
||||
RoutingArgsTTL = 60
|
||||
|
||||
_io_token_rate_limit_request_kwargs: contextvars.ContextVar[Optional[dict[str, Any]]] = contextvars.ContextVar(
|
||||
"io_token_rate_limit_request_kwargs",
|
||||
default=None,
|
||||
)
|
||||
|
||||
ITPM_RESERVED_KEY = "_litellm_itpm_reserved"
|
||||
OTPM_RESERVED_KEY = "_litellm_otpm_reserved"
|
||||
ITPM_CACHE_KEY = "_litellm_itpm_cache_key"
|
||||
OTPM_CACHE_KEY = "_litellm_otpm_cache_key"
|
||||
|
||||
|
||||
def set_io_token_rate_limit_request_kwargs(kwargs: Optional[dict[str, Any]]) -> None:
|
||||
# The reservation sentinels are server-only, but `metadata` is caller
|
||||
# controlled on proxy requests. Strip any client-supplied copies here (this
|
||||
# runs before the router stashes its own reservation) so a forged
|
||||
# reservation can't drive the post-call reconcile/refund against an
|
||||
# arbitrary counter and bypass the configured limits.
|
||||
_clear_reservation_from_kwargs(kwargs)
|
||||
_io_token_rate_limit_request_kwargs.set(kwargs)
|
||||
|
||||
|
||||
def get_io_token_rate_limit_request_kwargs() -> Optional[dict[str, Any]]:
|
||||
return _io_token_rate_limit_request_kwargs.get()
|
||||
|
||||
|
||||
def seconds_until_minute_reset() -> int:
|
||||
dt = get_utc_datetime()
|
||||
return max(1, 60 - dt.second)
|
||||
|
||||
|
||||
def get_deployment_io_token_limits(
|
||||
deployment: dict,
|
||||
) -> tuple[Optional[int], Optional[int]]:
|
||||
itpm = deployment.get("itpm")
|
||||
otpm = deployment.get("otpm")
|
||||
litellm_params = deployment.get("litellm_params") or {}
|
||||
model_info = deployment.get("model_info") or {}
|
||||
if itpm is None:
|
||||
itpm = litellm_params.get("itpm")
|
||||
if otpm is None:
|
||||
otpm = litellm_params.get("otpm")
|
||||
if itpm is None:
|
||||
itpm = model_info.get("itpm")
|
||||
if otpm is None:
|
||||
otpm = model_info.get("otpm")
|
||||
return itpm, otpm
|
||||
|
||||
|
||||
def deployment_has_io_token_limits(deployment: dict) -> bool:
|
||||
itpm, otpm = get_deployment_io_token_limits(deployment)
|
||||
return itpm is not None or otpm is not None
|
||||
|
||||
|
||||
def _get_cache_keys(deployment: dict, current_minute: str) -> Optional[tuple[str, str]]:
|
||||
model_id = deployment.get("model_info", {}).get("id")
|
||||
deployment_name = deployment.get("litellm_params", {}).get("model")
|
||||
# Without both a deployment id and model name the key would collapse to a
|
||||
# shared "None:None" bucket across misconfigured deployments, so bail out.
|
||||
if model_id is None or deployment_name is None:
|
||||
return None
|
||||
itpm_key = RouterCacheEnum.ITPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute)
|
||||
otpm_key = RouterCacheEnum.OTPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute)
|
||||
return itpm_key, otpm_key
|
||||
|
||||
|
||||
def _estimate_input_tokens(request_kwargs: Optional[dict[str, Any]], model: str = "") -> int:
|
||||
if not request_kwargs:
|
||||
return 0
|
||||
messages = request_kwargs.get("messages")
|
||||
prompt = request_kwargs.get("prompt")
|
||||
input_text = request_kwargs.get("input")
|
||||
# token_counter can raise from any of its tokenizer backends; this is a
|
||||
# best-effort estimate for the ITPM reservation and must never fail the
|
||||
# underlying request. Passing the deployment model name uses a model-specific
|
||||
# tokenizer when available, reducing the reservation over/under-estimate window
|
||||
# between pre-call and post-call reconcile.
|
||||
with contextlib.suppress(Exception):
|
||||
return max(0, int(token_counter(model=model, messages=messages, text=prompt or input_text)))
|
||||
return 0
|
||||
|
||||
|
||||
def _model_max_output_tokens(model_name: str) -> Optional[int]:
|
||||
# litellm.get_model_info raises a bare Exception for an unrecognized model;
|
||||
# this lookup is a fallback default and must never fail the request.
|
||||
with contextlib.suppress(Exception):
|
||||
info = litellm.get_model_info(model=model_name)
|
||||
model_max = info.get("max_output_tokens") or info.get("max_tokens")
|
||||
if model_max is not None:
|
||||
return max(0, int(model_max))
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_max_tokens(request_kwargs: Optional[dict[str, Any]], 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))
|
||||
|
||||
model_name = (deployment.get("litellm_params") or {}).get("model")
|
||||
if model_name:
|
||||
model_max = _model_max_output_tokens(model_name)
|
||||
if model_max is not None:
|
||||
return model_max
|
||||
return 4096
|
||||
|
||||
|
||||
def _get_usage_tokens(usage: Any) -> tuple[int, int, int]:
|
||||
if usage is None:
|
||||
return 0, 0, 0
|
||||
if hasattr(usage, "prompt_tokens") or hasattr(usage, "input_tokens"):
|
||||
prompt = int(getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", 0) or 0)
|
||||
completion = int(getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", 0) or 0)
|
||||
cached = 0
|
||||
details = getattr(usage, "prompt_tokens_details", None)
|
||||
if details is not None:
|
||||
cached = int(getattr(details, "cached_tokens", 0) or 0)
|
||||
if not cached:
|
||||
cached = int(getattr(usage, "cache_read_input_tokens", 0) or 0)
|
||||
return prompt, completion, cached
|
||||
if isinstance(usage, dict):
|
||||
prompt = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
|
||||
completion = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0)
|
||||
details = usage.get("prompt_tokens_details") or {}
|
||||
cached = int(details.get("cached_tokens", 0) or 0) if isinstance(details, dict) else 0
|
||||
if not cached:
|
||||
cached = int(usage.get("cache_read_input_tokens") or 0)
|
||||
return prompt, completion, cached
|
||||
return 0, 0, 0
|
||||
|
||||
|
||||
def _extract_response_usage(response_obj: Any) -> Any:
|
||||
if isinstance(response_obj, dict):
|
||||
return response_obj.get("usage")
|
||||
return getattr(response_obj, "usage", None)
|
||||
|
||||
|
||||
def _usage_is_present(usage: Any) -> bool:
|
||||
"""
|
||||
True only if usage carries an actual input/output breakdown.
|
||||
|
||||
``total_tokens`` alone is deliberately excluded: ``_get_usage_tokens`` has
|
||||
no way to split a bare total into input vs. output, so treating it as
|
||||
"present" would resolve to (0, 0) and refund the full reservation as if
|
||||
zero tokens were used.
|
||||
"""
|
||||
if usage is None:
|
||||
return False
|
||||
fields = ("prompt_tokens", "completion_tokens", "input_tokens", "output_tokens")
|
||||
if isinstance(usage, dict):
|
||||
return any(key in usage for key in fields)
|
||||
return any(hasattr(usage, key) for key in fields)
|
||||
|
||||
|
||||
def _resolve_reconcile_usage_tokens(
|
||||
kwargs: Any,
|
||||
response_obj: Any,
|
||||
) -> tuple[int, int, bool]:
|
||||
"""
|
||||
Resolve billable input and output tokens for post-call reconcile.
|
||||
|
||||
Prefer the response usage object; fall back to standard_logging_object token
|
||||
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 = _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
|
||||
|
||||
if isinstance(kwargs, dict):
|
||||
standard_logging_object = kwargs.get("standard_logging_object")
|
||||
if isinstance(standard_logging_object, dict):
|
||||
prompt_tokens = int(standard_logging_object.get("prompt_tokens") or 0)
|
||||
completion_tokens = int(standard_logging_object.get("completion_tokens") or 0)
|
||||
cached_tokens = 0
|
||||
metadata = standard_logging_object.get("metadata")
|
||||
if isinstance(metadata, dict):
|
||||
cached_tokens = int(metadata.get("cache_read_input_tokens") or 0)
|
||||
# Same rationale as _usage_is_present: a bare total_tokens with no
|
||||
# prompt/completion breakdown can't be split, so it isn't treated
|
||||
# as resolved usage - the reservation is kept instead of refunded.
|
||||
if prompt_tokens or completion_tokens:
|
||||
return max(0, prompt_tokens - cached_tokens), completion_tokens, True
|
||||
|
||||
return 0, 0, False
|
||||
|
||||
|
||||
def _stash_reservation_in_metadata(
|
||||
request_kwargs: Optional[dict[str, Any]],
|
||||
*,
|
||||
itpm_reserved: int,
|
||||
otpm_reserved: int,
|
||||
itpm_cache_key: Optional[str],
|
||||
otpm_cache_key: Optional[str],
|
||||
) -> None:
|
||||
if not request_kwargs:
|
||||
return
|
||||
reservation = {
|
||||
ITPM_RESERVED_KEY: itpm_reserved,
|
||||
OTPM_RESERVED_KEY: otpm_reserved,
|
||||
ITPM_CACHE_KEY: itpm_cache_key,
|
||||
OTPM_CACHE_KEY: otpm_cache_key,
|
||||
}
|
||||
for channel in ("metadata", "litellm_metadata"):
|
||||
existing = request_kwargs.get(channel)
|
||||
if isinstance(existing, dict):
|
||||
existing.update(reservation)
|
||||
elif channel == "metadata":
|
||||
request_kwargs[channel] = dict(reservation)
|
||||
|
||||
|
||||
def _extract_reservation(reservation: dict[str, Any]) -> tuple[int, int, Optional[str], Optional[str]]:
|
||||
itpm_cache_key = reservation.get(ITPM_CACHE_KEY)
|
||||
otpm_cache_key = reservation.get(OTPM_CACHE_KEY)
|
||||
return (
|
||||
int(reservation.get(ITPM_RESERVED_KEY, 0) or 0),
|
||||
int(reservation.get(OTPM_RESERVED_KEY, 0) or 0),
|
||||
itpm_cache_key if isinstance(itpm_cache_key, str) else None,
|
||||
otpm_cache_key if isinstance(otpm_cache_key, str) else None,
|
||||
)
|
||||
|
||||
|
||||
def _reservation_channels(kwargs: Any) -> tuple[Any, ...]:
|
||||
"""
|
||||
Places a reservation may live, in priority order: the top-level metadata
|
||||
channels win over litellm_params.metadata (so a top-level stash is never
|
||||
shadowed), which win over the standard_logging_object copy.
|
||||
"""
|
||||
if not isinstance(kwargs, dict):
|
||||
return ()
|
||||
channels = [kwargs.get("metadata"), kwargs.get("litellm_metadata")]
|
||||
litellm_params = kwargs.get("litellm_params")
|
||||
if isinstance(litellm_params, dict):
|
||||
channels.append(litellm_params.get("metadata"))
|
||||
standard_logging_object = kwargs.get("standard_logging_object")
|
||||
if isinstance(standard_logging_object, dict):
|
||||
channels.append(standard_logging_object.get("metadata"))
|
||||
return tuple(channels)
|
||||
|
||||
|
||||
def _read_reservation_from_kwargs(kwargs: Any) -> tuple[int, int, Optional[str], Optional[str]]:
|
||||
for channel_dict in _reservation_channels(kwargs):
|
||||
if isinstance(channel_dict, dict) and ITPM_RESERVED_KEY in channel_dict:
|
||||
return _extract_reservation(channel_dict)
|
||||
return 0, 0, None, None
|
||||
|
||||
|
||||
def _clear_reservation_from_kwargs(kwargs: Any) -> None:
|
||||
"""
|
||||
Remove the stashed reservation so a retry on a different (e.g. non-IO)
|
||||
deployment does not re-process the already-reconciled/refunded reservation.
|
||||
"""
|
||||
for channel_dict in _reservation_channels(kwargs):
|
||||
if isinstance(channel_dict, dict):
|
||||
for key in (ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY):
|
||||
channel_dict.pop(key, None)
|
||||
|
||||
|
||||
def _reservation_value(value: int, limit: Optional[int]) -> int:
|
||||
if limit is None:
|
||||
return 0
|
||||
if value > 0:
|
||||
return value
|
||||
# Estimation failed (empty messages, unsupported model, tokenizer error).
|
||||
# Reserve a minimal 1-token slot rather than the full limit: the latter
|
||||
# would let one request whose estimate failed fill the entire bucket,
|
||||
# serializing every concurrent request to the deployment until it
|
||||
# completes and reconciles against actual usage.
|
||||
return 1
|
||||
|
||||
|
||||
def _rate_limit_error(limit_label: str, limit: int, current: float) -> litellm.RateLimitError:
|
||||
return litellm.RateLimitError(
|
||||
message=f"Model rate limit exceeded. {limit_label} limit={limit}, current usage={current}",
|
||||
llm_provider="",
|
||||
model="",
|
||||
response=httpx.Response(
|
||||
status_code=429,
|
||||
content=(
|
||||
f"{RouterErrors.user_defined_ratelimit_error.value} "
|
||||
f"{limit_label} limit={limit}. current usage={current}."
|
||||
),
|
||||
headers={"retry-after": str(RoutingArgsTTL)},
|
||||
request=httpx.Request(
|
||||
method="io_token_rate_limit_check",
|
||||
url="https://github.com/BerriAI/litellm",
|
||||
),
|
||||
),
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
|
||||
def _sync_increment_with_rollback(
|
||||
dual_cache: DualCache,
|
||||
key: str,
|
||||
value: int,
|
||||
limit: Optional[int],
|
||||
*,
|
||||
limit_label: str,
|
||||
) -> None:
|
||||
if value <= 0 or limit is None:
|
||||
return
|
||||
current = dual_cache.increment_cache(
|
||||
key=key,
|
||||
value=value,
|
||||
ttl=RoutingArgsTTL,
|
||||
)
|
||||
if current is not None and current > limit:
|
||||
dual_cache.increment_cache(
|
||||
key=key,
|
||||
value=-value,
|
||||
ttl=RoutingArgsTTL,
|
||||
)
|
||||
raise _rate_limit_error(limit_label, limit, current)
|
||||
|
||||
|
||||
async def _increment_with_rollback(
|
||||
dual_cache: DualCache,
|
||||
key: str,
|
||||
value: int,
|
||||
limit: Optional[int],
|
||||
*,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
limit_label: str,
|
||||
) -> None:
|
||||
if value <= 0 or limit is None:
|
||||
return
|
||||
current = await dual_cache.async_increment_cache(
|
||||
key=key,
|
||||
value=value,
|
||||
ttl=RoutingArgsTTL,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
if current is not None and current > limit:
|
||||
await dual_cache.async_increment_cache(
|
||||
key=key,
|
||||
value=-value,
|
||||
ttl=RoutingArgsTTL,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
raise _rate_limit_error(limit_label, limit, current)
|
||||
|
||||
|
||||
def io_token_pre_call_check(
|
||||
dual_cache: DualCache,
|
||||
deployment: dict,
|
||||
) -> Optional[dict]:
|
||||
itpm_limit, otpm_limit = get_deployment_io_token_limits(deployment)
|
||||
if itpm_limit is None and otpm_limit is None:
|
||||
return deployment
|
||||
|
||||
request_kwargs = get_io_token_rate_limit_request_kwargs()
|
||||
_model = (deployment.get("litellm_params") or {}).get("model") or ""
|
||||
estimated_input = _estimate_input_tokens(request_kwargs, model=_model)
|
||||
max_tokens = _resolve_max_tokens(request_kwargs, deployment)
|
||||
|
||||
dt = get_utc_datetime()
|
||||
current_minute = dt.strftime("%H-%M")
|
||||
cache_keys = _get_cache_keys(deployment, current_minute)
|
||||
if cache_keys is None:
|
||||
return deployment
|
||||
itpm_key, otpm_key = cache_keys
|
||||
|
||||
itpm_reserved = 0
|
||||
otpm_reserved = 0
|
||||
|
||||
if itpm_limit is not None:
|
||||
itpm_reserved = _reservation_value(estimated_input, itpm_limit)
|
||||
_sync_increment_with_rollback(
|
||||
dual_cache,
|
||||
itpm_key,
|
||||
itpm_reserved,
|
||||
itpm_limit,
|
||||
limit_label="ITPM",
|
||||
)
|
||||
|
||||
if otpm_limit is not None:
|
||||
otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit)
|
||||
try:
|
||||
_sync_increment_with_rollback(
|
||||
dual_cache,
|
||||
otpm_key,
|
||||
otpm_reserved,
|
||||
otpm_limit,
|
||||
limit_label="OTPM",
|
||||
)
|
||||
except Exception:
|
||||
if itpm_reserved > 0:
|
||||
dual_cache.increment_cache(
|
||||
key=itpm_key,
|
||||
value=-itpm_reserved,
|
||||
ttl=RoutingArgsTTL,
|
||||
)
|
||||
raise
|
||||
|
||||
_stash_reservation_in_metadata(
|
||||
request_kwargs,
|
||||
itpm_reserved=itpm_reserved,
|
||||
otpm_reserved=otpm_reserved,
|
||||
itpm_cache_key=itpm_key if itpm_limit is not None else None,
|
||||
otpm_cache_key=otpm_key if otpm_limit is not None else None,
|
||||
)
|
||||
return deployment
|
||||
|
||||
|
||||
async def async_io_token_pre_call_check(
|
||||
dual_cache: DualCache,
|
||||
deployment: dict,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
) -> Optional[dict]:
|
||||
itpm_limit, otpm_limit = get_deployment_io_token_limits(deployment)
|
||||
if itpm_limit is None and otpm_limit is None:
|
||||
return deployment
|
||||
|
||||
request_kwargs = get_io_token_rate_limit_request_kwargs()
|
||||
_model = (deployment.get("litellm_params") or {}).get("model") or ""
|
||||
estimated_input = _estimate_input_tokens(request_kwargs, model=_model)
|
||||
max_tokens = _resolve_max_tokens(request_kwargs, deployment)
|
||||
|
||||
dt = get_utc_datetime()
|
||||
current_minute = dt.strftime("%H-%M")
|
||||
cache_keys = _get_cache_keys(deployment, current_minute)
|
||||
if cache_keys is None:
|
||||
return deployment
|
||||
itpm_key, otpm_key = cache_keys
|
||||
|
||||
itpm_reserved = 0
|
||||
otpm_reserved = 0
|
||||
|
||||
if itpm_limit is not None:
|
||||
itpm_reserved = _reservation_value(estimated_input, itpm_limit)
|
||||
await _increment_with_rollback(
|
||||
dual_cache,
|
||||
itpm_key,
|
||||
itpm_reserved,
|
||||
itpm_limit,
|
||||
parent_otel_span=parent_otel_span,
|
||||
limit_label="ITPM",
|
||||
)
|
||||
|
||||
if otpm_limit is not None:
|
||||
otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit)
|
||||
try:
|
||||
await _increment_with_rollback(
|
||||
dual_cache,
|
||||
otpm_key,
|
||||
otpm_reserved,
|
||||
otpm_limit,
|
||||
parent_otel_span=parent_otel_span,
|
||||
limit_label="OTPM",
|
||||
)
|
||||
except Exception:
|
||||
# Any failure reserving OTPM (a 429 or a transient cache error) must
|
||||
# release the ITPM reservation already made, or it stays inflated
|
||||
# until the TTL expires.
|
||||
if itpm_reserved > 0:
|
||||
await dual_cache.async_increment_cache(
|
||||
key=itpm_key,
|
||||
value=-itpm_reserved,
|
||||
ttl=RoutingArgsTTL,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
raise
|
||||
|
||||
_stash_reservation_in_metadata(
|
||||
request_kwargs,
|
||||
itpm_reserved=itpm_reserved,
|
||||
otpm_reserved=otpm_reserved,
|
||||
itpm_cache_key=itpm_key if itpm_limit is not None else None,
|
||||
otpm_cache_key=otpm_key if otpm_limit is not None else None,
|
||||
)
|
||||
return deployment
|
||||
|
||||
|
||||
def io_token_reconcile_success(
|
||||
dual_cache: DualCache,
|
||||
kwargs: Any,
|
||||
response_obj: Any,
|
||||
) -> None:
|
||||
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
|
||||
if itpm_key is None and otpm_key is None:
|
||||
return
|
||||
|
||||
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(kwargs, response_obj)
|
||||
|
||||
try:
|
||||
if usage_resolved:
|
||||
if itpm_key is not None:
|
||||
itpm_delta = 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 = completion_tokens - otpm_reserved
|
||||
if otpm_delta != 0:
|
||||
dual_cache.increment_cache(
|
||||
key=otpm_key,
|
||||
value=otpm_delta,
|
||||
ttl=RoutingArgsTTL,
|
||||
)
|
||||
else:
|
||||
verbose_router_logger.debug(
|
||||
"[IO TOKEN LIMIT] usage missing; keeping reservation "
|
||||
f"(itpm_reserved={itpm_reserved}, otpm_reserved={otpm_reserved})"
|
||||
)
|
||||
finally:
|
||||
_clear_reservation_from_kwargs(kwargs)
|
||||
|
||||
verbose_router_logger.debug(
|
||||
f"[IO TOKEN LIMIT] reconciled "
|
||||
f"(usage_resolved={usage_resolved}, itpm_reserved={itpm_reserved}, "
|
||||
f"billable_input={billable_input}, otpm_reserved={otpm_reserved}, output={completion_tokens})"
|
||||
)
|
||||
|
||||
|
||||
async def async_io_token_reconcile_success(
|
||||
dual_cache: DualCache,
|
||||
kwargs: Any,
|
||||
response_obj: Any,
|
||||
*,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
) -> None:
|
||||
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
|
||||
if itpm_key is None and otpm_key is None:
|
||||
return
|
||||
|
||||
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(kwargs, response_obj)
|
||||
|
||||
# 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.
|
||||
try:
|
||||
if usage_resolved:
|
||||
if itpm_key is not None:
|
||||
itpm_delta = 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 = 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,
|
||||
)
|
||||
else:
|
||||
verbose_router_logger.debug(
|
||||
"[IO TOKEN LIMIT] usage missing; keeping reservation "
|
||||
f"(itpm_reserved={itpm_reserved}, otpm_reserved={otpm_reserved})"
|
||||
)
|
||||
finally:
|
||||
_clear_reservation_from_kwargs(kwargs)
|
||||
|
||||
verbose_router_logger.debug(
|
||||
f"[IO TOKEN LIMIT] reconciled "
|
||||
f"(usage_resolved={usage_resolved}, itpm_reserved={itpm_reserved}, "
|
||||
f"billable_input={billable_input}, otpm_reserved={otpm_reserved}, output={completion_tokens})"
|
||||
)
|
||||
|
||||
|
||||
def io_token_refund_failure(
|
||||
dual_cache: DualCache,
|
||||
kwargs: Any,
|
||||
) -> None:
|
||||
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
|
||||
if itpm_key is None and otpm_key is None:
|
||||
return
|
||||
if itpm_key is not None and itpm_reserved > 0:
|
||||
dual_cache.increment_cache(
|
||||
key=itpm_key,
|
||||
value=-itpm_reserved,
|
||||
ttl=RoutingArgsTTL,
|
||||
)
|
||||
if otpm_key is not None and otpm_reserved > 0:
|
||||
dual_cache.increment_cache(
|
||||
key=otpm_key,
|
||||
value=-otpm_reserved,
|
||||
ttl=RoutingArgsTTL,
|
||||
)
|
||||
_clear_reservation_from_kwargs(kwargs)
|
||||
verbose_router_logger.debug(f"[IO TOKEN LIMIT] refunded ITPM={itpm_reserved} OTPM={otpm_reserved}")
|
||||
|
||||
|
||||
def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Optional[dict[str, Any]]) -> None:
|
||||
"""
|
||||
Synchronously refund and clear any reservation a previous deployment
|
||||
attempt stashed in ``kwargs``, before it's overwritten for the next
|
||||
attempt (retry/fallback).
|
||||
|
||||
``set_io_token_rate_limit_request_kwargs`` strips reservation sentinels
|
||||
from ``kwargs`` on every deployment pick (a security measure so a
|
||||
caller-forged reservation can't be replayed). Without this refund, a
|
||||
retry after a non-RateLimitError failure (e.g. an upstream 500) would
|
||||
wipe deployment A's still-unreconciled reservation before its failure
|
||||
event - which may be scheduled as a background task - gets a chance to
|
||||
refund it, permanently stranding the reservation until its TTL expires
|
||||
and causing false rate-limit errors for subsequent requests.
|
||||
|
||||
ponytail: uses sync ``DualCache.increment_cache`` which issues a blocking
|
||||
Redis INCR when a Redis backend is configured. This only triggers on
|
||||
streaming mid-stream retries (non-streaming failures await their failure
|
||||
handler before retrying, so the sentinels are already cleared). Upgrade
|
||||
path: make ``_update_kwargs_with_deployment`` async and switch to
|
||||
``async_io_token_refund_failure`` — requires touching all callers.
|
||||
"""
|
||||
if not kwargs:
|
||||
return
|
||||
io_token_refund_failure(dual_cache, kwargs)
|
||||
|
||||
|
||||
async def async_io_token_refund_failure(
|
||||
dual_cache: DualCache,
|
||||
kwargs: Any,
|
||||
*,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
) -> None:
|
||||
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
|
||||
if itpm_key is None and otpm_key is None:
|
||||
return
|
||||
if itpm_key is not None and itpm_reserved > 0:
|
||||
await dual_cache.async_increment_cache(
|
||||
key=itpm_key,
|
||||
value=-itpm_reserved,
|
||||
ttl=RoutingArgsTTL,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
if otpm_key is not None and otpm_reserved > 0:
|
||||
await dual_cache.async_increment_cache(
|
||||
key=otpm_key,
|
||||
value=-otpm_reserved,
|
||||
ttl=RoutingArgsTTL,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
_clear_reservation_from_kwargs(kwargs)
|
||||
verbose_router_logger.debug(f"[IO TOKEN LIMIT] refunded ITPM={itpm_reserved} OTPM={otpm_reserved}")
|
||||
|
||||
|
||||
def build_io_token_rate_limit_headers(
|
||||
*,
|
||||
itpm_limit: Optional[int],
|
||||
otpm_limit: Optional[int],
|
||||
current_itpm: Optional[int],
|
||||
current_otpm: Optional[int],
|
||||
) -> dict[str, int]:
|
||||
headers: dict[str, int] = {}
|
||||
reset = seconds_until_minute_reset()
|
||||
if itpm_limit is not None:
|
||||
usage = current_itpm or 0
|
||||
headers["x-ratelimit-limit-input-tokens"] = itpm_limit
|
||||
headers["x-ratelimit-remaining-input-tokens"] = max(0, itpm_limit - usage)
|
||||
headers["x-ratelimit-reset-input-tokens"] = reset
|
||||
if otpm_limit is not None:
|
||||
usage = current_otpm or 0
|
||||
headers["x-ratelimit-limit-output-tokens"] = otpm_limit
|
||||
headers["x-ratelimit-remaining-output-tokens"] = max(0, otpm_limit - usage)
|
||||
headers["x-ratelimit-reset-output-tokens"] = reset
|
||||
return headers
|
||||
|
|
@ -1,13 +1,16 @@
|
|||
"""
|
||||
Enforce TPM/RPM rate limits set on model deployments.
|
||||
Enforce TPM/RPM or separate ITPM/OTPM rate limits set on model deployments.
|
||||
|
||||
This pre-call check ensures that model-level TPM/RPM limits are enforced
|
||||
across all requests, regardless of routing strategy.
|
||||
When enabled via router_settings.optional_pre_call_checks: ["enforce_model_rate_limits"]
|
||||
|
||||
When enabled via `enforce_model_rate_limits: true` in litellm_settings,
|
||||
requests that exceed the configured TPM/RPM limits will receive a 429 error.
|
||||
- tpm/rpm: combined TPM + optional RPM (legacy)
|
||||
- itpm/otpm: separate input/output tokens per minute
|
||||
|
||||
When a deployment sets both itpm/otpm and tpm/rpm, both are enforced. A warning
|
||||
is logged the first time such a deployment is seen.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
|
@ -16,6 +19,17 @@ import litellm
|
|||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
|
||||
async_io_token_pre_call_check,
|
||||
async_io_token_reconcile_success,
|
||||
async_io_token_refund_failure,
|
||||
deployment_has_io_token_limits,
|
||||
get_io_token_rate_limit_request_kwargs,
|
||||
io_token_pre_call_check,
|
||||
io_token_reconcile_success,
|
||||
io_token_refund_failure,
|
||||
ITPM_RESERVED_KEY,
|
||||
)
|
||||
from litellm.types.router import RouterErrors
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
|
@ -34,7 +48,7 @@ class RoutingArgs:
|
|||
|
||||
class ModelRateLimitingCheck(CustomLogger):
|
||||
"""
|
||||
Pre-call check that enforces TPM/RPM limits on model deployments.
|
||||
Pre-call check that enforces TPM/RPM or ITPM/OTPM limits on model deployments.
|
||||
|
||||
This check runs before each request and raises a RateLimitError
|
||||
if the deployment has exceeded its configured TPM or RPM limits.
|
||||
|
|
@ -45,6 +59,42 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
|
||||
def __init__(self, dual_cache: DualCache):
|
||||
self.dual_cache = dual_cache
|
||||
# model_ids already warned about itpm/otpm + tpm/rpm on the same deployment,
|
||||
# so the warning is logged once per deployment rather than per request.
|
||||
self._io_token_conflict_warned_ids: set[str] = set()
|
||||
|
||||
def _warn_io_token_and_tpm_rpm_coexist_once(self, deployment: dict) -> None:
|
||||
tpm_limit, rpm_limit = self._get_deployment_limits(deployment)
|
||||
if tpm_limit is None and rpm_limit is None:
|
||||
return
|
||||
model_id = deployment.get("model_info", {}).get("id")
|
||||
# Dedup per deployment id; if there is no id (degenerate config) don't
|
||||
# collapse every such deployment onto one key - warn each time instead.
|
||||
if model_id is not None:
|
||||
if model_id in self._io_token_conflict_warned_ids:
|
||||
return
|
||||
self._io_token_conflict_warned_ids.add(str(model_id))
|
||||
verbose_router_logger.warning(
|
||||
f"Deployment '{model_id}' configures itpm/otpm alongside tpm/rpm; "
|
||||
"both limit types are enforced on this deployment"
|
||||
)
|
||||
|
||||
def _refund_io_token_reservation_if_any(self) -> None:
|
||||
request_kwargs = get_io_token_rate_limit_request_kwargs()
|
||||
if request_kwargs is not None:
|
||||
io_token_refund_failure(self.dual_cache, request_kwargs)
|
||||
|
||||
async def _async_refund_io_token_reservation_if_any(
|
||||
self,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
) -> None:
|
||||
request_kwargs = get_io_token_rate_limit_request_kwargs()
|
||||
if request_kwargs is not None:
|
||||
await async_io_token_refund_failure(
|
||||
self.dual_cache,
|
||||
request_kwargs,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
||||
def _get_deployment_limits(self, deployment: Dict) -> tuple[Optional[int], Optional[int]]:
|
||||
"""
|
||||
|
|
@ -93,6 +143,15 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
Raises RateLimitError if deployment exceeds TPM/RPM limits.
|
||||
"""
|
||||
try:
|
||||
io_reservation_made = False
|
||||
if deployment_has_io_token_limits(deployment):
|
||||
self._warn_io_token_and_tpm_rpm_coexist_once(deployment)
|
||||
io_token_pre_call_check(
|
||||
self.dual_cache,
|
||||
deployment,
|
||||
)
|
||||
io_reservation_made = True
|
||||
|
||||
tpm_limit, rpm_limit = self._get_deployment_limits(deployment)
|
||||
|
||||
# If no limits are set, allow the request
|
||||
|
|
@ -149,6 +208,8 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
return deployment
|
||||
|
||||
except litellm.RateLimitError:
|
||||
if io_reservation_made:
|
||||
self._refund_io_token_reservation_if_any()
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.pre_call_check: {str(e)}")
|
||||
|
|
@ -159,9 +220,19 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
"""
|
||||
Async pre-call check for model rate limits.
|
||||
|
||||
Raises RateLimitError if deployment exceeds TPM/RPM limits.
|
||||
Raises RateLimitError if deployment exceeds TPM/RPM or ITPM/OTPM limits.
|
||||
"""
|
||||
try:
|
||||
io_reservation_made = False
|
||||
if deployment_has_io_token_limits(deployment):
|
||||
self._warn_io_token_and_tpm_rpm_coexist_once(deployment)
|
||||
await async_io_token_pre_call_check(
|
||||
self.dual_cache,
|
||||
deployment,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
io_reservation_made = True
|
||||
|
||||
tpm_limit, rpm_limit = self._get_deployment_limits(deployment)
|
||||
|
||||
# If no limits are set, allow the request
|
||||
|
|
@ -225,6 +296,8 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
return deployment
|
||||
|
||||
except litellm.RateLimitError:
|
||||
if io_reservation_made:
|
||||
await self._async_refund_io_token_reservation_if_any(parent_otel_span=parent_otel_span)
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_pre_call_check: {str(e)}")
|
||||
|
|
@ -232,14 +305,31 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
return deployment
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
"""
|
||||
Track TPM usage after successful request.
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
_get_parent_otel_span_from_kwargs,
|
||||
)
|
||||
|
||||
This updates the TPM counter with the actual tokens used.
|
||||
Always tracks tokens - the pre-call check handles enforcement.
|
||||
"""
|
||||
try:
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = 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 = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {}
|
||||
kwargs_metadata = 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.
|
||||
|
||||
if standard_logging_object is None:
|
||||
return
|
||||
|
||||
|
|
@ -272,6 +362,19 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
except Exception as e:
|
||||
verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_log_success_event: {str(e)}")
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
_get_parent_otel_span_from_kwargs,
|
||||
)
|
||||
|
||||
# Never fail the primary logging pipeline over an io-token refund error.
|
||||
with contextlib.suppress(Exception):
|
||||
await async_io_token_refund_failure(
|
||||
self.dual_cache,
|
||||
kwargs,
|
||||
parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs),
|
||||
)
|
||||
|
||||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
"""
|
||||
Sync version of tracking TPM usage after successful request.
|
||||
|
|
@ -279,6 +382,18 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
"""
|
||||
try:
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object")
|
||||
slo_metadata = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {}
|
||||
kwargs_metadata = 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.
|
||||
|
||||
if standard_logging_object is None:
|
||||
return
|
||||
|
||||
|
|
@ -304,3 +419,10 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
|
||||
except Exception as e:
|
||||
verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.log_success_event: {str(e)}")
|
||||
|
||||
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
with contextlib.suppress(Exception):
|
||||
io_token_refund_failure(
|
||||
self.dual_cache,
|
||||
kwargs,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -214,6 +214,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
|
|||
custom_llm_provider: Optional[str] = None
|
||||
tpm: Optional[int] = None
|
||||
rpm: Optional[int] = None
|
||||
itpm: Optional[int] = None
|
||||
otpm: Optional[int] = None
|
||||
timeout: Optional[Union[float, str, httpx.Timeout]] = None # if str, pass in as os.environ/
|
||||
stream_timeout: Optional[Union[float, str]] = (
|
||||
None # timeout when making stream=True calls, if str, pass in as os.environ/
|
||||
|
|
@ -359,6 +361,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
|
|||
custom_llm_provider: Optional[str]
|
||||
tpm: Optional[int]
|
||||
rpm: Optional[int]
|
||||
itpm: Optional[int]
|
||||
otpm: Optional[int]
|
||||
order: Optional[int]
|
||||
weight: Optional[int]
|
||||
max_parallel_requests: Optional[int]
|
||||
|
|
@ -552,6 +556,8 @@ class ModelGroupInfo(BaseModel):
|
|||
] = Field(default="chat")
|
||||
tpm: Optional[int] = None
|
||||
rpm: Optional[int] = None
|
||||
itpm: Optional[int] = None
|
||||
otpm: Optional[int] = None
|
||||
supports_parallel_function_calling: bool = Field(default=False)
|
||||
supports_vision: bool = Field(default=False)
|
||||
supports_web_search: bool = Field(default=False)
|
||||
|
|
@ -749,6 +755,8 @@ class RoutingStrategy(enum.Enum):
|
|||
class RouterCacheEnum(enum.Enum):
|
||||
TPM = "global_router:{id}:{model}:tpm:{current_minute}"
|
||||
RPM = "global_router:{id}:{model}:rpm:{current_minute}"
|
||||
ITPM = "global_router:{id}:{model}:itpm:{current_minute}"
|
||||
OTPM = "global_router:{id}:{model}:otpm:{current_minute}"
|
||||
|
||||
|
||||
class GenericBudgetWindowDetails(BaseModel):
|
||||
|
|
|
|||
|
|
@ -3120,6 +3120,8 @@ all_litellm_params = (
|
|||
"client",
|
||||
"rpm",
|
||||
"tpm",
|
||||
"itpm",
|
||||
"otpm",
|
||||
"max_parallel_requests",
|
||||
"input_cost_per_token",
|
||||
"output_cost_per_token",
|
||||
|
|
|
|||
|
|
@ -324,7 +324,7 @@
|
|||
"limit": 883
|
||||
},
|
||||
"UP006": {
|
||||
"limit": 12870
|
||||
"limit": 12869
|
||||
},
|
||||
"UP007": {
|
||||
"limit": 2570
|
||||
|
|
|
|||
|
|
@ -523,6 +523,55 @@ async def test_deployment_callback_on_success(sync_mode):
|
|||
assert tpm_key is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment():
|
||||
"""
|
||||
An IO-limited deployment (itpm/otpm, no tpm/rpm) must still record TPM usage
|
||||
in the router's routing counter so TPM-aware routing strategies see its real
|
||||
load in mixed model groups; its itpm/otpm enforcement runs separately.
|
||||
"""
|
||||
import time
|
||||
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "opus",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "sk-fake",
|
||||
"itpm": 1000,
|
||||
},
|
||||
"model_info": {"id": "io-100"},
|
||||
}
|
||||
]
|
||||
router = Router(model_list=model_list)
|
||||
|
||||
standard_logging_payload = create_standard_logging_payload()
|
||||
standard_logging_payload["total_tokens"] = 100
|
||||
standard_logging_payload["model_id"] = "io-100"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"deployment": "openai/gpt-4o-mini",
|
||||
"model_group": "opus",
|
||||
},
|
||||
"model_info": {"id": "io-100"},
|
||||
},
|
||||
"standard_logging_object": standard_logging_payload,
|
||||
}
|
||||
response = litellm.ModelResponse(model="openai/gpt-4o-mini", usage={"total_tokens": 100})
|
||||
|
||||
tpm_key = await router.deployment_callback_on_success(
|
||||
kwargs=kwargs,
|
||||
completion_response=response,
|
||||
start_time=time.time(),
|
||||
end_time=time.time(),
|
||||
)
|
||||
|
||||
# The IO deployment is no longer skipped: its TPM routing counter is tracked.
|
||||
assert tpm_key is not None
|
||||
assert await router.cache.async_get_cache(key=tpm_key) == 100
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_callback_on_failure(model_list):
|
||||
"""Test if the '_deployment_callback_on_failure' function is working correctly"""
|
||||
|
|
@ -923,6 +972,227 @@ async def test_set_response_headers_subtracts_in_flight_delta(model_list):
|
|||
assert headers["x-ratelimit-limit-requests"] == 100
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_list):
|
||||
"""
|
||||
The in-flight replay applies only to the post-incremented TPM/RPM counters
|
||||
(`x-ratelimit-remaining-tokens` / `-requests`). The ITPM/OTPM counters are
|
||||
incremented at reservation time (pre-call), so the input/output token
|
||||
headers already reflect this request and must pass through untouched.
|
||||
"""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class _Usage(BaseModel):
|
||||
total_tokens: int = 30
|
||||
prompt_tokens: int = 20
|
||||
completion_tokens: int = 10
|
||||
|
||||
class _Resp(BaseModel):
|
||||
usage: _Usage = _Usage()
|
||||
_hidden_params: dict = {}
|
||||
|
||||
router = Router(model_list=model_list)
|
||||
router.get_remaining_model_group_usage = AsyncMock(
|
||||
return_value={
|
||||
"x-ratelimit-remaining-tokens": 1000,
|
||||
"x-ratelimit-remaining-requests": 100,
|
||||
"x-ratelimit-remaining-input-tokens": 1000,
|
||||
"x-ratelimit-remaining-output-tokens": 500,
|
||||
}
|
||||
)
|
||||
|
||||
resp = _Resp()
|
||||
resp._hidden_params = {}
|
||||
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
|
||||
|
||||
headers = resp._hidden_params["additional_headers"]
|
||||
# TPM/RPM headers replay the in-flight increment...
|
||||
assert headers["x-ratelimit-remaining-tokens"] == 970
|
||||
assert headers["x-ratelimit-remaining-requests"] == 99
|
||||
# ...but the reservation-based input/output headers pass through unchanged.
|
||||
assert headers["x-ratelimit-remaining-input-tokens"] == 1000
|
||||
assert headers["x-ratelimit-remaining-output-tokens"] == 500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_model_group_io_token_usage_sums_across_deployments():
|
||||
"""
|
||||
get_model_group_io_token_usage must sum ITPM/OTPM across every deployment
|
||||
in the model group (not just the first), reading the same per-deployment
|
||||
cache keys the pre-call reservation writes to.
|
||||
"""
|
||||
from litellm.types.router import RouterCacheEnum
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "opus",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"itpm": 1000,
|
||||
"otpm": 500,
|
||||
},
|
||||
"model_info": {"id": "io-usage-dep-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "opus",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"itpm": 1000,
|
||||
"otpm": 500,
|
||||
},
|
||||
"model_info": {"id": "io-usage-dep-2"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
keys_and_values = [
|
||||
(
|
||||
RouterCacheEnum.ITPM.value.format(
|
||||
id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute
|
||||
),
|
||||
30,
|
||||
),
|
||||
(
|
||||
RouterCacheEnum.OTPM.value.format(
|
||||
id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute
|
||||
),
|
||||
10,
|
||||
),
|
||||
(
|
||||
RouterCacheEnum.ITPM.value.format(
|
||||
id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute
|
||||
),
|
||||
70,
|
||||
),
|
||||
(
|
||||
RouterCacheEnum.OTPM.value.format(
|
||||
id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute
|
||||
),
|
||||
20,
|
||||
),
|
||||
]
|
||||
for key, value in keys_and_values:
|
||||
await router.cache.async_increment_cache(key=key, value=value, ttl=60)
|
||||
|
||||
current_itpm, current_otpm = await router.get_model_group_io_token_usage("opus")
|
||||
|
||||
assert current_itpm == 100
|
||||
assert current_otpm == 30
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_model_group_io_token_usage_no_deployments_returns_none():
|
||||
router = Router(model_list=[])
|
||||
current_itpm, current_otpm = await router.get_model_group_io_token_usage(
|
||||
"nonexistent-group"
|
||||
)
|
||||
assert current_itpm is None
|
||||
assert current_otpm is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_remaining_model_group_usage_merges_io_and_tpm_headers(model_list):
|
||||
"""
|
||||
A model group with both itpm/otpm and tpm/rpm limits must expose the
|
||||
standard remaining-tokens/requests headers alongside the input/output token
|
||||
headers, so clients and prometheus gauges relying on either still get data.
|
||||
"""
|
||||
from unittest.mock import Mock
|
||||
|
||||
from litellm.types.router import ModelGroupInfo
|
||||
|
||||
router = Router(model_list=model_list)
|
||||
router._cached_get_model_group_info = Mock(
|
||||
return_value=ModelGroupInfo(
|
||||
model_group="gpt-3.5-turbo",
|
||||
providers=["openai"],
|
||||
itpm=2000,
|
||||
otpm=1000,
|
||||
tpm=5000,
|
||||
rpm=50,
|
||||
)
|
||||
)
|
||||
router.get_model_group_io_token_usage = AsyncMock(return_value=(100, 40))
|
||||
router.get_model_group_usage = AsyncMock(return_value=(500, 5))
|
||||
|
||||
headers = await router.get_remaining_model_group_usage("gpt-3.5-turbo")
|
||||
|
||||
assert headers["x-ratelimit-remaining-input-tokens"] == 1900
|
||||
assert headers["x-ratelimit-remaining-output-tokens"] == 960
|
||||
assert headers["x-ratelimit-remaining-tokens"] == 4500
|
||||
assert headers["x-ratelimit-remaining-requests"] == 45
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_response_headers_native_input_token_header_does_not_suppress_router_headers(model_list):
|
||||
"""
|
||||
A provider that natively returns `x-ratelimit-remaining-input-tokens` must
|
||||
not suppress the router's own remaining-tokens/requests headers for a
|
||||
non-IO model group.
|
||||
"""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class _Usage(BaseModel):
|
||||
total_tokens: int = 42
|
||||
|
||||
class _Resp(BaseModel):
|
||||
usage: _Usage = _Usage()
|
||||
_hidden_params: dict = {}
|
||||
|
||||
router = Router(model_list=model_list)
|
||||
router.get_remaining_model_group_usage = AsyncMock(
|
||||
return_value={
|
||||
"x-ratelimit-remaining-tokens": 1000,
|
||||
"x-ratelimit-remaining-requests": 100,
|
||||
}
|
||||
)
|
||||
|
||||
resp = _Resp()
|
||||
resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-input-tokens": 5}}
|
||||
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
|
||||
|
||||
headers = resp._hidden_params["additional_headers"]
|
||||
assert headers["x-ratelimit-remaining-tokens"] == 958
|
||||
assert headers["x-ratelimit-remaining-requests"] == 99
|
||||
# the provider's native header is left untouched
|
||||
assert headers["x-ratelimit-remaining-input-tokens"] == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_response_headers_native_token_header_does_not_suppress_io_headers(model_list):
|
||||
from pydantic import BaseModel
|
||||
|
||||
class _Usage(BaseModel):
|
||||
total_tokens: int = 42
|
||||
|
||||
class _Resp(BaseModel):
|
||||
usage: _Usage = _Usage()
|
||||
_hidden_params: dict = {}
|
||||
|
||||
router = Router(model_list=model_list)
|
||||
router.get_remaining_model_group_usage = AsyncMock(
|
||||
return_value={
|
||||
"x-ratelimit-remaining-tokens": 1000,
|
||||
"x-ratelimit-remaining-requests": 100,
|
||||
"x-ratelimit-remaining-input-tokens": 900,
|
||||
"x-ratelimit-remaining-output-tokens": 450,
|
||||
}
|
||||
)
|
||||
|
||||
resp = _Resp()
|
||||
resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-tokens": 5}}
|
||||
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
|
||||
|
||||
headers = resp._hidden_params["additional_headers"]
|
||||
assert headers["x-ratelimit-remaining-tokens"] == 5
|
||||
assert headers["x-ratelimit-remaining-requests"] == 99
|
||||
assert headers["x-ratelimit-remaining-input-tokens"] == 900
|
||||
assert headers["x-ratelimit-remaining-output-tokens"] == 450
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_response_headers_handles_missing_usage(model_list):
|
||||
"""
|
||||
|
|
@ -952,6 +1222,72 @@ async def test_set_response_headers_handles_missing_usage(model_list):
|
|||
assert headers["x-ratelimit-remaining-requests"] == 99
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_response_headers_dict_anthropic_messages_response(model_list):
|
||||
"""Anthropic /v1/messages returns a dict; IO rate-limit headers must attach."""
|
||||
router = Router(model_list=model_list)
|
||||
router.get_remaining_model_group_usage = AsyncMock(
|
||||
return_value={
|
||||
"x-ratelimit-limit-input-tokens": 25,
|
||||
"x-ratelimit-remaining-input-tokens": 20,
|
||||
"x-ratelimit-limit-output-tokens": 100,
|
||||
"x-ratelimit-remaining-output-tokens": 95,
|
||||
}
|
||||
)
|
||||
|
||||
resp = {
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"usage": {"input_tokens": 5, "output_tokens": 1},
|
||||
}
|
||||
await router.set_response_headers(response=resp, model_group="io-itpm-strict")
|
||||
|
||||
assert "_hidden_params" in resp
|
||||
headers = resp["_hidden_params"]["additional_headers"]
|
||||
assert headers["x-litellm-model-group"] == "io-itpm-strict"
|
||||
assert headers["x-ratelimit-limit-input-tokens"] == 25
|
||||
assert headers["x-ratelimit-remaining-input-tokens"] == 20
|
||||
assert headers["x-ratelimit-remaining-output-tokens"] == 95
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_response_headers_wraps_bare_async_generator(model_list):
|
||||
"""
|
||||
Streaming responses that never go through Router.make_call's usual
|
||||
object-based wrappers (e.g. the Anthropic /v1/messages -> Responses API
|
||||
bridge, which yields a raw async generator with no `_hidden_params` slot)
|
||||
must still get IO rate-limit headers attached via a thin wrapper.
|
||||
"""
|
||||
|
||||
async def _raw_generator():
|
||||
yield {"type": "message_start"}
|
||||
yield {"type": "message_stop"}
|
||||
|
||||
router = Router(model_list=model_list)
|
||||
router.get_remaining_model_group_usage = AsyncMock(
|
||||
return_value={
|
||||
"x-ratelimit-limit-input-tokens": 25,
|
||||
"x-ratelimit-remaining-input-tokens": 20,
|
||||
}
|
||||
)
|
||||
|
||||
wrapped = await router.set_response_headers(response=_raw_generator(), model_group="io-itpm-strict")
|
||||
|
||||
assert hasattr(wrapped, "_hidden_params")
|
||||
headers = wrapped._hidden_params["additional_headers"]
|
||||
assert headers["x-litellm-model-group"] == "io-itpm-strict"
|
||||
assert headers["x-ratelimit-limit-input-tokens"] == 25
|
||||
assert headers["x-ratelimit-remaining-input-tokens"] == 20
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
assert isinstance(wrapped, AsyncIterator)
|
||||
chunks = [chunk async for chunk in wrapped]
|
||||
assert chunks == [{"type": "message_start"}, {"type": "message_stop"}]
|
||||
|
||||
|
||||
def test_get_all_deployments(model_list):
|
||||
"""Test if the 'get_all_deployments' function is working correctly"""
|
||||
router = Router(model_list=model_list)
|
||||
|
|
|
|||
|
|
@ -4090,7 +4090,7 @@ class TestResponseCostHeaderForTypedDictResponses:
|
|||
logging_obj._on_deferred_stream_complete = None
|
||||
return logging_obj
|
||||
|
||||
async def _drive_non_streaming(self, *, monkeypatch, response, logging_obj, route_type):
|
||||
async def _drive_non_streaming(self, *, monkeypatch, response, logging_obj, route_type, return_result=False):
|
||||
import litellm.proxy.common_request_processing as crp
|
||||
from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth
|
||||
|
||||
|
|
@ -4119,7 +4119,7 @@ class TestResponseCostHeaderForTypedDictResponses:
|
|||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
):
|
||||
await processing_obj.base_process_llm_request(
|
||||
result = await processing_obj.base_process_llm_request(
|
||||
request=MagicMock(spec=Request, headers={}),
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=RealUserAPIKeyAuth(api_key="sk-test"),
|
||||
|
|
@ -4131,6 +4131,8 @@ class TestResponseCostHeaderForTypedDictResponses:
|
|||
llm_router=None,
|
||||
skip_pre_call_logic=True,
|
||||
)
|
||||
if return_result:
|
||||
return fastapi_response, result
|
||||
return fastapi_response
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -4352,3 +4354,46 @@ class TestResponseCostHeaderForTypedDictResponses:
|
|||
|
||||
assert "x-litellm-response-cost" not in fastapi_response.headers
|
||||
recompute.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_messages_typeddict_does_not_leak_hidden_params_into_response_body(self, monkeypatch):
|
||||
"""
|
||||
Router.set_response_headers now writes rate-limit headers onto dict-shaped
|
||||
responses (e.g. Anthropic /v1/messages, whose AnthropicMessagesResponse is a
|
||||
TypedDict) via response["_hidden_params"] = ... . Unlike a pydantic model's
|
||||
private attribute, that key is indistinguishable from any other dict key and
|
||||
would otherwise serialize verbatim into the client-facing JSON body, leaking
|
||||
response_cost/model_id/api_base/fallback errors. base_process_llm_request
|
||||
must strip it before returning the response to the endpoint layer.
|
||||
"""
|
||||
from litellm.types.utils import AnthropicMessagesResponse
|
||||
|
||||
response = AnthropicMessagesResponse(
|
||||
id="msg_1",
|
||||
type="message",
|
||||
role="assistant",
|
||||
content=[{"type": "text", "text": "hi"}],
|
||||
model="claude-haiku-4-5",
|
||||
usage={"input_tokens": 10, "output_tokens": 5},
|
||||
)
|
||||
response["_hidden_params"] = {
|
||||
"additional_headers": {"x-ratelimit-limit-input-tokens": "25"},
|
||||
"response_cost": 0.00123,
|
||||
"model_id": "internal-deployment-id",
|
||||
}
|
||||
logging_obj = self._build_logging_obj(
|
||||
model_call_details={"response_cost": 0.00123},
|
||||
response_cost_calculator=MagicMock(return_value=999.0),
|
||||
)
|
||||
|
||||
fastapi_response, result = await self._drive_non_streaming(
|
||||
monkeypatch=monkeypatch,
|
||||
response=response,
|
||||
logging_obj=logging_obj,
|
||||
route_type="anthropic_messages",
|
||||
return_result=True,
|
||||
)
|
||||
|
||||
assert "_hidden_params" not in result
|
||||
assert fastapi_response.headers["x-ratelimit-limit-input-tokens"] == "25"
|
||||
assert fastapi_response.headers["x-litellm-response-cost"] == "0.00123"
|
||||
|
|
|
|||
|
|
@ -140,3 +140,25 @@ def test_get_fallback_errors_from_headers_invalid_json_returns_empty():
|
|||
def test_get_fallback_errors_from_headers_missing_key_returns_empty():
|
||||
result = get_fallback_errors_from_headers({})
|
||||
assert result == []
|
||||
|
||||
|
||||
def test_get_hidden_params_dict_with_dict_response():
|
||||
response = {"id": "msg_1", "usage": {"input_tokens": 1, "output_tokens": 2}}
|
||||
assert get_hidden_params_dict(response) == {}
|
||||
|
||||
hidden_params = get_hidden_params_dict(response, create=True)
|
||||
assert hidden_params == {}
|
||||
assert response["_hidden_params"] == {}
|
||||
|
||||
response["_hidden_params"] = {"additional_headers": {"x-test": "1"}}
|
||||
assert get_hidden_params_dict(response) == {
|
||||
"additional_headers": {"x-test": "1"},
|
||||
}
|
||||
|
||||
|
||||
def test_add_fallback_headers_to_dict_response():
|
||||
response = {"id": "msg_1"}
|
||||
result = add_fallback_headers_to_response(response=response, attempted_fallbacks=1)
|
||||
|
||||
assert result is response
|
||||
assert response["_hidden_params"]["additional_headers"]["x-litellm-attempted-fallbacks"] == 1
|
||||
|
|
|
|||
979
tests/test_litellm/test_router/test_io_token_rate_limits.py
Normal file
979
tests/test_litellm/test_router/test_io_token_rate_limits.py
Normal file
|
|
@ -0,0 +1,979 @@
|
|||
"""
|
||||
Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
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,
|
||||
_reservation_value,
|
||||
_resolve_max_tokens,
|
||||
async_io_token_pre_call_check,
|
||||
async_io_token_reconcile_success,
|
||||
build_io_token_rate_limit_headers,
|
||||
deployment_has_io_token_limits,
|
||||
get_io_token_rate_limit_request_kwargs,
|
||||
io_token_reconcile_success,
|
||||
io_token_refund_failure,
|
||||
refund_stale_reservation_before_retry,
|
||||
set_io_token_rate_limit_request_kwargs,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
|
||||
ModelRateLimitingCheck,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
|
||||
class TestIOTokenRateLimitHelpers:
|
||||
def test_deployment_has_io_token_limits(self):
|
||||
assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}})
|
||||
assert not deployment_has_io_token_limits({"litellm_params": {"model": "x"}})
|
||||
|
||||
def test_reservation_value_minimal_when_estimate_fails(self):
|
||||
# A failed/empty estimate (0) must reserve a minimal slot, not the
|
||||
# entire limit - otherwise one request whose estimate failed fills
|
||||
# the whole bucket and blocks every concurrent request until it
|
||||
# completes and reconciles.
|
||||
assert _reservation_value(0, 100) == 1
|
||||
assert _reservation_value(0, 1) == 1
|
||||
# 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
|
||||
|
||||
def test_build_io_token_rate_limit_headers(self):
|
||||
headers = build_io_token_rate_limit_headers(
|
||||
itpm_limit=200,
|
||||
otpm_limit=40,
|
||||
current_itpm=15,
|
||||
current_otpm=4,
|
||||
)
|
||||
assert headers["x-ratelimit-limit-input-tokens"] == 200
|
||||
assert headers["x-ratelimit-remaining-input-tokens"] == 185
|
||||
assert headers["x-ratelimit-limit-output-tokens"] == 40
|
||||
assert headers["x-ratelimit-remaining-output-tokens"] == 36
|
||||
|
||||
|
||||
class TestModelRateLimitingCheckIOTokens:
|
||||
@pytest.mark.asyncio
|
||||
async def test_itpm_reservation_and_reconcile(self):
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"itpm": 100,
|
||||
"otpm": 50,
|
||||
},
|
||||
"model_info": {"id": "io-test-id"},
|
||||
"model_name": "opus",
|
||||
}
|
||||
|
||||
request_kwargs = {
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 10,
|
||||
"metadata": {},
|
||||
}
|
||||
set_io_token_rate_limit_request_kwargs(request_kwargs)
|
||||
await check.async_pre_call_check(deployment)
|
||||
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}"
|
||||
otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}"
|
||||
|
||||
kwargs = {
|
||||
"standard_logging_object": {
|
||||
"model_id": "io-test-id",
|
||||
"hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"},
|
||||
"metadata": dict(request_kwargs["metadata"]),
|
||||
},
|
||||
"metadata": request_kwargs["metadata"],
|
||||
}
|
||||
response = ModelResponse(
|
||||
choices=[
|
||||
{
|
||||
"message": {"role": "assistant", "content": "hi"},
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8),
|
||||
)
|
||||
await check.async_log_success_event(kwargs, response, None, None)
|
||||
|
||||
current_itpm = await dual_cache.async_get_cache(key=itpm_key)
|
||||
current_otpm = await dual_cache.async_get_cache(key=otpm_key)
|
||||
# ITPM tracks input tokens only (billable prompt tokens), not output.
|
||||
assert current_itpm == 5
|
||||
assert current_otpm == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_itpm_limit_raises_429(self):
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"itpm": 5,
|
||||
},
|
||||
"model_info": {"id": "io-limit-id"},
|
||||
"model_name": "opus",
|
||||
}
|
||||
|
||||
# ITPM enforces input tokens only; the prompt alone must exceed the limit,
|
||||
# a large max_tokens must not contribute to the ITPM reservation.
|
||||
set_io_token_rate_limit_request_kwargs(
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "hello world this is a longer prompt that exceeds the tiny itpm limit",
|
||||
}
|
||||
],
|
||||
"max_tokens": 10,
|
||||
"metadata": {},
|
||||
}
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.RateLimitError) as exc_info:
|
||||
await check.async_pre_call_check(deployment)
|
||||
|
||||
assert "ITPM limit=5" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self):
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
otpm_limit = 10
|
||||
max_tokens = 4
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"otpm": otpm_limit,
|
||||
},
|
||||
"model_info": {"id": "io-otpm-race-id"},
|
||||
"model_name": "opus",
|
||||
}
|
||||
|
||||
set_io_token_rate_limit_request_kwargs(
|
||||
{
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": max_tokens,
|
||||
"metadata": {},
|
||||
}
|
||||
)
|
||||
|
||||
async def _attempt():
|
||||
try:
|
||||
await check.async_pre_call_check(deployment)
|
||||
return True
|
||||
except litellm.RateLimitError:
|
||||
return False
|
||||
|
||||
results = await asyncio.gather(*[_attempt() for _ in range(8)])
|
||||
successes = sum(1 for r in results if r)
|
||||
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}"
|
||||
current_otpm = await dual_cache.async_get_cache(key=otpm_key)
|
||||
|
||||
# Atomic reservation must never let concurrent requests overshoot the limit.
|
||||
assert current_otpm is not None
|
||||
assert current_otpm <= otpm_limit
|
||||
assert successes == otpm_limit // max_tokens
|
||||
assert current_otpm == successes * max_tokens
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_itpm_estimate_failure_reserves_minimal_not_full_limit(self):
|
||||
"""
|
||||
When input-token estimation yields 0 (no messages/prompt/input field,
|
||||
unsupported model, tokenizer error), the reservation must be a
|
||||
minimal 1 token, not the entire itpm limit. Otherwise the first
|
||||
request whose estimate fails fills the whole bucket and every
|
||||
concurrent request is rejected until it completes - effectively
|
||||
serializing traffic to the deployment.
|
||||
"""
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
itpm_limit = 5
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"itpm": itpm_limit,
|
||||
},
|
||||
"model_info": {"id": "io-itpm-estimate-fail-id"},
|
||||
"model_name": "opus",
|
||||
}
|
||||
|
||||
# No messages/prompt/input field -> _estimate_input_tokens returns 0.
|
||||
set_io_token_rate_limit_request_kwargs(
|
||||
{
|
||||
"max_tokens": 5,
|
||||
"metadata": {},
|
||||
}
|
||||
)
|
||||
|
||||
async def _attempt():
|
||||
try:
|
||||
await check.async_pre_call_check(deployment)
|
||||
return True
|
||||
except litellm.RateLimitError:
|
||||
return False
|
||||
|
||||
results = await asyncio.gather(*[_attempt() for _ in range(8)])
|
||||
successes = sum(1 for r in results if r)
|
||||
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}"
|
||||
current_itpm = await dual_cache.async_get_cache(key=itpm_key)
|
||||
|
||||
# A minimal 1-token reservation per request lets itpm_limit concurrent
|
||||
# requests through, instead of a single request starving the rest.
|
||||
assert current_itpm is not None
|
||||
assert current_itpm <= itpm_limit
|
||||
assert successes == itpm_limit
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self):
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}"
|
||||
await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60)
|
||||
|
||||
# Production kwargs commonly carry litellm_params.metadata; the stashed
|
||||
# reservation lives in the top-level metadata and must still be found.
|
||||
kwargs = {
|
||||
"standard_logging_object": {
|
||||
"model_id": "io-lp-id",
|
||||
"hidden_params": {"litellm_model_name": "bedrock_mantle/test"},
|
||||
"metadata": {},
|
||||
},
|
||||
"metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key},
|
||||
"litellm_params": {"metadata": {"user_api_key_hash": "abc123"}},
|
||||
}
|
||||
await check.async_log_failure_event(kwargs, None, None, None)
|
||||
|
||||
current = await dual_cache.async_get_cache(key=itpm_key)
|
||||
assert current == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_tracks_actual_usage_when_estimate_zero(self):
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"itpm": 100,
|
||||
},
|
||||
"model_info": {"id": "io-zero-est-id"},
|
||||
"model_name": "opus",
|
||||
}
|
||||
|
||||
request_kwargs = {"max_tokens": 5, "metadata": {}}
|
||||
set_io_token_rate_limit_request_kwargs(request_kwargs)
|
||||
await check.async_pre_call_check(deployment)
|
||||
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}"
|
||||
# A failed/zero estimate reserves a minimal 1 token, not the full
|
||||
# itpm limit, so it doesn't starve concurrent requests.
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 1
|
||||
|
||||
kwargs = {
|
||||
"standard_logging_object": {
|
||||
"model_id": "io-zero-est-id",
|
||||
"hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"},
|
||||
"metadata": dict(request_kwargs["metadata"]),
|
||||
},
|
||||
"metadata": request_kwargs["metadata"],
|
||||
}
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage=Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7),
|
||||
)
|
||||
await check.async_log_success_event(kwargs, response, None, None)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 7
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self):
|
||||
"""
|
||||
A zero/failed estimate reserves a minimal 1 token rather than the
|
||||
full itpm limit, so up to itpm_limit such calls are allowed
|
||||
concurrently instead of the first one claiming the entire bucket.
|
||||
"""
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"itpm": 2,
|
||||
},
|
||||
"model_info": {"id": "io-zero-cap-id"},
|
||||
"model_name": "opus",
|
||||
}
|
||||
request_kwargs = {"max_tokens": 5, "metadata": {}}
|
||||
set_io_token_rate_limit_request_kwargs(request_kwargs)
|
||||
await check.async_pre_call_check(deployment)
|
||||
|
||||
# Second zero-estimate call still fits within the itpm=2 limit.
|
||||
set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}})
|
||||
await check.async_pre_call_check(deployment)
|
||||
|
||||
# A third exceeds the limit and is rejected.
|
||||
set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}})
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await check.async_pre_call_check(deployment)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self):
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"otpm": 5,
|
||||
},
|
||||
"model_info": {"id": "io-zero-output-id"},
|
||||
"model_name": "opus",
|
||||
}
|
||||
zero_output_kwargs = {
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 0,
|
||||
"metadata": {},
|
||||
}
|
||||
set_io_token_rate_limit_request_kwargs(zero_output_kwargs)
|
||||
await check.async_pre_call_check(deployment)
|
||||
|
||||
zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY]
|
||||
assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0
|
||||
assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0
|
||||
|
||||
normal_output_kwargs = {
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 5,
|
||||
"metadata": {},
|
||||
}
|
||||
set_io_token_rate_limit_request_kwargs(normal_output_kwargs)
|
||||
await check.async_pre_call_check(deployment)
|
||||
|
||||
normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY]
|
||||
assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5
|
||||
|
||||
def test_sync_io_pre_call_reserves_and_reconciles(self):
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"itpm": 100,
|
||||
"otpm": 50,
|
||||
},
|
||||
"model_info": {"id": "io-sync-id"},
|
||||
"model_name": "opus",
|
||||
}
|
||||
request_kwargs = {
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 10,
|
||||
"metadata": {},
|
||||
}
|
||||
set_io_token_rate_limit_request_kwargs(request_kwargs)
|
||||
check.pre_call_check(deployment)
|
||||
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}"
|
||||
otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}"
|
||||
kwargs = {
|
||||
"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": request_kwargs["metadata"],
|
||||
}
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8),
|
||||
)
|
||||
check.log_success_event(kwargs, response, None, None)
|
||||
|
||||
assert dual_cache.get_cache(key=itpm_key) == 5
|
||||
assert dual_cache.get_cache(key=otpm_key) == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_runs_via_success_event_without_model_id(self):
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34"
|
||||
await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60)
|
||||
|
||||
# standard_logging_object has no model_id (only the TPM path needs it);
|
||||
# IO reconciliation must still run off the stashed cache key.
|
||||
kwargs = {
|
||||
"standard_logging_object": {
|
||||
"hidden_params": {"litellm_model_name": "bedrock_mantle/test"},
|
||||
"metadata": {},
|
||||
},
|
||||
"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key},
|
||||
}
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3),
|
||||
)
|
||||
await check.async_log_success_event(kwargs, response, None, None)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_clears_reservation_so_retry_is_not_poisoned(self):
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}"
|
||||
await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60)
|
||||
|
||||
# Shared request metadata carrying the first (IO) deployment's reservation.
|
||||
metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key}
|
||||
fail_kwargs = {
|
||||
"metadata": metadata,
|
||||
"standard_logging_object": {
|
||||
"model_id": "io-first",
|
||||
"hidden_params": {"litellm_model_name": "bedrock_mantle/test"},
|
||||
"metadata": {},
|
||||
},
|
||||
}
|
||||
await check.async_log_failure_event(fail_kwargs, None, None, None)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 0
|
||||
assert ITPM_RESERVED_KEY not in metadata
|
||||
assert ITPM_CACHE_KEY not in metadata
|
||||
|
||||
# Retry succeeds on a non-IO fallback deployment reusing the same metadata.
|
||||
retry_kwargs = {
|
||||
"metadata": metadata,
|
||||
"standard_logging_object": {
|
||||
"model_id": "non-io-second",
|
||||
"hidden_params": {"litellm_model_name": "openai/gpt-4o-mini"},
|
||||
"metadata": {},
|
||||
"total_tokens": 12,
|
||||
},
|
||||
}
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage=Usage(prompt_tokens=6, completion_tokens=6, total_tokens=12),
|
||||
)
|
||||
await check.async_log_success_event(retry_kwargs, response, None, None)
|
||||
|
||||
# The first deployment's ITPM counter is not driven negative...
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 0
|
||||
# ...and the non-IO deployment's TPM usage is tracked normally.
|
||||
tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}"
|
||||
assert await dual_cache.async_get_cache(key=tpm_key) == 12
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_reservation_refunded_before_retry_overwrites_it(self):
|
||||
"""
|
||||
A retry reuses the same mutable kwargs dict for the next deployment.
|
||||
If deployment A's failure event hasn't run yet (e.g. it was scheduled
|
||||
as a background task) when the retry calls
|
||||
set_io_token_rate_limit_request_kwargs for deployment B, the router
|
||||
must first synchronously refund + clear A's reservation via
|
||||
refund_stale_reservation_before_retry - otherwise A's counter stays
|
||||
elevated by the reservation until its TTL expires, and the
|
||||
now-orphaned sentinels must not leak into B's accounting either.
|
||||
"""
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}"
|
||||
await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60)
|
||||
|
||||
# Deployment A's still-unreconciled reservation, stashed on the shared
|
||||
# kwargs dict the retry loop reuses.
|
||||
shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}}
|
||||
|
||||
# Router calls this before overwriting kwargs for deployment B's attempt -
|
||||
# simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs.
|
||||
refund_stale_reservation_before_retry(dual_cache, shared_kwargs)
|
||||
|
||||
# A's reservation is refunded immediately, not left stranded for a
|
||||
# background failure task that may run arbitrarily later (or never,
|
||||
# if the sentinels get cleared out from under it first).
|
||||
assert await dual_cache.async_get_cache(key=itpm_key_a) == 0
|
||||
assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"]
|
||||
assert ITPM_CACHE_KEY not in shared_kwargs["metadata"]
|
||||
|
||||
# A's own (now-late) failure event finds nothing left to refund and
|
||||
# is a safe no-op, since the sentinels were already cleared above.
|
||||
io_token_refund_failure(dual_cache, shared_kwargs)
|
||||
assert await dual_cache.async_get_cache(key=itpm_key_a) == 0
|
||||
|
||||
# The retry proceeds to stash deployment B's own reservation on the
|
||||
# same dict; it starts clean, unaffected by A's cleared sentinels.
|
||||
set_io_token_rate_limit_request_kwargs(shared_kwargs)
|
||||
itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}"
|
||||
shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4
|
||||
shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b
|
||||
await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60)
|
||||
assert await dual_cache.async_get_cache(key=itpm_key_b) == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_supplied_reservation_keys_are_stripped(self):
|
||||
# metadata is caller-controlled; the server-only reservation sentinels
|
||||
# must be removed before the router captures the request kwargs.
|
||||
forged = {
|
||||
"metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"},
|
||||
"litellm_metadata": {OTPM_RESERVED_KEY: 7},
|
||||
"litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}},
|
||||
}
|
||||
set_io_token_rate_limit_request_kwargs(forged)
|
||||
stored = get_io_token_rate_limit_request_kwargs()
|
||||
|
||||
assert ITPM_RESERVED_KEY not in stored["metadata"]
|
||||
assert ITPM_CACHE_KEY not in stored["metadata"]
|
||||
assert OTPM_RESERVED_KEY not in stored["litellm_metadata"]
|
||||
assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_forged_reservation_cannot_decrement_counter(self):
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
victim_key = "global_router:victim:model:itpm:00-00"
|
||||
await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60)
|
||||
|
||||
# A caller forges a reservation pointing at another deployment's counter.
|
||||
kwargs = {
|
||||
"metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key},
|
||||
"standard_logging_object": {
|
||||
"model_id": "m",
|
||||
"hidden_params": {"litellm_model_name": "model"},
|
||||
"metadata": {},
|
||||
"total_tokens": 2,
|
||||
},
|
||||
}
|
||||
# The router sanitizes the request kwargs before the call runs.
|
||||
set_io_token_rate_limit_request_kwargs(kwargs)
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2),
|
||||
)
|
||||
await check.async_log_success_event(kwargs, response, None, None)
|
||||
|
||||
# The forged reservation was stripped, so the victim counter is untouched.
|
||||
assert await dual_cache.async_get_cache(key=victim_key) == 100
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_otpm_reservation_error_rolls_back_itpm(self):
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
class _OtpmFailCache(DualCache):
|
||||
async def async_increment_cache(self, key, **kwargs):
|
||||
if ":otpm:" in key:
|
||||
raise RuntimeError("transient cache error")
|
||||
return await super().async_increment_cache(key=key, **kwargs)
|
||||
|
||||
dual_cache = _OtpmFailCache()
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"itpm": 1000,
|
||||
"otpm": 1000,
|
||||
},
|
||||
"model_info": {"id": "io-rollback-id"},
|
||||
"model_name": "opus",
|
||||
}
|
||||
set_io_token_rate_limit_request_kwargs(
|
||||
{
|
||||
"messages": [{"role": "user", "content": "hello world"}],
|
||||
"max_tokens": 5,
|
||||
"metadata": {},
|
||||
}
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
await async_io_token_pre_call_check(dual_cache, deployment)
|
||||
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}"
|
||||
# A transient OTPM error must release the ITPM reservation, not leak it.
|
||||
assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_clears_stash_even_when_increment_errors(self):
|
||||
class _ItpmFailCache(DualCache):
|
||||
async def async_increment_cache(self, key, **kwargs):
|
||||
if ":itpm:" in key:
|
||||
raise RuntimeError("transient cache error")
|
||||
return await super().async_increment_cache(key=key, **kwargs)
|
||||
|
||||
dual_cache = _ItpmFailCache()
|
||||
metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"}
|
||||
kwargs = {"metadata": metadata}
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
await async_io_token_reconcile_success(dual_cache, kwargs, response)
|
||||
|
||||
# The stash is cleared even though reconciliation raised, so a duplicate
|
||||
# success event can't re-process it.
|
||||
assert ITPM_RESERVED_KEY not in metadata
|
||||
assert ITPM_CACHE_KEY not in metadata
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog):
|
||||
import logging
|
||||
|
||||
check = ModelRateLimitingCheck(dual_cache=DualCache())
|
||||
deployment = {
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000},
|
||||
"model_info": {},
|
||||
}
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
|
||||
check._warn_io_token_and_tpm_rpm_coexist_once(deployment)
|
||||
check._warn_io_token_and_tpm_rpm_coexist_once(deployment)
|
||||
|
||||
warnings = [r for r in caplog.records if "both limit types are enforced" in r.message]
|
||||
# id-less deployments are not collapsed onto a single dedup key.
|
||||
assert len(warnings) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_deployment_id_skips_io_reservation(self):
|
||||
dual_cache = DualCache()
|
||||
deployment = {
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100},
|
||||
"model_info": {}, # no id -> cannot build a per-deployment cache key
|
||||
"model_name": "opus",
|
||||
}
|
||||
request_kwargs = {
|
||||
"messages": [{"role": "user", "content": "hello world"}],
|
||||
"max_tokens": 5,
|
||||
"metadata": {},
|
||||
}
|
||||
set_io_token_rate_limit_request_kwargs(request_kwargs)
|
||||
|
||||
result = await async_io_token_pre_call_check(dual_cache, deployment)
|
||||
|
||||
assert result is deployment
|
||||
# No reservation is stashed, so nothing lands in a shared None:None bucket.
|
||||
assert ITPM_RESERVED_KEY not in request_kwargs["metadata"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_uses_reservation_minute_key(self):
|
||||
dual_cache = DualCache()
|
||||
# Reservation was made on a fixed minute key; a call that finishes in a
|
||||
# later minute must reconcile against that same key, never a key built
|
||||
# from the response-time minute.
|
||||
itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99"
|
||||
await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60)
|
||||
|
||||
kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}}
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4),
|
||||
)
|
||||
await async_io_token_reconcile_success(dual_cache, kwargs, response)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_missing_usage_keeps_reservation(self):
|
||||
dual_cache = DualCache()
|
||||
itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00"
|
||||
otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00"
|
||||
await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60)
|
||||
await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60)
|
||||
|
||||
kwargs = {
|
||||
"metadata": {
|
||||
ITPM_RESERVED_KEY: 8,
|
||||
OTPM_RESERVED_KEY: 5,
|
||||
ITPM_CACHE_KEY: itpm_key,
|
||||
OTPM_CACHE_KEY: otpm_key,
|
||||
}
|
||||
}
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}],
|
||||
)
|
||||
|
||||
await async_io_token_reconcile_success(dual_cache, kwargs, response)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 8
|
||||
assert await dual_cache.async_get_cache(key=otpm_key) == 5
|
||||
assert ITPM_RESERVED_KEY not in kwargs["metadata"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_total_tokens_only_keeps_reservation(self):
|
||||
"""
|
||||
A response usage object with only total_tokens (no prompt/completion
|
||||
breakdown) can't be split into input/output, so it must be treated the
|
||||
same as missing usage: keep the reservation instead of resolving to
|
||||
(0, 0) and refunding it in full.
|
||||
"""
|
||||
dual_cache = DualCache()
|
||||
itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00"
|
||||
otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00"
|
||||
await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60)
|
||||
await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60)
|
||||
|
||||
kwargs = {
|
||||
"metadata": {
|
||||
ITPM_RESERVED_KEY: 8,
|
||||
OTPM_RESERVED_KEY: 5,
|
||||
ITPM_CACHE_KEY: itpm_key,
|
||||
OTPM_CACHE_KEY: otpm_key,
|
||||
}
|
||||
}
|
||||
response = {"type": "message", "usage": {"total_tokens": 13}}
|
||||
|
||||
await async_io_token_reconcile_success(dual_cache, kwargs, response)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 8
|
||||
assert await dual_cache.async_get_cache(key=otpm_key) == 5
|
||||
|
||||
def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self):
|
||||
dual_cache = DualCache()
|
||||
itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00"
|
||||
dual_cache.set_cache(key=itpm_key, value=10, ttl=60)
|
||||
|
||||
kwargs = {
|
||||
"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key},
|
||||
"standard_logging_object": {"total_tokens": 4},
|
||||
}
|
||||
response = {"type": "message", "role": "assistant", "content": []}
|
||||
|
||||
io_token_reconcile_success(dual_cache, kwargs, response)
|
||||
|
||||
assert dual_cache.get_cache(key=itpm_key) == 10
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_falls_back_to_standard_logging_object(self):
|
||||
dual_cache = DualCache()
|
||||
itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00"
|
||||
await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60)
|
||||
|
||||
kwargs = {
|
||||
"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key},
|
||||
"standard_logging_object": {
|
||||
"prompt_tokens": 4,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 4,
|
||||
},
|
||||
}
|
||||
response = {"type": "message", "role": "assistant", "content": []}
|
||||
|
||||
await async_io_token_reconcile_success(dual_cache, kwargs, response)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 4
|
||||
|
||||
def test_sync_reconcile_anthropic_dict_usage(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=otpm_key, value=4, ttl=60)
|
||||
|
||||
kwargs = {
|
||||
"metadata": {
|
||||
ITPM_RESERVED_KEY: 6,
|
||||
OTPM_RESERVED_KEY: 4,
|
||||
ITPM_CACHE_KEY: itpm_key,
|
||||
OTPM_CACHE_KEY: otpm_key,
|
||||
}
|
||||
}
|
||||
response = {
|
||||
"type": "message",
|
||||
"usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1},
|
||||
}
|
||||
|
||||
io_token_reconcile_success(dual_cache, kwargs, response)
|
||||
|
||||
assert dual_cache.get_cache(key=itpm_key) == 2
|
||||
assert dual_cache.get_cache(key=otpm_key) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog):
|
||||
import logging
|
||||
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
model_id = "io-mixed-id"
|
||||
deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7"
|
||||
deployment = {
|
||||
"litellm_params": {
|
||||
"model": deployment_name,
|
||||
"itpm": 100,
|
||||
"rpm": 1,
|
||||
},
|
||||
"model_info": {"id": model_id},
|
||||
"model_name": "opus",
|
||||
}
|
||||
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}"
|
||||
itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}"
|
||||
await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60)
|
||||
|
||||
request_kwargs = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 5,
|
||||
"metadata": {},
|
||||
}
|
||||
set_io_token_rate_limit_request_kwargs(request_kwargs)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await check.async_pre_call_check(deployment)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=rpm_key) == 6
|
||||
assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0
|
||||
assert ITPM_RESERVED_KEY not in request_kwargs["metadata"]
|
||||
assert any("both limit types are enforced" in record.message for record in caplog.records)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_io_success_still_tracks_tpm_for_mixed_deployment(self):
|
||||
"""
|
||||
A deployment with itpm/otpm AND tpm/rpm must have BOTH counters updated on
|
||||
success, otherwise the tpm_key the pre-call check reads is never written
|
||||
and the tpm_limit can never be enforced.
|
||||
"""
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
model_id = "io-tpm-mixed-id"
|
||||
deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7"
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}"
|
||||
tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}"
|
||||
await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60)
|
||||
|
||||
kwargs = {
|
||||
"metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key},
|
||||
"standard_logging_object": {
|
||||
"model_id": model_id,
|
||||
"total_tokens": 7,
|
||||
"hidden_params": {"litellm_model_name": deployment_name},
|
||||
},
|
||||
}
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3),
|
||||
)
|
||||
|
||||
await check.async_log_success_event(kwargs, response, None, None)
|
||||
|
||||
# ITPM reconciled down from the 5-token reservation to actual usage (3).
|
||||
assert await dual_cache.async_get_cache(key=itpm_key) == 3
|
||||
# TPM tracking must still run so the tpm/rpm pre-call path can enforce it.
|
||||
assert await dual_cache.async_get_cache(key=tpm_key) == 7
|
||||
|
||||
def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self):
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
model_id = "io-tpm-mixed-sync-id"
|
||||
deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7"
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}"
|
||||
tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}"
|
||||
dual_cache.set_cache(key=itpm_key, value=5, ttl=60)
|
||||
|
||||
kwargs = {
|
||||
"metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key},
|
||||
"standard_logging_object": {
|
||||
"model_id": model_id,
|
||||
"total_tokens": 7,
|
||||
"hidden_params": {"litellm_model_name": deployment_name},
|
||||
},
|
||||
}
|
||||
response = ModelResponse(
|
||||
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3),
|
||||
)
|
||||
|
||||
check.log_success_event(kwargs, response, None, None)
|
||||
|
||||
assert dual_cache.get_cache(key=itpm_key) == 3
|
||||
assert dual_cache.get_cache(key=tpm_key) == 7
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_refunds_itpm_reservation(self):
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
dual_cache = DualCache()
|
||||
check = ModelRateLimitingCheck(dual_cache=dual_cache)
|
||||
minute = get_utc_datetime().strftime("%H-%M")
|
||||
itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}"
|
||||
await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60)
|
||||
|
||||
reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}
|
||||
kwargs = {
|
||||
"standard_logging_object": {
|
||||
"model_id": "io-refund-id",
|
||||
"hidden_params": {"litellm_model_name": "bedrock_mantle/test"},
|
||||
"metadata": dict(reservation),
|
||||
},
|
||||
"metadata": dict(reservation),
|
||||
}
|
||||
await check.async_log_failure_event(kwargs, None, None, None)
|
||||
|
||||
current = await dual_cache.async_get_cache(key=itpm_key)
|
||||
assert current == 0
|
||||
|
||||
|
||||
class TestRouterIOTokenIntegration:
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_group_info_aggregates_io_limits(self):
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "opus",
|
||||
"litellm_params": {
|
||||
"model": "bedrock_mantle/anthropic.claude-opus-4-7",
|
||||
"itpm": 100,
|
||||
"otpm": 20,
|
||||
},
|
||||
}
|
||||
],
|
||||
optional_pre_call_checks=["enforce_model_rate_limits"],
|
||||
)
|
||||
info = router.get_model_group_info("opus")
|
||||
assert info is not None
|
||||
assert info.itpm == 100
|
||||
assert info.otpm == 20
|
||||
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -25430,6 +25430,8 @@ export interface components {
|
|||
input_cost_per_video_per_second_above_15s_interval?: number | null;
|
||||
/** Input Cost Per Video Per Second Above 8S Interval */
|
||||
input_cost_per_video_per_second_above_8s_interval?: number | null;
|
||||
/** Itpm */
|
||||
itpm?: number | null;
|
||||
/** Litellm Credential Name */
|
||||
litellm_credential_name?: string | null;
|
||||
/** Litellm Trace Id */
|
||||
|
|
@ -25465,6 +25467,8 @@ export interface components {
|
|||
ocr_cost_per_page?: number | null;
|
||||
/** Organization */
|
||||
organization?: string | null;
|
||||
/** Otpm */
|
||||
otpm?: number | null;
|
||||
/** Output Cost Per Audio Per Second */
|
||||
output_cost_per_audio_per_second?: number | null;
|
||||
/** Output Cost Per Audio Token */
|
||||
|
|
@ -27261,6 +27265,8 @@ export interface components {
|
|||
* @default false
|
||||
*/
|
||||
is_public_model_group: boolean;
|
||||
/** Itpm */
|
||||
itpm?: number | null;
|
||||
/** Max Input Tokens */
|
||||
max_input_tokens?: number | null;
|
||||
/** Max Output Tokens */
|
||||
|
|
@ -27272,6 +27278,8 @@ export interface components {
|
|||
mode: string | ("chat" | "embedding" | "completion" | "image_generation" | "audio_transcription" | "rerank" | "moderations") | null;
|
||||
/** Model Group */
|
||||
model_group: string;
|
||||
/** Otpm */
|
||||
otpm?: number | null;
|
||||
/** Output Cost Per Token */
|
||||
output_cost_per_token?: number | null;
|
||||
/** Providers */
|
||||
|
|
@ -33216,6 +33224,8 @@ export interface components {
|
|||
input_cost_per_video_per_second_above_15s_interval?: number | null;
|
||||
/** Input Cost Per Video Per Second Above 8S Interval */
|
||||
input_cost_per_video_per_second_above_8s_interval?: number | null;
|
||||
/** Itpm */
|
||||
itpm?: number | null;
|
||||
/** Litellm Credential Name */
|
||||
litellm_credential_name?: string | null;
|
||||
/** Litellm Trace Id */
|
||||
|
|
@ -33251,6 +33261,8 @@ export interface components {
|
|||
ocr_cost_per_page?: number | null;
|
||||
/** Organization */
|
||||
organization?: string | null;
|
||||
/** Otpm */
|
||||
otpm?: number | null;
|
||||
/** Output Cost Per Audio Per Second */
|
||||
output_cost_per_audio_per_second?: number | null;
|
||||
/** Output Cost Per Audio Token */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue