diff --git a/litellm/router.py b/litellm/router.py index a0ff3e0fd31..e0ed9b85ec1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -348,9 +348,13 @@ class Router: provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None. deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. prompt_prefix_affinity_tokens (int): Number of canonical prompt-prefix - tokens used for deterministic prompt-prefix routing. Defaults to 2048. + tokens used for deterministic prompt-prefix routing. This is + measured against tokenization of the canonical JSON representation + (not the provider-specific "raw prompt token count"). Defaults to 2048. prompt_prefix_affinity_min_tokens (int): Minimum canonical prompt token - count before prompt-prefix affinity applies. Defaults to + count before prompt-prefix affinity applies. This is measured + against tokenization of the canonical JSON representation (not the + provider-specific "raw prompt token count"). Defaults to MINIMUM_PROMPT_CACHE_TOKEN_COUNT. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. Returns: diff --git a/litellm/router_utils/pre_call_checks/prompt_prefix_affinity_check.py b/litellm/router_utils/pre_call_checks/prompt_prefix_affinity_check.py index 5fddc8a4c8d..d41a80bda17 100644 --- a/litellm/router_utils/pre_call_checks/prompt_prefix_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_prefix_affinity_check.py @@ -11,6 +11,7 @@ mapping in Redis. All Router instances with the same config and deployment IDs will make the same routing decision. """ +import asyncio import hashlib import json from typing import Any, Dict, List, Optional, Tuple, cast @@ -28,9 +29,16 @@ class PromptPrefixAffinityCheck(CustomLogger): This is intended to improve prompt-cache hit rate for providers where prompt caching is scoped to the account/key behind a deployment. + + Note: `prefix_tokens` / `min_tokens` are measured against the tokenization of + the canonical JSON representation produced by `_build_canonical_prompt()` + (e.g. includes field names, separators, quotes), not the provider-specific + "raw prompt token count" after rendering the final prompt. """ CACHE_KEY_EXCLUDED_FIELDS = frozenset({"encrypted_content"}) + _CANONICAL_PROMPT_CHARS_PER_TOKEN_ESTIMATE = 32 + _CANONICAL_PROMPT_ABSOLUTE_MAX_CHARS = 1_000_000 def __init__( self, @@ -95,6 +103,26 @@ class PromptPrefixAffinityCheck(CustomLogger): ensure_ascii=False, ) + def _truncate_canonical_prompt_for_tokenization(self, canonical_prompt: str) -> str: + """ + Bound tokenization work for very large prompts. + + We only use the first `prefix_tokens` token IDs (and gate on `min_tokens`), + so we can safely cap the number of characters we feed into tokenization to + avoid worst-case CPU time on extremely large canonical prompts. + """ + + token_target = max(self.prefix_tokens, self.min_tokens) + if token_target <= 0: + return canonical_prompt + + soft_limit = token_target * self._CANONICAL_PROMPT_CHARS_PER_TOKEN_ESTIMATE + max_chars = min(soft_limit, self._CANONICAL_PROMPT_ABSOLUTE_MAX_CHARS) + if len(canonical_prompt) <= max_chars: + return canonical_prompt + + return canonical_prompt[:max_chars] + def _get_prefix_hash( self, model: str, @@ -182,9 +210,13 @@ class PromptPrefixAffinityCheck(CustomLogger): if canonical_prompt is None: return typed_healthy_deployments - prefix_hash = self._get_prefix_hash( + canonical_prompt_for_tokenization = self._truncate_canonical_prompt_for_tokenization( + canonical_prompt + ) + prefix_hash = await asyncio.to_thread( + self._get_prefix_hash, model=model, - canonical_prompt=canonical_prompt, + canonical_prompt=canonical_prompt_for_tokenization, ) if prefix_hash is None: return typed_healthy_deployments