diff --git a/litellm/constants.py b/litellm/constants.py index 1014b472c61..e98107f8da4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -279,6 +279,7 @@ MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = int(os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT # escape hatch when `MINIMUM_PROMPT_CACHE_TOKEN_COUNT` is explicitly set. MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE: int | None = get_env_int_or_none("MINIMUM_PROMPT_CACHE_TOKEN_COUNT") DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT = 1024 +DEFAULT_CHARS_PER_TOKEN = 4 MINIMUM_PROMPT_CACHE_TOKEN_COUNT = ( MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None @@ -1451,6 +1452,7 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_R SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job" EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job" +CACHE_WARMING_JOB_NAME = "complexity_router_cache_warming" SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index faedf8ae1a3..d1c6e085380 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -34,6 +34,8 @@ else: # breakpoints: "A maximum of 4 blocks with cache_control may be provided." MAX_CACHE_CONTROL_BLOCKS = 4 +EXPLICIT_PROMPT_CACHING_PROVIDERS = frozenset(("anthropic", "bedrock")) + class AnthropicCacheControlHook(CustomPromptManagement): def get_chat_completion_prompt( @@ -406,7 +408,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): except Exception: # noqa: BLE001 # unroutable model must never block the call, just skip auto-caching return [] - if provider not in ("anthropic", "bedrock"): + if provider not in EXPLICIT_PROMPT_CACHING_PROVIDERS: return [] from litellm.utils import supports_prompt_caching diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index cecc35ee1c1..8fdca178388 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -1,6 +1,7 @@ # What is this? ## Helper utilities import copy +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union import httpx @@ -214,6 +215,29 @@ def get_or_create_metadata_bucket( return metadata_key, metadata_bucket +def iter_request_metadata_dicts(request_kwargs: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + """Every metadata dict present at the TOP level of request kwargs, litellm_metadata first. + + Distinct from ``get_litellm_metadata_from_kwargs``, which reads the nested + ``kwargs["litellm_params"]`` level and returns one merged dict; callers that must inspect both + slots as written (routing affinity, cache warming) need them separately and unmerged. + """ + return tuple( # pyright: ignore[reportUnknownVariableType] # isinstance narrows object to dict[Unknown, Unknown] + metadata # pyright: ignore[reportUnknownArgumentType] # request metadata keys are runtime strings + for metadata_key in ("litellm_metadata", "metadata") + if isinstance(metadata := request_kwargs.get(metadata_key), dict) + ) + + +def get_request_metadata_field(request_kwargs: Mapping[str, object], field: str) -> str | None: + """First stringified value of ``field`` across the top-level metadata slots.""" + for metadata in iter_request_metadata_dicts(request_kwargs): + value = metadata.get(field) + if value is not None: + return str(value) + return None + + def get_litellm_metadata_from_kwargs(kwargs: dict): """ Helper to get litellm metadata from all litellm request kwargs diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0472b496b78..c8771528afa 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -498,7 +498,7 @@ async def common_checks( llm_router: Optional[Router], proxy_logging_obj: ProxyLogging, valid_token: Optional[UserAPIKeyAuth], - request: Request, + request: Optional[Request], skip_budget_checks: bool = False, project_object: Optional[LiteLLM_ProjectTableCachedObj] = None, ) -> bool: @@ -2509,11 +2509,17 @@ async def get_key_object( parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, check_cache_only: Optional[bool] = None, + check_db_only: Optional[bool] = None, ) -> UserAPIKeyAuth: """ - Check if team id in proxy Team Table - if valid, return LiteLLM_TeamTable object with defined limits - if not, then raise an error + + ``check_db_only`` skips the cache read, mirroring the flag get_team_object and get_user_object already + carry. A caller that must not act on a revoked key needs the authoritative row: the cache is written by + every worker and a key blocked, expired or deleted seconds ago still resolves to valid state from another + worker's entry until it expires. """ if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") @@ -2523,9 +2529,13 @@ async def get_key_object( # Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth # (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB. - user_api_key_auth = await user_api_key_cache.async_get_cache( - key=key, - model_type=UserAPIKeyAuth, + user_api_key_auth = ( + None + if check_db_only + else await user_api_key_cache.async_get_cache( + key=key, + model_type=UserAPIKeyAuth, + ) ) if user_api_key_auth is not None: return _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 719496785dd..b984ead2881 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -26,7 +26,7 @@ from typing import ( from litellm import DualCache from litellm._logging import verbose_proxy_logger -from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE +from litellm.constants import DEFAULT_CHARS_PER_TOKEN, DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -285,7 +285,6 @@ REDIS_NODE_HASHTAG_NAME = "all_keys" # When max_tokens is not specified in the request we still need to reserve # *some* output budget; these define that fallback estimate. DEFAULT_MAX_TOKENS_ESTIMATE = 4096 -DEFAULT_CHARS_PER_TOKEN = 4 # Fraction of the available output budget reserved as the upfront floor when # the request omits max_tokens. Applied to both DEFAULT_MAX_TOKENS_ESTIMATE # (baseline floor) and to the smallest configured TPM limit (capped floor for diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 673e73f72fb..4e784d48c46 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1140,6 +1140,125 @@ class LiteLLMProxyRequestSetup: ) return data + @staticmethod + def add_key_team_project_metadata( + data: dict[str, Any], + user_api_key_dict: UserAPIKeyAuth, + _metadata_variable_name: str, + ) -> dict[str, Any]: + """Every key, team and project control derived purely from the resolved identity: cache + controls, merged tags, spend_logs_metadata, guardrail opt-outs, disable_fallbacks, the + management-endpoint metadata subset, and the budget fields Prometheus reads. Request-free on + purpose, so a caller without a FastAPI Request (a cache-warming replay) reaches the same block + instead of re-deriving it and drifting from spend attribution.""" + ### KEY-LEVEL Controls + key_metadata = user_api_key_dict.metadata + data = LiteLLMProxyRequestSetup.add_key_level_controls( + key_metadata=key_metadata, + data=data, + _metadata_variable_name=_metadata_variable_name, + ) + ## TEAM-LEVEL SPEND LOGS/TAGS + team_metadata = user_api_key_dict.team_metadata or {} + if "tags" in team_metadata and team_metadata["tags"] is not None: + data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=team_metadata["tags"], + ) + if "disable_global_guardrails" in team_metadata and isinstance( + team_metadata["disable_global_guardrails"], bool + ): + data[_metadata_variable_name]["disable_global_guardrails"] = team_metadata["disable_global_guardrails"] + if "opted_out_global_guardrails" in team_metadata and isinstance( + team_metadata["opted_out_global_guardrails"], list + ): + data[_metadata_variable_name]["opted_out_global_guardrails"] = team_metadata["opted_out_global_guardrails"] + if "spend_logs_metadata" in team_metadata and isinstance(team_metadata["spend_logs_metadata"], dict): + if "spend_logs_metadata" in data[_metadata_variable_name] and isinstance( + data[_metadata_variable_name]["spend_logs_metadata"], dict + ): + for key, value in team_metadata["spend_logs_metadata"].items(): + if ( + key not in data[_metadata_variable_name]["spend_logs_metadata"] + ): # don't override k-v pair sent by request (user request) + data[_metadata_variable_name]["spend_logs_metadata"][key] = value + else: + data[_metadata_variable_name]["spend_logs_metadata"] = team_metadata["spend_logs_metadata"] + + ## PROJECT-LEVEL TAGS + project_metadata = user_api_key_dict.project_metadata or {} + if "tags" in project_metadata and project_metadata["tags"] is not None: + data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=project_metadata["tags"], + ) + + ## TEAM-LEVEL METADATA + data = LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( + data=data, + management_endpoint_metadata=team_metadata, + _metadata_variable_name=_metadata_variable_name, + ) + + # Team spend, budget - used by prometheus.py + data[_metadata_variable_name]["user_api_key_team_max_budget"] = user_api_key_dict.team_max_budget + data[_metadata_variable_name]["user_api_key_team_spend"] = user_api_key_dict.team_spend + data[_metadata_variable_name]["user_api_key_request_route"] = user_api_key_dict.request_route + + # API Key spend, budget - used by prometheus.py + data[_metadata_variable_name]["user_api_key_spend"] = user_api_key_dict.spend + data[_metadata_variable_name]["user_api_key_max_budget"] = user_api_key_dict.max_budget + data[_metadata_variable_name]["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget + data[_metadata_variable_name]["user_api_key_end_user_model_max_budget"] = ( + user_api_key_dict.end_user_model_max_budget + ) + + # User spend, budget - used by prometheus.py + # Follow same pattern as team and API key budgets + data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend + data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget + + data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) + data[_metadata_variable_name]["user_api_key_team_metadata"] = strip_callback_config( + user_api_key_dict.team_metadata + ) + data[_metadata_variable_name]["user_api_key_object_permission_id"] = getattr( + user_api_key_dict, "object_permission_id", None + ) + data[_metadata_variable_name]["user_api_key_team_object_permission_id"] = getattr( + user_api_key_dict, "team_object_permission_id", None + ) + return data + + @staticmethod + def apply_dynamic_logging_settings( + data: dict[str, Any], + user_api_key_dict: UserAPIKeyAuth, + proxy_config: ProxyConfig, + ) -> dict[str, Any]: + """Key and team scoped success/failure callbacks with their callback_vars, plus the key's + disabled-callback list. Request-free for the same reason as add_key_team_project_metadata: + without it a request never reaches the logger its own team configured.""" + # Team Callbacks controls + callback_settings_obj = _get_dynamic_logging_metadata( + user_api_key_dict=user_api_key_dict, proxy_config=proxy_config + ) + if callback_settings_obj is not None: + data["success_callback"] = callback_settings_obj.success_callback + data["failure_callback"] = callback_settings_obj.failure_callback + + if callback_settings_obj.callback_vars is not None: + # unpack callback_vars in data + for k, v in callback_settings_obj.callback_vars.items(): + data[k] = v + + # Add disabled callbacks from key metadata + if user_api_key_dict.metadata and "litellm_disabled_callbacks" in user_api_key_dict.metadata: + disabled_callbacks = user_api_key_dict.metadata["litellm_disabled_callbacks"] + if disabled_callbacks and isinstance(disabled_callbacks, list): + data["litellm_disabled_callbacks"] = disabled_callbacks + return data + @staticmethod def _merge_tags(request_tags: Optional[list], tags_to_add: Optional[list]) -> list: """ @@ -1608,79 +1727,11 @@ async def add_litellm_data_to_request( "global_max_parallel_requests", None ) - ### KEY-LEVEL Controls - key_metadata = user_api_key_dict.metadata - data = LiteLLMProxyRequestSetup.add_key_level_controls( - key_metadata=key_metadata, + data = LiteLLMProxyRequestSetup.add_key_team_project_metadata( data=data, + user_api_key_dict=user_api_key_dict, _metadata_variable_name=_metadata_variable_name, ) - ## TEAM-LEVEL SPEND LOGS/TAGS - team_metadata = user_api_key_dict.team_metadata or {} - if "tags" in team_metadata and team_metadata["tags"] is not None: - data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=team_metadata["tags"], - ) - if "disable_global_guardrails" in team_metadata and isinstance(team_metadata["disable_global_guardrails"], bool): - data[_metadata_variable_name]["disable_global_guardrails"] = team_metadata["disable_global_guardrails"] - if "opted_out_global_guardrails" in team_metadata and isinstance( - team_metadata["opted_out_global_guardrails"], list - ): - data[_metadata_variable_name]["opted_out_global_guardrails"] = team_metadata["opted_out_global_guardrails"] - if "spend_logs_metadata" in team_metadata and isinstance(team_metadata["spend_logs_metadata"], dict): - if "spend_logs_metadata" in data[_metadata_variable_name] and isinstance( - data[_metadata_variable_name]["spend_logs_metadata"], dict - ): - for key, value in team_metadata["spend_logs_metadata"].items(): - if ( - key not in data[_metadata_variable_name]["spend_logs_metadata"] - ): # don't override k-v pair sent by request (user request) - data[_metadata_variable_name]["spend_logs_metadata"][key] = value - else: - data[_metadata_variable_name]["spend_logs_metadata"] = team_metadata["spend_logs_metadata"] - - ## PROJECT-LEVEL TAGS - project_metadata = user_api_key_dict.project_metadata or {} - if "tags" in project_metadata and project_metadata["tags"] is not None: - data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=project_metadata["tags"], - ) - - ## TEAM-LEVEL METADATA - data = LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( - data=data, - management_endpoint_metadata=team_metadata, - _metadata_variable_name=_metadata_variable_name, - ) - - # Team spend, budget - used by prometheus.py - data[_metadata_variable_name]["user_api_key_team_max_budget"] = user_api_key_dict.team_max_budget - data[_metadata_variable_name]["user_api_key_team_spend"] = user_api_key_dict.team_spend - data[_metadata_variable_name]["user_api_key_request_route"] = user_api_key_dict.request_route - - # API Key spend, budget - used by prometheus.py - data[_metadata_variable_name]["user_api_key_spend"] = user_api_key_dict.spend - data[_metadata_variable_name]["user_api_key_max_budget"] = user_api_key_dict.max_budget - data[_metadata_variable_name]["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget - data[_metadata_variable_name]["user_api_key_end_user_model_max_budget"] = ( - user_api_key_dict.end_user_model_max_budget - ) - - # User spend, budget - used by prometheus.py - # Follow same pattern as team and API key budgets - data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend - data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget - - data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) - data[_metadata_variable_name]["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata) - data[_metadata_variable_name]["user_api_key_object_permission_id"] = getattr( - user_api_key_dict, "object_permission_id", None - ) - data[_metadata_variable_name]["user_api_key_team_object_permission_id"] = getattr( - user_api_key_dict, "team_object_permission_id", None - ) data[_metadata_variable_name]["headers"] = _headers data[_metadata_variable_name]["endpoint"] = str(request.url) # Carry the proxy-receive instant via metadata (like `endpoint`) so the @@ -1754,24 +1805,11 @@ async def add_litellm_data_to_request( tags_to_add=_user_tags, ) - # Team Callbacks controls - callback_settings_obj = _get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=proxy_config + data = LiteLLMProxyRequestSetup.apply_dynamic_logging_settings( + data=data, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, ) - if callback_settings_obj is not None: - data["success_callback"] = callback_settings_obj.success_callback - data["failure_callback"] = callback_settings_obj.failure_callback - - if callback_settings_obj.callback_vars is not None: - # unpack callback_vars in data - for k, v in callback_settings_obj.callback_vars.items(): - data[k] = v - - # Add disabled callbacks from key metadata - if user_api_key_dict.metadata and "litellm_disabled_callbacks" in user_api_key_dict.metadata: - disabled_callbacks = user_api_key_dict.metadata["litellm_disabled_callbacks"] - if disabled_callbacks and isinstance(disabled_callbacks, list): - data["litellm_disabled_callbacks"] = disabled_callbacks # Guardrails from key/team metadata and policy engine await move_guardrails_to_metadata( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c72e3d4ee5b..fab2e7028b3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -567,6 +567,7 @@ from litellm.router import ( LiteLLM_Params, ModelGroupInfo, ) +from litellm.router_strategy.complexity_router.cache_warming.refresher import CacheWarmingRefresher from litellm.scheduler import FlowItem, Scheduler from litellm.secret_managers.aws_secret_manager import load_aws_kms from litellm.secret_managers.google_kms import load_google_kms @@ -1106,6 +1107,7 @@ async def proxy_startup_event(app: FastAPI): await _tagged.strategy.load_state_from_db(prisma_client) _tagged.strategy._state_loaded = True asyncio.create_task(_adaptive_router_flusher_loop()) + asyncio.create_task(_complexity_cache_warming_loop()) ## [Optional] Initialize dd tracer ProxyStartupEvent._init_dd_tracer() @@ -3318,6 +3320,29 @@ async def _adaptive_router_flusher_loop(): verbose_proxy_logger.exception("adaptive_router flusher iteration failed") +_COMPLEXITY_CACHE_WARMING_TICK_SECONDS = 30 + + +async def _complexity_cache_warming_loop(refresher: CacheWarmingRefresher | None = None): + global llm_router, prisma_client + active_refresher = refresher if refresher is not None else CacheWarmingRefresher() + while True: + try: + await asyncio.sleep(_COMPLEXITY_CACHE_WARMING_TICK_SECONDS) + router = llm_router + if router is None: + continue + await active_refresher.run_tick( + llm_router=router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except asyncio.CancelledError: + raise + except Exception: # noqa: BLE001 # one failed tick must never kill the warming loop + verbose_proxy_logger.exception("complexity_router cache warming tick failed") + + async def _run_background_health_check(): """ Periodically run health checks in the background on the endpoints. diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index a6267453bf7..72bcb7f9dee 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -142,6 +142,48 @@ Technical code keywords are detected case-insensitively and include: - Infrastructure: `database`, `api`, `endpoint`, `docker`, `kubernetes` - Actions: `debug`, `implement`, `refactor`, `optimize` +## Cache Warming + +Provider prompt caches (Anthropic, Bedrock) are per-model, so a mid-session tier switch pays a fresh cache write on the new model and loses the cache-read discount. `cache_warming` keeps every tier model's prompt cache warm for active sessions: the proxy captures each session's latest payload and a background refresher replays it (`max_tokens=1`) against the other tier models before the provider's ~5 minute cache TTL expires. When the router later switches tiers, the switched-to model already has the session's prefix cached, and the routing pick prefers models whose cache is verifiably warm. + +```yaml +model_list: + - model_name: smart-router + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + tiers: + SIMPLE: fast-claude + COMPLEX: smart-claude + session_affinity: false # warming is the alternative to pinning; see below + cache_warming: + enabled: true + refresh_interval_seconds: 270 # keep under the provider cache TTL (Anthropic: 5 min) + session_ttl_seconds: 3600 + idle_timeout_seconds: 600 # stop warming a session this long after its last real request + max_sessions: 1000 + # warm_models: [fast-claude, smart-claude] # default: first member of each tier pool + +general_settings: + store_prompts_in_spend_logs: true # consent gate; warming stores full payloads in Redis + +router_settings: + redis_host: localhost + redis_port: 6379 +``` + +Requirements and semantics: + +- **Redis is required.** Session payloads, per-model warmth stamps, and a per-router session index live in Redis so all pods share them and a single pod (via a Redis cron lock) runs the replays. Without Redis, warming logs a warning once and no-ops; requests are unaffected. Sessions are tracked through the index rather than keyspace scans, so Redis Cluster is supported. +- **Prompt retention consent is a prerequisite.** Warming persists full request payloads (messages, system, tools) in Redis, so capture requires `store_prompts_in_spend_logs: true` and respects message redaction: with the flag off, or with `turn_off_message_logging` (globally or via the per-request redaction header) active, capture warns once and skips. +- **Only Anthropic and Bedrock models that support prompt caching are warmed.** Other models in the tier pools are left alone. Requests must carry a `metadata.session_id` and exceed the warm set's minimum cacheable token count (`prompt_cache_min_tokens`, default 1024) to be captured. +- **A replay is admitted like a request.** Each replay is assembled as a request body and put through the proxy's own admission entry points before it is dispatched: the budget reservation (the same call the auth layer makes after `common_checks`, so the key, team, user, end-user, organization and tag counters all apply and are reconciled to the replay's actual cost) and then `ProxyLogging.pre_call_hook` (so RPM, TPM, max-parallel and every configured guardrail apply, on the proxy's own shared counters rather than a private copy). A rejection skips that one replay and is retried on the next tick; the failure hook returns whatever the rejected replay had already reserved. Warming stops for keys that are deleted, blocked, or expired; key state is verified fresh each tick and when the database is unreachable the tick is skipped, so warming pauses until it is reachable again (caches re-warm on the next successful tick). +- **Warming cost is visible where the customer already looks for cost.** Warming writes no spend logs of its own, so the replay rows are the only record of warming cost that exists, and they carry the same identity block a real request on that key carries: key, team, user and end-user attribution, the key's and team's own `tags` and `spend_logs_metadata`, and the `litellm_cache_warming` tag alongside them so warming is both included in per-tag chargeback and filterable out of it. Replays also fan out to the key-scoped and team-scoped logging callbacks, so a team pointing its traffic at its own Langfuse sees warming there too rather than only in the proxy-wide logs. A `max_tokens=1` replay of a warm prefix bills roughly 10% of the input cost. Key-level cache controls, `disable_fallbacks` and the global-guardrail opt-outs are applied the same way, which means a key that declares its own `cache` controls overrides warming's response-cache bypass exactly as it overrides a caller's. +- **Guardrails run on replays.** Because a replay goes through `pre_call_hook`, the guardrails configured for the model group (globally and on the deployment) also run on its warming replays. A blocking guardrail costs one skipped warm; a content-rewriting guardrail makes the warm ineffective rather than wrong, because a rewritten prefix simply is not the prefix real traffic sends. Sessions on a key that declares `max_iterations` are skipped instead of warmed, because that limiter counts every request on a `session_id` and cannot be consulted without incrementing it, so warming would consume the caller's own iteration budget. +- **Multi-deployment groups require deployment affinity.** A replay routes by group name, so with several deployments it warms one member's cache while real traffic spreads across all of them: paying the cache-write premium against 1/N routing odds is worse than not warming, so such a group is skipped (with a one-time warning naming the group and both remedies) unless `DeploymentAffinityCheck` is active for it with the session_id mode, enabled globally via `router_settings.optional_pre_call_checks: ["session_affinity"]` or per group via `router_settings.model_group_affinity_config`. When active, replays carry the session's `session_id` (and the originating key hash), so the affinity check pins warming and real traffic to the same deployment. Single-deployment groups warm regardless. More than one deployment is used as the conservative stand-in for "more than one provider cache domain" +- **`max_sessions`** caps concurrently warmed sessions per auto-router, enforced atomically at capture; once reached, new sessions are not admitted until existing ones expire. +- **Interplay with `session_affinity`** (default on): affinity pins a session to its first-turn model, so no tier switch happens and warming buys nothing; with affinity on, captured sessions are still warmed but the pin decides routing. Disable `session_affinity` to let per-turn classification switch tiers and have warming make those switches cache hits. + ## Performance - **Classification time**: <1ms typical diff --git a/litellm/router_strategy/complexity_router/cache_warming/__init__.py b/litellm/router_strategy/complexity_router/cache_warming/__init__.py new file mode 100644 index 00000000000..0a2c1f4757a --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/__init__.py @@ -0,0 +1,33 @@ +from litellm.router_strategy.complexity_router.cache_warming.capture import capture_session +from litellm.router_strategy.complexity_router.cache_warming.eligibility import ( + min_prompt_cache_tokens_for_warm_set, + resolve_warm_models, +) +from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_RECORD_SCHEMA_VERSION, + CACHE_WARMING_REPLAY_MARKER_KEY, + CACHE_WARMING_REPLAY_TAG, + WARM_FRESHNESS_SLACK_SECONDS, + CacheWarmingAttribution, + CacheWarmingPayload, + CacheWarmingRecord, + compress_payload, + decompress_payload, +) + +__all__ = [ + "CacheWarmingStore", + "capture_session", + "min_prompt_cache_tokens_for_warm_set", + "resolve_warm_models", + "CACHE_WARMING_RECORD_SCHEMA_VERSION", + "CACHE_WARMING_REPLAY_MARKER_KEY", + "CACHE_WARMING_REPLAY_TAG", + "WARM_FRESHNESS_SLACK_SECONDS", + "CacheWarmingAttribution", + "CacheWarmingPayload", + "CacheWarmingRecord", + "compress_payload", + "decompress_payload", +] diff --git a/litellm/router_strategy/complexity_router/cache_warming/capture.py b/litellm/router_strategy/complexity_router/cache_warming/capture.py new file mode 100644 index 00000000000..877cd109b94 --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/capture.py @@ -0,0 +1,224 @@ +import asyncio +import json +from collections.abc import Mapping, Sequence +from functools import lru_cache +from typing import TYPE_CHECKING, Literal + +from litellm._logging import verbose_router_logger +from litellm.constants import DEFAULT_CHARS_PER_TOKEN +from litellm.router_strategy.complexity_router.cache_warming.eligibility import ( + min_prompt_cache_tokens_for_warm_set, + resolve_warm_models, +) +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_REPLAY_MARKER_KEY, + CACHE_WARMING_REPLAY_TAG, + CacheWarmingAttribution, + CacheWarmingPayload, + compress_payload, +) +from litellm.litellm_core_utils.core_helpers import ( + get_request_metadata_field, + iter_request_metadata_dicts, +) + +if TYPE_CHECKING: + from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter + +_MAX_UNCOMPRESSED_RATIO = 8 +_ATTRIBUTION_KEYS = ( + "user_api_key", + "user_api_key_hash", + "user_api_key_user_id", + "user_api_key_team_id", + "user_api_key_org_id", + "user_api_key_project_id", + "user_api_key_end_user_id", +) + + +@lru_cache(maxsize=64) +def _warn_privacy_gate_blocked(auto_router_model_name: str) -> None: + verbose_router_logger.warning( + "cache_warming is enabled for auto-router %s but prompt retention is not permitted " + "(store_prompts_in_spend_logs is off or message redaction is active); capture is skipped", + auto_router_model_name, + ) + + +@lru_cache(maxsize=4096) +def _warn_payload_too_large(auto_router_model_name: str, session_id: str) -> None: + verbose_router_logger.warning( + "cache_warming: session %s on auto-router %s exceeds max_payload_bytes; not warming this session", + session_id, + auto_router_model_name, + ) + + +def _capture_allowed(kwargs: Mapping[str, object]) -> bool: + """One consent predicate for one question: may this request's prompt content be + retained? Honors both halves of the operator's stated policy: the redaction + opt-out (turn_off_message_logging, including per-request and header forms) and + the prompt-retention opt-in (store_prompts_in_spend_logs; SDK use without the + proxy consents through cache_warming.enabled itself).""" + from litellm.litellm_core_utils.redact_messages import ( + should_redact_message_logging, # pyright: ignore[reportUnknownVariableType] # legacy-untyped helper + ) + + if should_redact_message_logging({"litellm_params": kwargs}): # mutable-ok: read-only view for the predicate + return False + try: + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + _should_store_prompts_and_responses_in_spend_logs, # pyright: ignore[reportPrivateUsage] # canonical proxy consent gate; no public accessor exists + ) + except ImportError: + return True + return _should_store_prompts_and_responses_in_spend_logs() + + +def _is_replay(metadata_dicts: Sequence[Mapping[str, object]]) -> bool: + for metadata in metadata_dicts: + if metadata.get(CACHE_WARMING_REPLAY_MARKER_KEY): + return True + tags = metadata.get("tags") + if isinstance(tags, list) and CACHE_WARMING_REPLAY_TAG in tags: + return True + return False + + +def _call_surface(request_kwargs: Mapping[str, object]) -> Literal["chat_completions", "anthropic_messages"]: + """Two in-band signals, checked in order. Through the proxy the logging object + rides the request into deployment selection and its call_type carries the route + (function_setup stamps the entry function's name), which is the reliable signal + because the fallback machinery strips original_function keys before the + pre-routing hook runs. The generic-dispatch stamp is kept as the SDK-direct + fallback for call shapes that never enter the proxy layer.""" + logging_call_type = getattr(request_kwargs.get("litellm_logging_obj"), "call_type", None) + if logging_call_type in ("anthropic_messages", "aanthropic_messages"): + return "anthropic_messages" + generic_function = request_kwargs.get("original_generic_function") + if getattr(generic_function, "__name__", None) == "anthropic_messages": + return "anthropic_messages" + return "chat_completions" + + +def _prompt_chars(payload: CacheWarmingPayload) -> int: + """Every character the provider will cache: prompt text via the shared flattener (which reads both + string and content-block shapes), plus tool schemas and tool_call arguments. Tool schemas dominate + agent sessions, so omitting them made those sessions fail the min-token gate and never warm. + The limiter's own _estimate_tokens_for_request omits tools too and is proxy-side, so no owner.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import get_str_from_messages + + system = payload.system if isinstance(payload.system, str) else list(payload.system or ()) + system_message = ({"role": "system", "content": system},) if payload.system is not None else () + text_chars = len(get_str_from_messages([*system_message, *payload.messages])) # pyright: ignore[reportArgumentType] # captured wire shape + tool_chars = sum(len(json.dumps(tool, default=str)) for tool in payload.tools or ()) + tool_call_chars = sum( + len(json.dumps(call, default=str)) + for message in payload.messages + for call in (message.get("tool_calls") or ()) + if isinstance(call, Mapping) + ) + return text_chars + tool_chars + tool_call_chars + + +def _gate_and_compress( + payload: CacheWarmingPayload, max_payload_bytes: int, min_tokens: int +) -> "tuple[str, str, int] | Literal['too_large', 'too_small']": + """Size gates plus compression, run off the event loop. The uncompressed bound + runs before any compression so adversarial highly-compressible content cannot + buy unbounded CPU with a small compressed result, and the chars/4 token + estimate deliberately avoids running a real tokenizer over a full multi-turn + conversation in the request path; counting prompt TEXT keeps the JSON envelope out of it.""" + serialized_chars = len(payload.model_dump_json()) + if serialized_chars > _MAX_UNCOMPRESSED_RATIO * max_payload_bytes: + return "too_large" + token_estimate = max(1, _prompt_chars(payload) // DEFAULT_CHARS_PER_TOKEN) + if token_estimate < min_tokens: + return "too_small" + blob, sha = compress_payload(payload) + if len(blob) > max_payload_bytes: + return "too_large" + return (blob, sha, token_estimate) + + +def _extract_attribution(metadata_dicts: Sequence[Mapping[str, object]]) -> CacheWarmingAttribution: + """Identity comes from the ONE proxy-stamped slot, never merged: the proxy strips only the + ``user_api_key_`` prefix, so a bare ``user_api_key`` in the other slot would forge billing identity. + ``user_api_key_hash`` cannot be client-injected, so it marks the authoritative slot.""" + stamped = next((metadata for metadata in metadata_dicts if "user_api_key_hash" in metadata), None) + if stamped is None: + return CacheWarmingAttribution() + return CacheWarmingAttribution( + **{ # mutable-ok: handed straight to pydantic, never retained + key: str(value) for key, value in stamped.items() if key in _ATTRIBUTION_KEYS and value is not None + } + ) + + +def _build_payload( + request_kwargs: Mapping[str, object], + messages: "Sequence[Mapping[str, object]] | None", + call_surface: Literal["chat_completions", "anthropic_messages"], + routed_model: str, +) -> CacheWarmingPayload | None: + if not messages: + return None + return CacheWarmingPayload.model_validate( + { # mutable-ok: pydantic input, never retained + "model": routed_model, + "messages": messages, + "system": request_kwargs.get("system") if call_surface == "anthropic_messages" else None, + "tools": request_kwargs.get("tools"), + "tool_choice": request_kwargs.get("tool_choice"), + "call_surface": call_surface, + } + ) + + +async def capture_session( + strategy: "ComplexityRouter", + request_kwargs: Mapping[str, object], + messages: "Sequence[Mapping[str, object]] | None", + routed_model: str, +) -> None: + metadata_dicts = iter_request_metadata_dicts(request_kwargs) + if _is_replay(metadata_dicts): + return + if not _capture_allowed(request_kwargs): + _warn_privacy_gate_blocked(strategy.model_name) + return + session_id = get_request_metadata_field(request_kwargs, "session_id") + if session_id is None: + return + payload = _build_payload(request_kwargs, messages, _call_surface(request_kwargs), routed_model) + if payload is None: + return + config = strategy.config.cache_warming + warm_models = resolve_warm_models(strategy.config) + gated = await asyncio.to_thread( + _gate_and_compress, payload, config.max_payload_bytes, min_prompt_cache_tokens_for_warm_set(warm_models) + ) + match gated: + case "too_large": + _warn_payload_too_large(strategy.model_name, session_id) + return + case "too_small": + return + case (blob, sha, token_estimate): + pass + store = strategy.get_cache_warming_store() + if store is None: + return + caller_scope = get_request_metadata_field(request_kwargs, "user_api_key_hash") or "unscoped" + await store.upsert_session( + caller_scope=caller_scope, + session_id=session_id, + payload_compressed=blob, + payload_sha256=sha, + token_estimate=token_estimate, + served_model=routed_model, + attribution=_extract_attribution(metadata_dicts), + ttl_seconds=config.session_ttl_seconds, + max_sessions=config.max_sessions, + ) diff --git a/litellm/router_strategy/complexity_router/cache_warming/eligibility.py b/litellm/router_strategy/complexity_router/cache_warming/eligibility.py new file mode 100644 index 00000000000..22d0cb300a5 --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/eligibility.py @@ -0,0 +1,22 @@ +from typing import TYPE_CHECKING + +from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT + +if TYPE_CHECKING: + from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig + + +def resolve_warm_models(config: "ComplexityRouterConfig") -> tuple[str, ...]: + explicit = config.cache_warming.warm_models + if explicit: + return tuple(dict.fromkeys(explicit)) + first_per_tier = (models if isinstance(models, str) else models[0] for models in config.tiers.values() if models) + return tuple(dict.fromkeys(first_per_tier)) + + +def min_prompt_cache_tokens_for_warm_set(warm_models: tuple[str, ...]) -> int: + from litellm.utils import get_prompt_cache_min_tokens + + if not warm_models: + return DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT + return min(get_prompt_cache_min_tokens(model) for model in warm_models) diff --git a/litellm/router_strategy/complexity_router/cache_warming/refresher.py b/litellm/router_strategy/complexity_router/cache_warming/refresher.py new file mode 100644 index 00000000000..d509ead716f --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/refresher.py @@ -0,0 +1,809 @@ +import asyncio +import contextlib +import time +import uuid +from collections.abc import Awaitable, Callable, Mapping, Sequence +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +import litellm +from litellm._logging import verbose_router_logger +from litellm.constants import CACHE_WARMING_JOB_NAME, LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.integrations.anthropic_cache_control_hook import EXPLICIT_PROMPT_CACHING_PROVIDERS +from litellm.router_strategy.complexity_router.cache_warming.eligibility import resolve_warm_models +from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_REPLAY_MARKER_KEY, + CACHE_WARMING_REPLAY_TAG, + CacheWarmingPayload, + CacheWarmingRecord, + decompress_payload, + needs_rewarming, + warn_once, +) + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import ( + RedisDistributedLock, + ) + from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import PrismaClient, ProxyLogging + from litellm.router import Router + from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter + from litellm.types.utils import CallTypesLiteral + +CACHE_WARMING_MAX_CONCURRENT_REPLAYS = 10 +CACHE_WARMING_LOCK_TTL_SECONDS = 60 +_RECONSTRUCTED_IDENTITY_FIELDS = ("user_id", "team_id", "org_id", "project_id") + +_REPLAY_MAX_OUTPUT_TOKENS = 1 + + +def _replay_surface(payload: CacheWarmingPayload) -> "tuple[str, CallTypesLiteral]": + """Route and call type as the owning endpoints declare them, so hooks branching on either see the + surface being replayed.""" + if payload.call_surface == "anthropic_messages": + return ("/v1/messages", "anthropic_messages") + return ("/v1/chat/completions", "acompletion") + + +def _replay_principal(key_state: "UserAPIKeyAuth | None", record: CacheWarmingRecord, route: str) -> "UserAPIKeyAuth": + """A replay's principal must be COMPLETE, because the authorization enumerations resolve tenancy by id and + a missing field is an absent ceiling rather than a denial. Three shapes of capture, one reconstruction: + + (a) virtual key: key_state was just read authoritatively from the database, so it is already complete + (b) keyless proxy caller (JWT, and anything else the proxy authenticated without a virtual key, where + api_key is None): UserAPIKeyAuth is litellm's principal type for those callers too and the proxy + stamped their tenancy, which the record carries, so every recorded identity field is restored below and + the team, user, organization and project gates resolve by id exactly as for a virtual key. Dropping any + of them is what previously let a JWT caller's replays run against an empty principal and escape every + tenant control + (c) direct SDK use with no proxy auth object: nothing was recorded, so there is no tenancy to preserve and + the principal is genuinely limitless, which is the same unattributed warming as before + + get_key_object fills token but not api_key, and every metadata consumer reads api_key, so it is set from + the token. end_user_id is request-scoped and survives only on the record.""" + from litellm.proxy._types import UserAPIKeyAuth + + attribution = record.attribution + base = key_state if key_state is not None else UserAPIKeyAuth(api_key=attribution.user_api_key) + return base.model_copy( + update={ # mutable-ok: pydantic model_copy input, never retained + "api_key": base.api_key or base.token, + **{ + field: getattr(base, field) or getattr(attribution, f"user_api_key_{field}") + for field in _RECONSTRUCTED_IDENTITY_FIELDS + }, + "end_user_id": attribution.user_api_key_end_user_id or base.end_user_id, + "request_route": route, + "budget_reservation": None, + } + ) + + +def _replay_body( + payload: CacheWarmingPayload, record: CacheWarmingRecord, model_group: str, route: str +) -> "dict[str, object]": + """no-cache is only sent when a response cache exists to bypass, because it is a cache control the key + may be forbidden to set (cache_control_check.py:36) and would otherwise refuse the replay for nothing. + Anthropic invalidates a cached prefix when tool_choice changes, so it and system ride along there. + litellm_call_id is per replay because the hanging-request checker keys its cache on it, and the empty + default would collide every replay onto one entry. The warming marker rides spend_logs_metadata, not + metadata.tags: tags are an input to deployment selection (enable_tag_filtering makes an unmatched tag + unroutable) and to policy (_reject_clientside_metadata_tags_check refuses any request carrying them), + while spend_logs_metadata exists to label spend rows, which is all this marker is for.""" + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + data: dict[str, object] = { # mutable-ok: the request body the proxy entry points mutate in place + "model": model_group, + "messages": [dict(message) for message in payload.messages], + "tools": list(payload.tools) if payload.tools is not None else None, + "tool_choice": dict(payload.tool_choice) if isinstance(payload.tool_choice, Mapping) else payload.tool_choice, + "max_tokens": _REPLAY_MAX_OUTPUT_TOKENS, + "stream": False, + "litellm_call_id": uuid.uuid4().hex, + **({"cache": {"no-cache": True}} if litellm.cache is not None else {}), + **( + {"system": list(payload.system) if isinstance(payload.system, tuple) else payload.system} + if payload.call_surface == "anthropic_messages" + else {} + ), + } + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=data, route=route) + data[get_metadata_variable_name_from_kwargs(data)] = { # mutable-ok: request metadata, never retained + CACHE_WARMING_REPLAY_MARKER_KEY: True, + **({"session_id": record.session_id} if record.session_id is not None else {}), + "spend_logs_metadata": {CACHE_WARMING_REPLAY_TAG: "true"}, # mutable-ok: request metadata, never retained + } + return data + + +async def _authorize_replay( + *, + principal: "UserAPIKeyAuth", + data: "dict[str, object]", + model_group: str, + route: str, + team: "LiteLLM_TeamTable | None", + user: "LiteLLM_UserTable | None", + end_user: "LiteLLM_EndUserTable | None", + project: "LiteLLM_ProjectTableCachedObj | None", + llm_router: "Router", + prisma_client: "PrismaClient", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging", + skip_budget_checks: bool, +) -> None: + """Authorization is not re-enumerated here. Production has exactly two enumerations of it, and a replay + runs both: ``_enforce_key_and_fallback_model_access`` for the key-level model allowlist and its fallback + targets, then ``common_checks`` for everything else (team blocked, team and member and user and project + model access, every budget scope including organization and tags, the global proxy budget, guardrail + modification, organization RBAC, vector stores, tool allowlists), plus the per-model budget gates the + auth builder owns for the key and the end user. The per-model gate is applied without the builder's + budget_fallbacks rewrite, because rerouting a replay to a different model would warm a cache nothing is + going to read; over budget means this group is simply not warmed. + + Both accept ``request=None`` by declared contract, because every Request dereference inside them is + already guarded (``_safe_get_request_headers``, ``_safe_get_request_query_params``, + ``get_request_route_template``, ``_enforce_user_param_check``). The one exception is the agent trace-id + header gate, which is reached only for a key whose agent sets ``require_trace_id_on_calls_by_agent`` and + which raises there; that denies the replay, which is the outcome a replay should get from a gate demanding + a client header it cannot have. + + Calling the enumerations rather than copying them is the point: three separate review rounds found a gate + a hand-written list had missed, and a hand-written list will always be one gate behind.""" + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.auth.user_api_key_auth import ( + _enforce_key_and_fallback_model_access, # pyright: ignore[reportPrivateUsage] # the key-level enumeration; its own docstring notes common_checks excludes these + get_global_proxy_spend, + ) + from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, model_max_budget_limiter + + await _enforce_key_and_fallback_model_access( + valid_token=principal, + request_data=data, + route=route, + request=None, + llm_model_list=llm_router.model_list, + llm_router=llm_router, + ) + await model_max_budget_limiter.is_key_within_model_budget(user_api_key_dict=principal, model=model_group) + end_user_model_budgets = principal.end_user_model_max_budget + if principal.end_user_id is not None and isinstance(end_user_model_budgets, dict) and end_user_model_budgets: + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=principal.end_user_id, + end_user_model_max_budget=end_user_model_budgets, + model=model_group, + ) + await common_checks( + request_body=data, + team_object=team, + user_object=user, + end_user_object=end_user, + global_proxy_spend=await get_global_proxy_spend( + litellm_proxy_admin_name=litellm_proxy_admin_name, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + token=principal.token or "", + proxy_logging_obj=proxy_logging_obj, + ), + general_settings=general_settings, + route=route, + llm_router=llm_router, + proxy_logging_obj=proxy_logging_obj, + valid_token=principal, + request=None, + skip_budget_checks=skip_budget_checks, + project_object=project, + ) + + +def _stamp_identity(data: "dict[str, object]", principal: "UserAPIKeyAuth") -> None: + """The request path's own three stampers in its order. Load-bearing: add_key_level_controls resets + data["cache"] and refills it from key metadata, so running after _replay_body lets the key's declared + controls win over warming's bypass exactly as they win over a caller's, and the reservation must + already be on the principal because metadata is the channel the cost callback reconciles it through + (proxy_track_cost_callback.py:446).""" + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + from litellm.proxy.proxy_server import proxy_config + + metadata_key = get_metadata_variable_name_from_kwargs(data) + LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data=data, user_api_key_dict=principal, _metadata_variable_name=metadata_key + ) + LiteLLMProxyRequestSetup.add_key_team_project_metadata( + data=data, user_api_key_dict=principal, _metadata_variable_name=metadata_key + ) + LiteLLMProxyRequestSetup.apply_dynamic_logging_settings( + data=data, user_api_key_dict=principal, proxy_config=proxy_config + ) + + +def _excluded_from_warming(key_state: "UserAPIKeyAuth", now: "datetime", proxy_logging_obj: "ProxyLogging") -> bool: + """Mirrors the canonical auth checks (user_api_key_auth.py:1717 and :2891) because common_checks owns + that policy for real traffic but needs a FastAPI Request. datetime.fromisoformat only accepts a + Z-suffixed expires from 3.11 while requires-python is 3.10, so the offset is normalized and an + unparseable value excludes rather than admits. A declared max-iterations ceiling excludes too: that + hook counts every request on a session_id with no read-only mode, so warming would spend the caller's + own loop budget.""" + from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler + + if key_state.blocked is True: + return True + iterations_hook = proxy_logging_obj.get_proxy_hook("max_iterations_limiter") + if ( + isinstance(iterations_hook, _PROXY_MaxIterationsHandler) + and iterations_hook._get_max_iterations(key_state) is not None # pyright: ignore[reportPrivateUsage] # the hook owns this predicate and exposes no public form + ): + warn_once( + "cache_warming: a key declares max_iterations, whose limiter counts every request on a " + "session_id and cannot be consulted without incrementing it; warming would spend the " + "caller's own iteration budget, so sessions on such keys are skipped" + ) + return True + expires = key_state.expires + if expires is None: + return False + try: + expiry = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires.replace("Z", "+00:00")) + except ValueError: + return True + if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: + expiry = expiry.replace(tzinfo=timezone.utc) + return expiry < now + + +def _missing_enforcement_dependency( + prisma_client: "PrismaClient | None", + user_api_key_cache: "DualCache | None", + proxy_logging_obj: "ProxyLogging", +) -> str | None: + """Names the first thing a replay's admission needs and cannot reach, or None. Every ceiling warming + claims to respect is enforced by one of these, so losing any single one turns admission into a no-op that + still spends the customer's money. They are therefore checked once, together, before any session is + considered, rather than as per-dependency arms that each fail open on their own path.""" + from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 + + if prisma_client is None: + return "the database client (key, team, user and project state)" + if user_api_key_cache is None: + return "the key cache (authorization objects and budget counters)" + if not isinstance( + proxy_logging_obj.get_proxy_hook("parallel_request_limiter"), _PROXY_MaxParallelRequestsHandler_v3 + ): + return "the v3 parallel-request limiter (rate limits)" + from litellm.proxy.proxy_server import spend_counter_cache + + if spend_counter_cache is None: + return "the shared spend counter plane (budget reservation and reconciliation)" + return None + + +def _deny_tick(missing: str) -> None: + warn_once( + f"cache_warming cannot enforce a replay's ceilings because {missing} is unavailable; no session is " + "warmed until it is reachable again, because admission without it would spend against limits and " + "budgets it cannot check" + ) + + +def _resolve_proxy_logging() -> "ProxyLogging": + from litellm.proxy.proxy_server import proxy_logging_obj + + return proxy_logging_obj + + +def _redis_lease_lock(redis_cache: object) -> "RedisDistributedLock": + """The MCP outbound-credentials distributed lock; promoting it to a neutral home is a follow-up.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import ( + RedisDistributedLock, + ) + + return RedisDistributedLock( + redis_cache.init_async_client(), # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownArgumentType] # RedisCache is legacy-untyped + namespace_key=redis_cache.check_and_fix_namespace, # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownArgumentType] # RedisCache is legacy-untyped + ) + + +async def _proxy_key_state( + hashed_token: str, prisma_client: "PrismaClient", user_api_key_cache: "DualCache" +) -> "UserAPIKeyAuth": + """Read past the local cache tier, not cache-first. Revocation clears the local entry and the Redis entry + (_delete_cache_key_object, auth_checks.py:1855) but cannot reach another pod's in-memory tier, so a + cache-first read can still see a key blocked, expired or deleted seconds ago as valid and keep billing it. + DualCache.async_get_cache offers no local-skip (its local_only flag skips Redis, the opposite), so the + authoritative row is the way past it; get_key_object writes the result back through _cache_key_object, so + the platform's cache shape is unchanged. Cost is one query per DISTINCT attributed key per 30-second tick, + bounded by max_sessions (default 1000). Warming is therefore strictly fresher than real request traffic on + another pod, which carries the same in-memory staleness.""" + from litellm.proxy.auth.auth_checks import get_key_object + + return await get_key_object( + hashed_token=hashed_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, # pyright: ignore[reportArgumentType] # UserApiKeyCache is a DualCache alias + parent_otel_span=None, + proxy_logging_obj=None, + check_cache_only=None, + check_db_only=True, + ) + + +def collect_warming_enabled_complexity_routers(llm_router: "Router") -> tuple["ComplexityRouter", ...]: + return tuple( + tagged.strategy + for tagged_list in llm_router.complexity_routers.values() + for tagged in tagged_list + if tagged.strategy.config.cache_warming.enabled + ) + + +def _deployment_provider(litellm_params: Mapping[str, object], deployment_model: str) -> str | None: + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + declared = litellm_params.get("custom_llm_provider") + if isinstance(declared, str) and declared: + return declared + try: + _, provider, _, _ = get_llm_provider(model=deployment_model) + except Exception: # noqa: BLE001 # unroutable deployment just isn't warmable + return None + return provider + + +def _session_affinity_active(llm_router: "Router", model_group: str) -> bool: + from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck + + return any( + callback._get_effective_flags(model_group)[2] # pyright: ignore[reportPrivateUsage] # no public accessor today; upstream should expose get_effective_flags + for callback in (llm_router.optional_callbacks or []) # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # Router.optional_callbacks is legacy-untyped + if isinstance(callback, DeploymentAffinityCheck) + ) + + +def _group_is_cache_warmable(llm_router: "Router", model_group: str) -> bool: + """Every selectable member must be warmable, because replays route by group name and the Router may + pick any member.""" + from litellm.utils import supports_prompt_caching + + deployments = llm_router.get_model_list(model_name=model_group) or [] + if not deployments: + return False + if len(deployments) > 1 and not _session_affinity_active(llm_router, model_group): + warn_once( + f"cache_warming: model group {model_group} has {len(deployments)} deployments and deployment " + "affinity with session_id is not active for it, so warming is skipped for this group; warming " + "without affinity cannot produce cache hits reliably because a replay warms one member's cache " + "while real traffic routes across all of them. Enable it globally with " + 'router_settings.optional_pre_call_checks: ["session_affinity"] or for this group only with ' + f'router_settings.model_group_affinity_config: {{"{model_group}": ["session_affinity"]}}' + ) + return False + return all( + isinstance(model, str) + and (provider := _deployment_provider(params, model)) in EXPLICIT_PROMPT_CACHING_PROVIDERS + and supports_prompt_caching(model=model, custom_llm_provider=provider) + for deployment in deployments + for params in ((deployment.get("litellm_params") or {}),) # pyright: ignore[reportUnknownMemberType] # DeploymentTypedDict fields are legacy-untyped + for model in (params.get("model"),) # pyright: ignore[reportUnknownMemberType] # DeploymentTypedDict fields are legacy-untyped + ) + + +def filter_cache_warmable(llm_router: "Router", model_groups: Sequence[str]) -> tuple[str, ...]: + return tuple(group for group in model_groups if _group_is_cache_warmable(llm_router, group)) + + +class CacheWarmingRefresher: + def __init__( + self, + max_concurrent_replays: int = CACHE_WARMING_MAX_CONCURRENT_REPLAYS, + lock_ttl_seconds: float = CACHE_WARMING_LOCK_TTL_SECONDS, + lock_factory: 'Callable[[object], "RedisDistributedLock"]' = _redis_lease_lock, + key_state_resolver: 'Callable[[str, "PrismaClient", "DualCache"], Awaitable["UserAPIKeyAuth"]]' = _proxy_key_state, + proxy_logging_resolver: 'Callable[[], "ProxyLogging"]' = _resolve_proxy_logging, + ) -> None: + self.max_concurrent_replays = max_concurrent_replays + self.lock_ttl_seconds = lock_ttl_seconds + self.lock_factory = lock_factory + self.key_state_resolver = key_state_resolver + self.proxy_logging_resolver = proxy_logging_resolver + + async def _hold_lease(self, lock: "RedisDistributedLock", token: str, lease_lost: asyncio.Event) -> None: + while True: + await asyncio.sleep(self.lock_ttl_seconds / 2) + if not await lock.extend(CACHE_WARMING_JOB_NAME, token, self.lock_ttl_seconds): + verbose_router_logger.warning( + "cache_warming pod lock was lost mid tick; finishing in-flight replays without admitting new ones" + ) + lease_lost.set() + return + + async def run_tick( + self, + *, + llm_router: "Router", + prisma_client: "PrismaClient | None", + user_api_key_cache: "DualCache | None" = None, + ) -> None: + warming_routers = collect_warming_enabled_complexity_routers(llm_router) + if not warming_routers: + return + warmable = tuple( + (complexity_router, store) + for complexity_router in warming_routers + if (store := complexity_router.get_cache_warming_store()) is not None and store.redis_cache is not None + ) + if not warmable: + return + from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import ( + LockAcquisition, + ) + + redis_cache = warmable[0][1].redis_cache + if redis_cache is None: + return + lock = self.lock_factory(redis_cache) + token = uuid.uuid4().hex + acquisition = await lock.acquire(CACHE_WARMING_JOB_NAME, token, self.lock_ttl_seconds) + if acquisition is not LockAcquisition.ACQUIRED: + if acquisition is LockAcquisition.ERROR: + _deny_tick( + "the pod lock backend (leader election, without which every pod warms concurrently and " + "warmth stamps cannot dedupe because each pod reads them before any pod writes)" + ) + return + lease_lost = asyncio.Event() + lease = asyncio.create_task(self._hold_lease(lock, token, lease_lost)) + try: + for complexity_router, store in warmable: + await self._warm_router_sessions( + llm_router=llm_router, + complexity_router=complexity_router, + store=store, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + lease_lost=lease_lost, + ) + finally: + lease.cancel() + with contextlib.suppress(asyncio.CancelledError): + await lease + await lock.release(CACHE_WARMING_JOB_NAME, token) + + async def _warm_router_sessions( + self, + *, + llm_router: "Router", + complexity_router: "ComplexityRouter", + store: CacheWarmingStore, + prisma_client: "PrismaClient | None", + user_api_key_cache: "DualCache | None", + lease_lost: asyncio.Event, + ) -> None: + config = complexity_router.config.cache_warming + session_keys = await store.list_session_keys(max_sessions=config.max_sessions) + if not session_keys: + return + if len(session_keys) >= config.max_sessions: + verbose_router_logger.debug( + "cache_warming: auto-router %s is at its max_sessions cap (%s); " + "new sessions are not admitted until existing ones expire", + complexity_router.model_name, + config.max_sessions, + ) + now = time.time() + records = tuple([(key, await store.get_record(key)) for key in session_keys]) + active = tuple( + (key, record) + for key, record in records + if record is not None and now - record.last_activity <= config.idle_timeout_seconds + ) + if not active: + return + warm_models = filter_cache_warmable(llm_router, resolve_warm_models(complexity_router.config)) + if not warm_models: + return + attributed = frozenset( + record.attribution.user_api_key + for _, record in active + if record.attribution.user_api_key is not None + and record.attribution.user_api_key != LITELLM_PROXY_MASTER_KEY_ALIAS + ) + proxy_logging_obj = self.proxy_logging_resolver() + missing = _missing_enforcement_dependency(prisma_client, user_api_key_cache, proxy_logging_obj) + if missing is not None: + _deny_tick(missing) + return + key_states = await self._fetch_key_states(prisma_client, user_api_key_cache, attributed) + if key_states is None: + return + checked_at = datetime.now(timezone.utc) + excluded_keys = frozenset( + key + for key in attributed + if (row := key_states.get(key)) is None or _excluded_from_warming(row, checked_at, proxy_logging_obj) + ) + semaphore = asyncio.Semaphore(self.max_concurrent_replays) + outcomes = await asyncio.gather( + *( + self._warm_session( + llm_router=llm_router, + store=store, + session_key=key, + record=record, + warm_models=warm_models, + refresh_interval_seconds=config.refresh_interval_seconds, + session_ttl_seconds=config.session_ttl_seconds, + semaphore=semaphore, + key_state=key_states.get(record.attribution.user_api_key or ""), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + lease_lost=lease_lost, + ) + for key, record in active + if record.attribution.user_api_key not in excluded_keys + ), + return_exceptions=True, + ) + for outcome in outcomes: + if isinstance(outcome, BaseException): + verbose_router_logger.warning("cache_warming session task failed", exc_info=outcome) + + async def _warm_session( + self, + *, + llm_router: "Router", + store: CacheWarmingStore, + session_key: str, + record: CacheWarmingRecord, + warm_models: tuple[str, ...], + refresh_interval_seconds: int, + session_ttl_seconds: int, + semaphore: asyncio.Semaphore, + key_state: "UserAPIKeyAuth | None", + prisma_client: "PrismaClient | None", + user_api_key_cache: "DualCache | None", + proxy_logging_obj: "ProxyLogging", + lease_lost: asyncio.Event, + ) -> None: + warmth = await store.get_warmth(session_key, warm_models) + now = time.time() + due_models = tuple( + model for model in warm_models if needs_rewarming(warmth.get(model, 0.0), now, refresh_interval_seconds) + ) + if not due_models: + return + payload = decompress_payload(record.payload_compressed) + for model_group in due_models: + async with semaphore: + if lease_lost.is_set(): + return + admitted = await self._admit_replay( + llm_router=llm_router, + payload=payload, + record=record, + model_group=model_group, + key_state=key_state, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if admitted is None: + continue + data, principal = admitted + attempted_at = time.time() + try: + await ( + llm_router.aanthropic_messages(**data) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # factory-generated router surface is legacy-untyped + if payload.call_surface == "anthropic_messages" + else llm_router.acompletion(**data) # pyright: ignore[reportUnknownMemberType, reportCallIssue, reportUnknownVariableType, reportArgumentType] # router overloads are legacy-untyped + ) + await proxy_logging_obj.update_request_status( + litellm_call_id=str(data.get("litellm_call_id") or ""), status="success" + ) + except Exception as exc: # noqa: BLE001 # one failing replay must not abort the tick + verbose_router_logger.warning( + "cache_warming replay failed for session %s model %s", session_key, model_group, exc_info=True + ) + await self._report_rejection( + data=data, principal=principal, exc=exc, proxy_logging_obj=proxy_logging_obj + ) + finally: + await store.mark_warm_attempt(session_key, model_group, attempted_at, session_ttl_seconds) + + async def _admit_replay( + self, + *, + llm_router: "Router", + payload: CacheWarmingPayload, + record: CacheWarmingRecord, + model_group: str, + key_state: "UserAPIKeyAuth | None", + prisma_client: "PrismaClient | None", + user_api_key_cache: "DualCache | None", + proxy_logging_obj: "ProxyLogging", + ) -> "tuple[dict[str, object], UserAPIKeyAuth] | None": + """A replay is a request, so it enters the same entry points real traffic does and inherits both + halves of every contract they own rather than reimplementing any of them. Order mirrors the request + path: budget reservation, identity stamping, the model group's own guardrails + (common_request_processing.py:1268), pre_call_hook, then the blocked-model gate route_request + applies before the router call (route_llm_request.py:453). A rejection skips only this replay.""" + from litellm.proxy.route_llm_request import ( + _raise_if_model_fully_blocked, # pyright: ignore[reportPrivateUsage] # the gate route_request itself calls + ) + from litellm.proxy.utils import _check_and_merge_model_level_guardrails + + route, call_type = _replay_surface(payload) + principal = _replay_principal(key_state, record, route) + data = _replay_body(payload, record, model_group, route) + try: + await self._authorize_and_reserve( + principal=principal, + data=data, + model_group=model_group, + route=route, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + _stamp_identity(data, principal) + admitted = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=principal, + data=_check_and_merge_model_level_guardrails( + data=data, llm_router=llm_router, trust_client_model_info=False + ), + call_type=call_type, + ) + _raise_if_model_fully_blocked(llm_router, model_group, principal.team_id) + return (admitted, principal) + except Exception as exc: # noqa: BLE001 # a declined replay skips only itself + verbose_router_logger.debug( + "cache_warming: the request path declined this replay on %s: %s", model_group, exc + ) + await self._report_rejection(data=data, principal=principal, exc=exc, proxy_logging_obj=proxy_logging_obj) + return None + + @staticmethod + @staticmethod + async def _authorize_and_reserve( + *, + principal: "UserAPIKeyAuth", + data: "dict[str, object]", + model_group: str, + route: str, + llm_router: "Router", + prisma_client: "PrismaClient | None", + user_api_key_cache: "DualCache", + proxy_logging_obj: "ProxyLogging", + ) -> None: + """Authorization then reservation, the order and the owners auth itself uses + (user_api_key_auth.py:2380-2396). The context objects come from the same resolvers because both + halves are derived from them: the team, user, end-user, organization and project counters, and the + team, project and per-member authorization gates. Passing None would leave those ceilings unbound and + those gates unenforced while the cost callback still charged the same scopes. These stay cache-first, so + they carry the same in-memory staleness a real request on another pod carries; only the key, which is + the credential warming spends against, is read past the local tier.""" + from litellm.proxy.auth.auth_checks import ( + get_end_user_object, + get_project_object, + get_team_object, + get_user_object, + ) + from litellm.proxy.auth.user_api_key_auth import ( + _reserve_budget_after_common_checks, # pyright: ignore[reportPrivateUsage] # the flow owner; reserve_budget_for_request alone drops the operator settings + _should_skip_budget_checks, # pyright: ignore[reportPrivateUsage] # same owner, and Request-optional + ) + from litellm.proxy.proxy_server import general_settings + + cache: UserApiKeyCache = user_api_key_cache # pyright: ignore[reportAssignmentType] # UserApiKeyCache is a DualCache alias + team = ( + await get_team_object(team_id=principal.team_id, prisma_client=prisma_client, user_api_key_cache=cache) + if principal.team_id is not None + else None + ) + user = ( + await get_user_object( + user_id=principal.user_id, prisma_client=prisma_client, user_api_key_cache=cache, user_id_upsert=False + ) + if principal.user_id is not None + else None + ) + end_user = ( + await get_end_user_object( + end_user_id=principal.end_user_id, prisma_client=prisma_client, user_api_key_cache=cache + ) + if principal.end_user_id is not None + else None + ) + project = ( + await get_project_object( + project_id=principal.project_id, prisma_client=prisma_client, user_api_key_cache=cache + ) + if principal.project_id is not None + else None + ) + skip_budget_checks = _should_skip_budget_checks( + request_data=data, route=route, request=None, llm_router=llm_router + ) + await _authorize_replay( + principal=principal, + data=data, + model_group=model_group, + route=route, + team=team, + user=user, + end_user=end_user, + project=project, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + skip_budget_checks=skip_budget_checks, + ) + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=principal, + request_data=data, + route=route, + llm_router=llm_router, + team_object=team, + user_object=user, + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + skip_budget_checks=skip_budget_checks, + general_settings=general_settings, + end_user_id=principal.end_user_id, + end_user_object=end_user, + ) + + @staticmethod + async def _report_rejection( + *, + data: "dict[str, object]", + principal: "UserAPIKeyAuth", + exc: Exception, + proxy_logging_obj: "ProxyLogging", + ) -> None: + """The same entry point the proxy calls when a request fails after pre_call_hook + (common_request_processing.py:2538); it returns the parallel slot, the reserved TPM tokens and the + budget reservation, which would otherwise be stranded until their TTLs and 429 real traffic. + Idempotent via the callbacks' own released markers.""" + with contextlib.suppress(Exception): + await proxy_logging_obj.post_call_failure_hook( + request_data=data, original_exception=exc, user_api_key_dict=principal + ) + + async def _fetch_key_states( + self, + prisma_client: "PrismaClient | None", + user_api_key_cache: "DualCache | None", + key_hashes: frozenset[str], + ) -> "Mapping[str, UserAPIKeyAuth] | None": + """Cache-aware key state via the proxy resolver; a key whose lookup raises (deleted keys raise + token_not_found_in_db) is absent from the map and therefore excluded, isolating one bad key.""" + if not key_hashes: + return {} + + async def _resolve(hashed_token: str) -> "tuple[str, UserAPIKeyAuth | None]": + try: + return (hashed_token, await self.key_state_resolver(hashed_token, prisma_client, user_api_key_cache)) + except Exception: # noqa: BLE001 # an unverifiable key stays fail-closed for that key only + verbose_router_logger.warning( + "cache_warming could not verify a key's state; skipping its sessions this tick", exc_info=True + ) + return (hashed_token, None) + + resolved = await asyncio.gather(*(_resolve(hashed_token) for hashed_token in key_hashes)) + return {hashed_token: state for hashed_token, state in resolved if state is not None} diff --git a/litellm/router_strategy/complexity_router/cache_warming/store.py b/litellm/router_strategy/complexity_router/cache_warming/store.py new file mode 100644 index 00000000000..1a77bbedd9b --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/store.py @@ -0,0 +1,222 @@ +import time +from collections.abc import Awaitable, Mapping +from functools import lru_cache +from typing import Callable + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_router_logger +from litellm.caching.redis_cache import RedisCache +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_RECORD_SCHEMA_VERSION, + CacheWarmingAttribution, + CacheWarmingRecord, +) + +_WARMTH_KEY_PREFIX = "complexity_router_cache_warmth:v1" + +_CAPTURE_SCRIPT = """ +local sessions_key = KEYS[1] +local index_key = KEYS[2] +local member = ARGV[1] +local record_json = ARGV[2] +local now = tonumber(ARGV[3]) +local expires_at = tonumber(ARGV[4]) +local max_sessions = tonumber(ARGV[5]) +local expired = redis.call('ZRANGEBYSCORE', index_key, 0, now) +if #expired > 0 then + redis.call('HDEL', sessions_key, unpack(expired)) + redis.call('ZREMRANGEBYSCORE', index_key, 0, now) +end +if not redis.call('ZSCORE', index_key, member) and redis.call('ZCARD', index_key) >= max_sessions then + return 0 +end +redis.call('HSET', sessions_key, member, record_json) +redis.call('ZADD', index_key, expires_at, member) +redis.call('EXPIREAT', sessions_key, math.ceil(expires_at)) +redis.call('EXPIREAT', index_key, math.ceil(expires_at)) +return 1 +""" + +_LIST_LIVE_SESSIONS_SCRIPT = """ +local index_key = KEYS[1] +local now = tonumber(ARGV[1]) +local limit = tonumber(ARGV[2]) +return redis.call('ZRANGEBYSCORE', index_key, '(' .. now, '+inf', 'LIMIT', 0, limit) +""" + +_GET_RECORD_SCRIPT = """ +return redis.call('HGET', KEYS[1], ARGV[1]) +""" + +_MEMBERS_ADAPTER: TypeAdapter[tuple[str | bytes, ...]] = TypeAdapter(tuple[str | bytes, ...]) + + +@lru_cache(maxsize=64) +def _warn_redis_missing(auto_router_model_name: str) -> None: + verbose_router_logger.warning( + "cache_warming is enabled for auto-router %s but the router cache has no Redis; " + "cache warming is inactive until Redis is configured", + auto_router_model_name, + ) + + +@lru_cache(maxsize=64) +def _warn_session_cap_reached(auto_router_model_name: str) -> None: + verbose_router_logger.warning( + "cache_warming: auto-router %s reached max_sessions; new sessions are not captured until " + "existing records expire or go idle", + auto_router_model_name, + ) + + +def _parse_record(raw: object) -> CacheWarmingRecord | None: + if not isinstance(raw, (str, bytes)): + return None + try: + record = CacheWarmingRecord.model_validate_json(raw) + except ValidationError: + return None + if record.schema_version != CACHE_WARMING_RECORD_SCHEMA_VERSION: + return None + return record + + +def _parse_warmth(raw: object) -> float | None: + if isinstance(raw, (int, float)): + return float(raw) + if isinstance(raw, (str, bytes)): + try: + return float(raw) + except ValueError: + return None + return None + + +class CacheWarmingStore: + """Per-session capture state for provider prompt-cache warming. + + Session records and their expiry index live in two hash-tagged keys on one + Redis Cluster slot, so every capture is a single atomic Lua operation that + prunes expired sessions, enforces the max_sessions cap exactly, and writes + the record all-or-nothing. No partial write states exist, so no + compensation or fencing is needed anywhere, and the cap bounds the slot's + footprint by construction. Warmth stamps are plain single-writer keys with + their own TTL. Records are last-writer-wins by design: the latest turn is + the correct replay payload. A script fault raises and fails closed, never + an empty result mistaken for capacity.""" + + def __init__(self, redis_cache: RedisCache | None, auto_router_model_name: str) -> None: + self.redis_cache = redis_cache + self.auto_router_model_name = auto_router_model_name + register = redis_cache.async_register_script if redis_cache is not None else None + self._capture: Callable[..., Awaitable[object]] | None = register(_CAPTURE_SCRIPT) if register else None + self._list_live: Callable[..., Awaitable[object]] | None = ( + register(_LIST_LIVE_SESSIONS_SCRIPT) if register else None + ) + self._get: Callable[..., Awaitable[object]] | None = register(_GET_RECORD_SCRIPT) if register else None + + @staticmethod + def record_key(auto_router_model_name: str, caller_scope: str, session_id: str) -> str: + """Scoped by auto-router as well as caller and session. The record hash is already per-router through + its hash-tagged container, but warmth keys are derived from this identity and live at the top level, + so without the router in it two warming auto-routers sharing one Redis read each other's warmth.""" + return f"{auto_router_model_name}:{caller_scope}:{session_id}" + + @staticmethod + def warmth_key(record_key: str, model_group: str) -> str: + """Hash-tagged into the same slot family as the record it belongs to, so a Cluster keeps a session's + record and its warmth stamps on one node and the router scope cannot be dropped.""" + auto_router_model_name, _, session_scope = record_key.partition(":") + return f"{{cache_warm:v1:{auto_router_model_name}}}:{_WARMTH_KEY_PREFIX}:{session_scope}:{model_group}" + + def sessions_key(self) -> str: + return f"{{cache_warm:v1:{self.auto_router_model_name}}}:sessions" + + def index_key(self) -> str: + return f"{{cache_warm:v1:{self.auto_router_model_name}}}:index" + + def _require_redis(self) -> RedisCache | None: + if self.redis_cache is None: + _warn_redis_missing(self.auto_router_model_name) + return None + return self.redis_cache + + async def get_record(self, key: str) -> CacheWarmingRecord | None: + if self._require_redis() is None or self._get is None: + return None + raw = await self._get(keys=[self.sessions_key()], args=[key]) + return _parse_record(raw) + + async def upsert_session( + self, + *, + caller_scope: str, + session_id: str, + payload_compressed: str, + payload_sha256: str, + token_estimate: int, + served_model: str, + attribution: CacheWarmingAttribution, + ttl_seconds: int, + max_sessions: int, + ) -> None: + redis_cache = self._require_redis() + if redis_cache is None or self._capture is None: + return + key = self.record_key(self.auto_router_model_name, caller_scope, session_id) + now = time.time() + record = CacheWarmingRecord( + schema_version=CACHE_WARMING_RECORD_SCHEMA_VERSION, + payload_compressed=payload_compressed, + payload_sha256=payload_sha256, + token_estimate=token_estimate, + last_activity=now, + served_model=served_model, + session_id=session_id, + attribution=attribution, + auto_router_model_name=self.auto_router_model_name, + ) + try: + admitted = await self._capture( + keys=[self.sessions_key(), self.index_key()], + args=[key, record.model_dump_json(), now, now + ttl_seconds, max_sessions], + ) + except Exception: # noqa: BLE001 # a capture fault fails closed: no capture beats an uncapped write + verbose_router_logger.warning("cache_warming capture script failed; skipping capture", exc_info=True) + return + if admitted != 1: + _warn_session_cap_reached(self.auto_router_model_name) + return + await self.mark_warm_attempt(key, served_model, attempted_at=now, ttl_seconds=ttl_seconds) + + async def mark_warm_attempt(self, key: str, model_group: str, attempted_at: float, ttl_seconds: int) -> None: + redis_cache = self._require_redis() + if redis_cache is None: + return + await redis_cache.async_set_cache( # pyright: ignore[reportUnknownMemberType] # RedisCache is legacy-untyped + key=self.warmth_key(key, model_group), value=attempted_at, ttl=ttl_seconds + ) + + async def get_warmth(self, key: str, model_groups: tuple[str, ...]) -> Mapping[str, float]: + redis_cache = self._require_redis() + if redis_cache is None: + return {} # mutable-ok: fresh per-call result, not shared state + return { # mutable-ok: fresh per-call result, not shared state + model_group: stamp + for model_group in model_groups + if ( + stamp := _parse_warmth( + await redis_cache.async_get_cache(self.warmth_key(key, model_group)) # pyright: ignore[reportUnknownMemberType] # RedisCache is legacy-untyped + ) + ) + is not None + } + + async def list_session_keys(self, max_sessions: int) -> tuple[str, ...]: + if self._require_redis() is None or self._list_live is None: + return () + members = _MEMBERS_ADAPTER.validate_python( + await self._list_live(keys=[self.index_key()], args=[time.time(), max_sessions]) + ) + return tuple(member.decode("utf-8") if isinstance(member, bytes) else member for member in members) diff --git a/litellm/router_strategy/complexity_router/cache_warming/types.py b/litellm/router_strategy/complexity_router/cache_warming/types.py new file mode 100644 index 00000000000..7308495f514 --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/types.py @@ -0,0 +1,91 @@ +import base64 +import hashlib +import json +import zlib +from functools import lru_cache +from collections.abc import Mapping +from typing import Literal + +from pydantic import BaseModel, ConfigDict + +from litellm._logging import verbose_router_logger + +CACHE_WARMING_REPLAY_MARKER_KEY = "litellm_cache_warming" +CACHE_WARMING_REPLAY_TAG = "litellm_cache_warming" +CACHE_WARMING_RECORD_SCHEMA_VERSION = 1 +WARM_FRESHNESS_SLACK_SECONDS = 60 +# Anthropic and Bedrock hold a cached prefix for about five minutes. Freshness is that TTL and nothing else: +# the operator's idle_timeout_seconds and refresh_interval_seconds may legally exceed it, and a model whose +# stamp is older than the TTL is cold no matter which of them says otherwise. +PROVIDER_PROMPT_CACHE_TTL_SECONDS = 300 + + +def is_cache_fresh(warmed_at: float, now: float) -> bool: + """The one definition of "the provider still holds this prefix", read by both the refresher's due-model + calculation and the router's warm-aware pick, so a model can never be preferred as warm by one while the + other treats it as stale.""" + return now - warmed_at < PROVIDER_PROMPT_CACHE_TTL_SECONDS + + +def needs_rewarming(warmed_at: float, now: float, refresh_interval_seconds: int) -> bool: + """Due when the operator's interval has elapsed or the provider TTL is about to lapse, whichever comes + first, so a refresh_interval longer than the TTL cannot open a window where the pick still believes a + model is warm.""" + return now - warmed_at >= min( + refresh_interval_seconds, PROVIDER_PROMPT_CACHE_TTL_SECONDS - WARM_FRESHNESS_SLACK_SECONDS + ) + + +class CacheWarmingPayload(BaseModel): + model_config = ConfigDict(extra="forbid") + + model: str + messages: tuple[Mapping[str, object], ...] + system: str | tuple[Mapping[str, object], ...] | None = None + tools: tuple[Mapping[str, object], ...] | None = None + tool_choice: str | Mapping[str, object] | None = None + call_surface: Literal["chat_completions", "anthropic_messages"] + + +class CacheWarmingAttribution(BaseModel): + model_config = ConfigDict(extra="forbid") + + user_api_key: str | None = None + user_api_key_hash: str | None = None + user_api_key_user_id: str | None = None + user_api_key_team_id: str | None = None + user_api_key_org_id: str | None = None + user_api_key_project_id: str | None = None + user_api_key_end_user_id: str | None = None + + +class CacheWarmingRecord(BaseModel): + model_config = ConfigDict(extra="forbid") + + schema_version: int + payload_compressed: str + payload_sha256: str + token_estimate: int + last_activity: float + served_model: str + session_id: str | None = None + attribution: CacheWarmingAttribution + auto_router_model_name: str + + +def compress_payload(payload: CacheWarmingPayload) -> tuple[str, str]: + raw = payload.model_dump_json().encode("utf-8") + blob = base64.b64encode(zlib.compress(raw)).decode("ascii") + sha = hashlib.sha256(json.dumps(payload.model_dump(), sort_keys=True).encode("utf-8")).hexdigest() + return blob, sha + + +def decompress_payload(blob_b64: str) -> CacheWarmingPayload: + raw = zlib.decompress(base64.b64decode(blob_b64.encode("ascii"))) + return CacheWarmingPayload.model_validate_json(raw) + + +@lru_cache(maxsize=4096) +def warn_once(message: str) -> None: + """Package-wide warn-once; keyed on the formatted message, so per-entity messages dedupe per entity.""" + verbose_router_logger.warning(message) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 933c6d170cf..592cd75c686 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -18,7 +18,8 @@ from __future__ import annotations import asyncio import random import re -from collections.abc import Mapping +import time +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Union, cast from pydantic import BaseModel @@ -26,6 +27,7 @@ from pydantic import BaseModel from litellm._logging import verbose_router_logger from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import get_request_metadata_field from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.types.utils import ( ModelResponse, @@ -50,6 +52,7 @@ if TYPE_CHECKING: from litellm.router import Router from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter + from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore from litellm.types.router import PreRoutingHookResponse else: Router = Any @@ -619,6 +622,9 @@ class ComplexityRouter(CustomLogger): request_kwargs: dict, ) -> str: if not self.config.plugins: + warm_pick = await self._warm_aware_pick(self._tier_pools().get(tier.value, []), request_kwargs) + if warm_pick is not None: + return warm_pick return self.get_model_for_tier(tier) from litellm.types.router import RoutingContext @@ -641,8 +647,37 @@ class ComplexityRouter(CustomLogger): # silently bypassed. Raise instead, matching the Router-level plugin # pipeline's own fail-closed behavior for the same situation. raise ValueError(f"No candidate models left for tier {tier_key} after routing-plugin filtering") + warm_pick = await self._warm_aware_pick(context.candidate_models, request_kwargs) + if warm_pick is not None: + return warm_pick return self._pick_from_tier_value(context.candidate_models, tier_key) + async def _warm_aware_pick(self, pool: Sequence[str], request_kwargs: Mapping[str, object]) -> str | None: + config = self.config.cache_warming + if not config.enabled or len(pool) <= 1: + return None + session_id = get_request_metadata_field(request_kwargs, "session_id") + if session_id is None: + return None + store = self.get_cache_warming_store() + if store is None or store.redis_cache is None: + return None + from litellm.router_strategy.complexity_router.cache_warming.types import is_cache_fresh + + caller_scope = get_request_metadata_field(request_kwargs, "user_api_key_hash") or "unscoped" + record_key = store.record_key(self.model_name, caller_scope, session_id) + record = await store.get_record(record_key) + if record is None: + return None + warmth = await store.get_warmth(record_key, tuple(pool)) + now = time.time() + warmed = frozenset(model for model, warmed_at in warmth.items() if is_cache_fresh(warmed_at, now)) + served = frozenset((record.served_model,)) if is_cache_fresh(record.last_activity, now) else frozenset[str]() + candidates = tuple(model for model in pool if model in warmed | served) + if not candidates: + return None + return random.choice(candidates) + def _ensure_adaptive_router(self) -> Any | None: if not self.config.adaptive: return None @@ -1047,42 +1082,33 @@ class ComplexityRouter(CustomLogger): return user_message, system_prompt - @staticmethod - def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: - """Metadata may land on `metadata` or `litellm_metadata` depending on the - endpoint, mirroring DeploymentAffinityCheck's precedence.""" - return [ - metadata - for metadata_key in ("litellm_metadata", "metadata") - if isinstance(metadata := request_kwargs.get(metadata_key), dict) - ] + async def _capture_session( + self, request_kwargs: Mapping[str, object], messages: Sequence[Mapping[str, object]] | None, routed_model: str + ) -> None: + if not self.config.cache_warming.enabled: + return + from litellm.router_strategy.complexity_router.cache_warming.capture import capture_session - @staticmethod - def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: - """Resolve a client-supplied session_id.""" - for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): - session_id = metadata.get("session_id") - if session_id is not None: - return str(session_id) - return None + try: + await capture_session(self, request_kwargs, messages, routed_model) + except Exception: # noqa: BLE001 # a capture failure must never fail the user's request + verbose_router_logger.exception("cache_warming capture failed; the request continues unaffected") - @staticmethod - def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None: - """Resolve the proxy-derived API key hash, the same trust boundary - DeploymentAffinityCheck uses for its own key-based affinity (not the - client-supplied OpenAI `user` param, which isn't authenticated).""" - for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): - user_key = metadata.get("user_api_key_hash") - if user_key is not None: - return str(user_key) - return None + def get_cache_warming_store(self) -> CacheWarmingStore | None: + if not self.config.cache_warming.enabled: + return None + from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore + + router = self.litellm_router_instance + redis_cache = router.cache.redis_cache if router is not None and router.cache is not None else None + return CacheWarmingStore(redis_cache=redis_cache, auto_router_model_name=self.model_name) def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> str: # Namespace by the caller's API key hash so two different callers reusing the # same client-supplied session_id can't poison each other's routing pin. Falls # back to "unscoped" only when there's no authenticated caller to scope by # (e.g. direct Router usage without the proxy layer). - caller_scope = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" + caller_scope = get_request_metadata_field(request_kwargs, "user_api_key_hash") or "unscoped" return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" async def async_pre_routing_hook( @@ -1114,7 +1140,7 @@ class ComplexityRouter(CustomLogger): metadata[RETURN_RAW_MODEL_NAME_METADATA_KEY] = True use_session_affinity = self.config.session_affinity and not self.config.plugins - session_id = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None + session_id = get_request_metadata_field(request_kwargs, "session_id") if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None if cache_key is not None: @@ -1149,6 +1175,7 @@ class ComplexityRouter(CustomLogger): kwargs_metadata = request_kwargs.setdefault("metadata", {}) if isinstance(kwargs_metadata, dict): kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = routed_model + await self._capture_session(request_kwargs, messages, routed_model) escalated = routed_model != pinned_model cause: RoutingDecisionCause = "session_affinity_escalation" if escalated else "session_affinity_pin" verbose_router_logger.info( @@ -1179,6 +1206,8 @@ class ComplexityRouter(CustomLogger): value=response.model, ttl=self.config.session_affinity_ttl_seconds, ) + if response is not None: + await self._capture_session(request_kwargs, messages, response.model) return response async def _classify_and_route( diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 7437138fbb7..48471f5f654 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -248,6 +248,54 @@ class ClassifierLLMConfig(BaseModel): ) +class CacheWarmingConfig(BaseModel): + """Configuration for multi-model provider prompt-cache warming.""" + + enabled: bool = Field( + default=False, + description=( + "Capture each session's latest payload and keep provider prompt caches warm on every " + "tier model via background max_tokens=1 replays, so mid-session tier switches stay " + "cache hits. Requires Redis on the router cache and store_prompts_in_spend_logs=True" + ), + ) + refresh_interval_seconds: int = Field( + default=270, + gt=0, + description="Replay cadence per model; keep under the provider's cache TTL (Anthropic: 5 minutes)", + ) + session_ttl_seconds: int = Field( + default=3600, + gt=0, + description="TTL for the stored session payload record in Redis", + ) + idle_timeout_seconds: int = Field( + default=600, + gt=0, + description="Stop warming a session after this long without a real request; resumes on the next turn", + ) + max_sessions: int = Field( + default=1000, + gt=0, + description=( + "Exact upper bound on concurrently warmed sessions per auto-router, enforced by an atomic Redis " + "admission index; new sessions past the cap are not captured until existing ones expire" + ), + ) + max_payload_bytes: int = Field( + default=1_048_576, + gt=0, + description="Skip capturing sessions whose compressed payload exceeds this size", + ) + warm_models: tuple[str, ...] | None = Field( + default=None, + description=( + "Explicit model groups to keep warm; defaults to the first member of each tier pool. " + "Only Anthropic/Bedrock models that support prompt caching are warmed" + ), + ) + + class ComplexityRouterConfig(BaseModel): """Configuration for the ComplexityRouter.""" @@ -399,6 +447,11 @@ class ComplexityRouterConfig(BaseModel): description="TTL for the session affinity pin; refreshed on every cache hit", ) + cache_warming: CacheWarmingConfig = Field( + default_factory=CacheWarmingConfig, + description="Multi-model provider prompt-cache warming; disabled by default", + ) + plugins: list[RoutingPlugin] | None = Field( default=None, description="RoutingPlugin instances that narrow the classified tier's candidate models before selection", @@ -466,6 +519,15 @@ class ComplexityRouterConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_cache_warming_adaptive_combo(self) -> "ComplexityRouterConfig": + if self.cache_warming.enabled and self.adaptive: + raise ValueError( + "cache_warming and adaptive=True cannot both be set: adaptive's bandit selection doesn't yet " + "consult warm-cache state. Disable adaptive or disable cache_warming." + ) + return self + # Combined default config DEFAULT_COMPLEXITY_CONFIG = ComplexityRouterConfig() diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index c3f58935ef0..23d31de76c2 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -221,62 +221,19 @@ class DeploymentAffinityCheck(CustomLogger): def get_session_affinity_cache_key(cls, model_group: str, session_id: str) -> str: return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{session_id}" - @staticmethod - def _get_user_key_from_metadata_dict(metadata: dict) -> Optional[str]: - # NOTE: affinity is keyed on the *API key hash* provided by the proxy (not the - # OpenAI `user` parameter, which is an end-user identifier). - user_key = metadata.get("user_api_key_hash") - if user_key is None: - return None - return str(user_key) - - @staticmethod - def _get_session_id_from_metadata_dict(metadata: dict) -> Optional[str]: - session_id = metadata.get("session_id") - if session_id is None: - return None - return str(session_id) - - @staticmethod - def _iter_metadata_dicts(request_kwargs: dict) -> List[dict]: - """ - Return all metadata dicts available on the request. - - Depending on the endpoint, Router may populate `metadata` or `litellm_metadata`. - Users may also send one or both, so we check both (rather than using `or`). - """ - metadata_dicts: List[dict] = [] - for key in ("litellm_metadata", "metadata"): - md = request_kwargs.get(key) - if isinstance(md, dict): - metadata_dicts.append(md) - return metadata_dicts - @staticmethod def _get_user_key_from_request_kwargs(request_kwargs: dict) -> Optional[str]: - """ - Extract a stable affinity key from request kwargs. + """Affinity is keyed on the proxy-provided API key HASH, never the OpenAI ``user`` param + (an end-user identifier). Metadata precedence is owned by core_helpers.""" + from litellm.litellm_core_utils.core_helpers import get_request_metadata_field - Source (proxy): `metadata.user_api_key_hash` - - Note: the OpenAI `user` parameter is an end-user identifier and is intentionally - not used for deployment affinity. - """ - # Check metadata dicts (Proxy usage) - for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): - user_key = DeploymentAffinityCheck._get_user_key_from_metadata_dict(metadata=metadata) - if user_key is not None: - return user_key - - return None + return get_request_metadata_field(request_kwargs, "user_api_key_hash") # pyright: ignore[reportUnknownArgumentType] # legacy-untyped call site @staticmethod def _get_session_id_from_request_kwargs(request_kwargs: dict) -> Optional[str]: - for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): - session_id = DeploymentAffinityCheck._get_session_id_from_metadata_dict(metadata=metadata) - if session_id is not None: - return session_id - return None + from litellm.litellm_core_utils.core_helpers import get_request_metadata_field + + return get_request_metadata_field(request_kwargs, "session_id") # pyright: ignore[reportUnknownArgumentType] # legacy-untyped call site @staticmethod def _find_deployment_by_model_id(healthy_deployments: List[dict], model_id: str) -> Optional[dict]: @@ -416,7 +373,9 @@ class DeploymentAffinityCheck(CustomLogger): - LiteLLM runs async success callbacks via a background logging worker for performance. - We want affinity to be immediately available for subsequent requests. """ - metadata_dicts = self._iter_metadata_dicts(kwargs) + from litellm.litellm_core_utils.core_helpers import iter_request_metadata_dicts + + metadata_dicts = iter_request_metadata_dicts(kwargs) # Extract deployment_model_name first — needed for both per-group flag resolution # and cache key scoping. diff --git a/tests/test_litellm/proxy/proxy_server/test_background_health.py b/tests/test_litellm/proxy/proxy_server/test_background_health.py index dca93e137ac..1d42c9e6eb2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_background_health.py +++ b/tests/test_litellm/proxy/proxy_server/test_background_health.py @@ -8,6 +8,7 @@ Pins covered: - ``_get_endpoint_exception_status`` - ``_write_health_state_to_router_cache`` - ``_adaptive_router_flusher_loop`` +- ``_complexity_cache_warming_loop`` - ``_run_background_health_check`` """ @@ -22,6 +23,7 @@ import pytest import litellm.proxy.proxy_server as proxy_server from litellm.proxy.proxy_server import ( _adaptive_router_flusher_loop, + _complexity_cache_warming_loop, _get_endpoint_exception_status, _get_process_rss_mb, _run_background_health_check, @@ -427,6 +429,76 @@ async def test_adaptive_router_flusher_loop_times_out_when_sleep_real(monkeypatc await asyncio.wait_for(_adaptive_router_flusher_loop(), timeout=0.2) +# --------------------------------------------------------------------------- +# _complexity_cache_warming_loop +# --------------------------------------------------------------------------- + + +class _RecordingRefresher: + def __init__(self, fail_on_call: int | None = None): + self.calls: list[dict] = [] + self.fail_on_call = fail_on_call + + async def run_tick(self, *, llm_router, prisma_client, user_api_key_cache=None): + self.calls.append( + {"llm_router": llm_router, "prisma_client": prisma_client, "user_api_key_cache": user_api_key_cache} + ) + if self.fail_on_call == len(self.calls): + raise RuntimeError("tick boom") + + +def _cancel_sleep_after(monkeypatch, iterations: int, on_call=None): + call_count = {"n": 0} + _real_sleep = asyncio.sleep + + async def _short_sleep(_seconds): + call_count["n"] += 1 + if on_call is not None: + on_call(call_count["n"]) + if call_count["n"] > iterations: + raise asyncio.CancelledError() + await _real_sleep(0) + + monkeypatch.setattr(proxy_server.asyncio, "sleep", _short_sleep) + + +@pytest.mark.asyncio +async def test_complexity_cache_warming_loop_survives_tick_exception(monkeypatch): + refresher = _RecordingRefresher(fail_on_call=1) + monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) + _cancel_sleep_after(monkeypatch, iterations=2) + + with pytest.raises(asyncio.CancelledError): + await _complexity_cache_warming_loop(refresher=refresher) + + assert len(refresher.calls) == 2 + + +@pytest.mark.asyncio +async def test_complexity_cache_warming_loop_noop_while_router_none_and_resolves_router_fresh(monkeypatch): + refresher = _RecordingRefresher() + fake_router = MagicMock() + fake_prisma = MagicMock() + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + + def _swap_router_in(call_number: int): + if call_number == 2: + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + + _cancel_sleep_after(monkeypatch, iterations=2, on_call=_swap_router_in) + + with pytest.raises(asyncio.CancelledError): + await _complexity_cache_warming_loop(refresher=refresher) + + assert len(refresher.calls) == 1 + assert refresher.calls[0]["llm_router"] is fake_router + assert refresher.calls[0]["prisma_client"] is fake_prisma + assert refresher.calls[0]["user_api_key_cache"] is proxy_server.user_api_key_cache + + + + # --------------------------------------------------------------------------- # _run_background_health_check # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/router_strategy/complexity_router/__init__.py b/tests/test_litellm/router_strategy/complexity_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/__init__.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py new file mode 100644 index 00000000000..968d178e4c1 --- /dev/null +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py @@ -0,0 +1,176 @@ +import json +import logging +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +import litellm +from litellm.router_strategy.complexity_router.cache_warming.capture import ( + _warn_payload_too_large, + _warn_privacy_gate_blocked, +) +from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_REPLAY_MARKER_KEY, + decompress_payload, +) +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter + +from tests.test_litellm.router_strategy.complexity_router.cache_warming.test_store import FakeRedisCache + +LONG_SYSTEM = "All deployment manifests must declare resource ceilings before rollout. " * 200 + +SESSIONS_KEY = "{cache_warm:v1:smart-router}:sessions" + + +@pytest.fixture(autouse=True) +def _prompt_retention_consent(monkeypatch): + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + + +def anthropic_messages(**kwargs: object) -> None: + raise AssertionError("marker function; never called") + + +def _complexity_router(redis: FakeRedisCache | None, **cache_warming_overrides: object) -> ComplexityRouter: + router_instance = MagicMock() + router_instance.cache = SimpleNamespace(redis_cache=redis) + return ComplexityRouter( + model_name="smart-router", + litellm_router_instance=router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "gpt-5-mini", "COMPLEX": "claude-sonnet-4-5"}, + "cache_warming": {"enabled": True, **cache_warming_overrides}, + }, + ) + + +MESSAGES = [ + {"role": "system", "content": LONG_SYSTEM}, + {"role": "user", "content": "summarize rule 7"}, +] + + +def _kwargs(**overrides: object) -> dict: + base: dict = { + "model": "smart-router", + "metadata": { + "session_id": "sess-1", + "user_api_key_hash": "hash-1", + "user_api_key": "hash-1", + "user_api_key_team_id": "team-9", + }, + } + return {**base, **overrides} + + +def _stored_records(redis: FakeRedisCache) -> list[dict]: + return [json.loads(value) for value in redis.hashes.get(SESSIONS_KEY, {}).values()] + + + + + + + + + + + + + + + + + + + + + + + + + + +@pytest.mark.asyncio +async def test_second_turn_overwrites_payload_and_preserves_other_model_warmth(): + redis = FakeRedisCache() + router = _complexity_router(redis) + await router._capture_session(_kwargs(), MESSAGES, "claude-sonnet-4-5") + key = CacheWarmingStore.record_key("smart-router", "hash-1", "sess-1") + first = json.loads(redis.hashes[SESSIONS_KEY][key]) + redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")] = json.dumps(123.0) + await router._capture_session( + _kwargs(), MESSAGES + [{"role": "user", "content": "and rule 8?"}], "claude-sonnet-4-5" + ) + second = json.loads(redis.hashes[SESSIONS_KEY][key]) + assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")]) == 123.0 + assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "claude-sonnet-4-5")]) > 0 + assert second["payload_sha256"] != first["payload_sha256"] + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +@pytest.mark.asyncio +@pytest.mark.parametrize("opt_out", ["no-retention-consent", "global-redaction", "per-request-redaction"]) +async def test_capture_honors_every_form_of_the_operators_prompt_retention_policy(opt_out, monkeypatch): + """Capture persists full prompts, so it requires the retention opt-in and respects message redaction in + every form it can be expressed.""" + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs() + _warn_privacy_gate_blocked.cache_clear() + if opt_out == "no-retention-consent": + monkeypatch.delenv("STORE_PROMPTS_IN_SPEND_LOGS", raising=False) + elif opt_out == "global-redaction": + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + else: + kwargs["metadata"]["headers"] = {"x-litellm-enable-message-redaction": True} + await router._capture_session(kwargs, MESSAGES, "claude-sonnet-4-5") + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_warming_does_not_accept_a_forged_billing_identity(): + """Identity comes from the one proxy-stamped slot, marked by user_api_key_hash which a caller cannot + inject because the proxy strips the user_api_key_ prefixed fields. A bare user_api_key survives that + strip and the slots are read litellm_metadata first, so merging them would let a caller name the + master-key alias and have the refresher skip its key-state verification entirely.""" + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs(litellm_metadata={"session_id": "sess-1", "user_api_key": LITELLM_PROXY_MASTER_KEY_ALIAS}) + await router._capture_session(kwargs, MESSAGES, "claude-sonnet-4-5") + record = _stored_records(redis)[0] + assert record["attribution"]["user_api_key"] == "hash-1" + assert record["attribution"]["user_api_key_hash"] == "hash-1" diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py new file mode 100644 index 00000000000..a3c28d56de1 --- /dev/null +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py @@ -0,0 +1,337 @@ +"""Refresher suite. + +Deliberately small. The live-proxy run in the PR body is the workflow evidence; what is here is the set of +behaviors whose regression would harm the customer's own traffic and would not be obvious from a replay +succeeding: warming must not double-charge a key's TPM, must not reset a key's rate-limit window, must not +collide the operator's hanging-request tracking, and must be refused by every ceiling the request path +enforces. The remaining two cover the shape of a replay on each surface and the two brakes on cost. +""" + +import time +from datetime import datetime, timedelta, timezone + +import pytest + +from litellm.caching.dual_cache import DualCache +from litellm.router_strategy.complexity_router.cache_warming.refresher import filter_cache_warmable +from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_REPLAY_MARKER_KEY, + CACHE_WARMING_REPLAY_TAG, +) + +from tests.test_litellm.router_strategy.complexity_router.cache_warming.test_store import FakeRedisCache +from tests.test_litellm.router_strategy.complexity_router.cache_warming.warming_rig import ( + UNIFORM_POOL, + FakeKeyDirectory, + FakeLeaseLock, + ReplayRouter, + affinity_check, + key_state, + priced_rig, + proxy_logging_with_hooks, + real_limiter, + refresher, + registered_callbacks, + replayed_models, + seed_session, + team, + tick, + warming_rig, + warmth_stamp, +) + +_PAST = (datetime.now(timezone.utc) - timedelta(hours=1)).replace(microsecond=0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "arm,warmed", + [ + ("healthy", True), + ("blocked-key", False), + ("expired-key", False), + ("unparseable-expiry", False), + ("blocked-team", False), + ("team-model-denied", False), + ("key-model-denied", False), + ("key-model-over-budget", False), + ("over-budget", False), + ("rate-limited", False), + ], +) +async def test_every_ceiling_the_request_path_enforces_gates_warming(arm, warmed): + """A replay is admitted through the request path's own entry points, so each of these refusals comes from + that path's owner rather than a second implementation here: the blocked and expired key checks mirroring + the canonical auth checks, then both of production's own authorization enumerations + (_enforce_key_and_fallback_model_access for the key-level model allowlist, common_checks for team, member, + user, project, every budget scope and the tool and vector-store allowlists), the budget reservation on the + shared spend counters, and the v3 limiter's own counters. The key-model-denied arm is the one that proves + warming runs the key-level enumeration and not only the team one; the unparseable-expiry arm is the + fail-closed half of a value that stays a plain string on a cached key object. + """ + from litellm.proxy.proxy_server import spend_counter_cache + + redis = FakeRedisCache() + limiter, counters = real_limiter() + key_cache = DualCache() + fields: dict = {} + token = "k" + + if arm == "blocked-key": + fields = {"blocked": True} + elif arm == "expired-key": + fields = {"expires": _PAST.isoformat().replace("+00:00", "Z")} + elif arm == "unparseable-expiry": + fields = {"expires": "not-a-timestamp"} + elif arm in ("blocked-team", "team-model-denied"): + team_object = team("t", blocked=arm == "blocked-team", models=[] if arm == "blocked-team" else ["fast-claude"]) + await key_cache.async_set_cache(key="team_id:t", value=team_object) + fields = {"team_id": "t"} + elif arm == "key-model-denied": + fields = {"models": ["fast-claude"]} + elif arm == "key-model-over-budget": + from litellm.proxy.hooks.model_max_budget_limiter import VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX + from litellm.proxy.proxy_server import user_api_key_cache as proxy_key_cache + + fields = {"model_max_budget": {"smart-claude": {"budget_limit": 1.0, "time_period": "1d"}}} + await proxy_key_cache.async_set_cache( + key=f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{token}:smart-claude:1d", value=5.0 + ) + elif arm == "over-budget": + fields = {"spend": 100.0, "max_budget": 100.0} + elif arm == "rate-limited": + fields = {"rpm_limit": 1} + await counters.async_set_cache(key="{api_key:k}:window", value=str(int(time.time()))) + await counters.async_set_cache(key="{api_key:k}:requests", value=1) + + priced = arm == "over-budget" + llm_router = priced_rig(redis) if priced else warming_rig(redis=redis)[0] + served = "claude-haiku-4-5" if priced else "fast-claude" + target = "claude-sonnet-4-5" if priced else "smart-claude" + counter_key = f"spend:key:{token}" + if arm == "over-budget": + await spend_counter_cache.async_set_cache(key=counter_key, value=100.0) + seed_session(redis, user_api_key=token, served_model=served, warmth={served: time.time()}) + keys = FakeKeyDirectory({token: key_state(token=token, **fields)}) + try: + with registered_callbacks(limiter): + await tick(llm_router, active=refresher(keys=keys, limiter=limiter), user_api_key_cache=key_cache) + assert (target in replayed_models(llm_router)) is warmed + finally: + spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) + + +@pytest.mark.asyncio +async def test_warming_does_not_double_charge_a_keys_tpm(): + """Both halves of the limiter's TPM contract on one replay: pre_call_hook reserves upfront, and the + limiter's own success callback finds that reservation through the request metadata the replay carried and + settles the counter to actual usage. Without the stash riding along the counter ends at reservation plus + actual, so the customer's next real request is throttled against tokens nobody used.""" + from datetime import datetime as _datetime + + from litellm.types.utils import ModelResponse, Usage + + limiter, counters = real_limiter() + llm_router, redis = warming_rig(redis=FakeRedisCache()) + seed_session(redis, user_api_key="tpm", warmth={"fast-claude": time.time()}) + keys = FakeKeyDirectory({"tpm": key_state(token="tpm", tpm_limit=100_000)}) + with registered_callbacks(limiter): + await tick(llm_router, active=refresher(keys=keys, limiter=limiter)) + tokens_key = limiter.create_rate_limit_keys("api_key", "tpm", "tokens") + assert int(await counters.async_get_cache(key=tokens_key) or 0) > 0 + replay_metadata = llm_router.completion_calls[0]["metadata"] + assert "_litellm_tpm_reserved_tokens" in replay_metadata + await limiter.async_log_success_event( + kwargs={"metadata": replay_metadata, "standard_logging_object": {"metadata": {"user_api_key_hash": "tpm"}}}, + response_obj=ModelResponse(usage=Usage(prompt_tokens=2000, completion_tokens=1, total_tokens=2001)), + start_time=_datetime.now(), + end_time=_datetime.now(), + ) + assert int(await counters.async_get_cache(key=tokens_key)) == 2001 + + +@pytest.mark.asyncio +async def test_warming_does_not_reset_a_keys_rate_limit_window(): + """The descriptor carries the limiter's own configured window, so an hour-long window opened two minutes + ago keeps accumulating. A hardcoded 60 would read as expired and overwrite the shared counters, + permanently un-enforcing the operator's real limit for that key.""" + limiter, counters = real_limiter(window_size=3600) + assert limiter.window_size == 3600 + llm_router, redis = warming_rig(redis=FakeRedisCache()) + seed_session(redis, user_api_key="hourly", warmth={"fast-claude": time.time()}) + keys = FakeKeyDirectory({"hourly": key_state(token="hourly", rpm_limit=100)}) + opened_at = int(time.time()) - 120 + await counters.async_set_cache(key="{api_key:hourly}:window", value=str(opened_at)) + await counters.async_set_cache(key="{api_key:hourly}:requests", value=5) + with registered_callbacks(limiter): + await tick(llm_router, active=refresher(keys=keys, limiter=limiter)) + assert len(llm_router.completion_calls) == 1 + assert int(await counters.async_get_cache(key="{api_key:hourly}:window")) == opened_at + assert int(await counters.async_get_cache(key="{api_key:hourly}:requests")) == 6 + + +@pytest.mark.asyncio +async def test_warming_does_not_collide_hanging_request_tracking(): + """The hanging-request checker keys its cache on litellm_call_id and clears an entry only when a request + status is recorded. Replays sharing the empty default would collide on one entry and alert the operator + about a request that never existed, indefinitely.""" + proxy_logging = proxy_logging_with_hooks() + proxy_logging.alerting = ["slack"] + llm_router, redis = warming_rig(redis=FakeRedisCache()) + seed_session(redis) + await tick(llm_router, active=refresher(proxy_logging=proxy_logging)) + call_ids = [call["litellm_call_id"] for call in llm_router.completion_calls] + assert len(call_ids) == 2 and len(set(call_ids)) == 2 + for call_id in call_ids: + assert ( + await proxy_logging.internal_usage_cache.async_get_cache( + key=f"request_status:{call_id}", litellm_parent_otel_span=None, local_only=True + ) + == "success" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface,channel", [("chat_completions", "metadata"), ("anthropic_messages", "litellm_metadata")]) +async def test_a_due_session_is_replayed_on_its_own_surface_and_stamped_warm(surface, channel): + """The replay reaches the surface it was captured from, through that surface's own metadata channel, with + generation held at the floor and the session_id that lets deployment affinity pin the replay to the same + member real traffic will hit. Anthropic invalidates a cached prefix when tool_choice changes, so it and + the system block ride along there. The warmth stamp is what paces the next tick.""" + llm_router, redis = warming_rig(redis=FakeRedisCache()) + record_key = seed_session( + redis, call_surface=surface, warmth={"fast-claude": time.time()}, tool_choice={"type": "auto"} + ) + await tick(llm_router) + calls = llm_router.anthropic_calls if surface == "anthropic_messages" else llm_router.completion_calls + other = llm_router.completion_calls if surface == "anthropic_messages" else llm_router.anthropic_calls + assert len(calls) == 1 and other == [] + call = calls[0] + assert call["model"] == "smart-claude" + assert call["max_tokens"] == 1 and call["stream"] is False + assert call["tool_choice"] == {"type": "auto"} + assert call[channel][CACHE_WARMING_REPLAY_MARKER_KEY] is True + assert call[channel]["session_id"] == "sess-1" + assert call[channel]["user_api_key"] == "hash-1" + assert call[channel]["spend_logs_metadata"] == {CACHE_WARMING_REPLAY_TAG: "true"} + assert "tags" not in call[channel] + if surface == "anthropic_messages": + assert call["system"] == "You are a policy assistant" + stamp = warmth_stamp(redis, record_key, "smart-claude") + assert stamp is not None and stamp > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "last_activity_offset,warmth_offset,replayed", + [(-601, None, []), (None, -10, []), (None, -300, ["smart-claude"])], + ids=["idle-session-drops-out", "recently-warmed-model-waits", "stale-warmth-is-refreshed"], +) +async def test_session_pacing_bounds_how_often_warming_spends(last_activity_offset, warmth_offset, replayed): + """Two independent brakes on cost: a session idle past its timeout stops being warmed at all, and a model + warmed inside the refresh interval is not warmed again.""" + llm_router, redis = warming_rig(redis=FakeRedisCache()) + now = time.time() + seed_session( + redis, + last_activity=now + last_activity_offset if last_activity_offset is not None else None, + warmth={"fast-claude": now, "smart-claude": now + warmth_offset} if warmth_offset is not None else None, + ) + await tick(llm_router) + assert replayed_models(llm_router) == replayed + + +@pytest.mark.parametrize("with_affinity,warmable", [(False, False), (True, True)]) +def test_a_multi_deployment_group_is_only_warmed_when_deployment_affinity_pins_it(with_affinity, warmable): + """Warming a pool without affinity pays the cache-write premium against 1/N routing odds, which is worse + than not warming, so such a group is skipped rather than degrading silently.""" + llm_router = ReplayRouter(model_list=UNIFORM_POOL) + if with_affinity: + llm_router.optional_callbacks = [affinity_check()] + assert filter_cache_warmable(llm_router, ["uniform"]) == (("uniform",) if warmable else ()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unavailable", ["database", "key-cache", "rate-limiter", "lock-backend"]) +async def test_no_session_is_warmed_while_any_enforcement_dependency_is_unavailable(unavailable, caplog): + """Every ceiling warming claims to respect is enforced by one of these, so losing any single one turns + admission into a no-op that still spends. They are checked together, once, before any session is + considered, rather than as per-dependency arms that each fail open on their own path. The warning is + asserted as well as the silence, because a replay that dies further down on a missing dependency also + makes no call and would let a removed guard pass unnoticed.""" + import logging + + from litellm.proxy.utils import ProxyLogging + from litellm.router_strategy.complexity_router.cache_warming.types import warn_once + + warn_once.cache_clear() + llm_router, redis = warming_rig(redis=FakeRedisCache()) + seed_session(redis, user_api_key="k") + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + if unavailable == "database": + await tick(llm_router, prisma=None) + elif unavailable == "key-cache": + await refresher().run_tick(llm_router=llm_router, prisma_client=object(), user_api_key_cache=None) + elif unavailable == "rate-limiter": + await tick(llm_router, active=refresher(proxy_logging=ProxyLogging(user_api_key_cache=DualCache()))) + else: + from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import ( + LockAcquisition, + ) + + await tick(llm_router, active=refresher(lock=FakeLeaseLock(acquisition=LockAcquisition.ERROR))) + assert llm_router.completion_calls == [] + assert any("cannot enforce a replay's ceilings" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_warmth_is_not_shared_between_two_auto_routers_on_one_redis(): + """Warmth keys derive from the session's scoped identity, which includes the auto-router. Without that a + second warming router reads the first one's stamps and skips replays it never made, so its tiers go + cold while it believes them warm.""" + redis = FakeRedisCache() + first, _ = warming_rig(redis=redis) + seed_session(redis, warmth={"fast-claude": time.time(), "smart-claude": time.time()}) + await tick(first) + assert first.completion_calls == [] + + other = CacheWarmingStore(redis_cache=redis, auto_router_model_name="other-router") + record = other.record_key("other-router", "hash-1", "sess-1") + assert await other.get_warmth(record, ("fast-claude", "smart-claude")) == {} + + + + +@pytest.mark.asyncio +@pytest.mark.parametrize("team_blocked,warmed", [(True, False), (False, True)]) +async def test_a_keyless_proxy_caller_is_authorized_through_its_reconstructed_tenancy(team_blocked, warmed): + """A JWT caller has api_key None, so before this it fell through to the unattributed path and every tenant + control was skipped. Its tenancy is recorded at capture, so the principal is rebuilt from it and the team + gate binds: a blocked team stops the replays, an open team still warms.""" + llm_router, redis = warming_rig(redis=FakeRedisCache()) + key_cache = DualCache() + await key_cache.async_set_cache( + key="team_id:jwt-team", value=team("jwt-team", blocked=team_blocked, models=[]) + ) + seed_session( + redis, + user_api_key=None, + caller_scope="jwt-user", + team_id="jwt-team", + warmth={"fast-claude": time.time()}, + ) + await tick(llm_router, active=refresher(keys=FakeKeyDirectory({})), user_api_key_cache=key_cache) + assert bool(llm_router.completion_calls) is warmed + if warmed: + assert llm_router.completion_calls[0]["metadata"]["user_api_key_team_id"] == "jwt-team" + + +@pytest.mark.asyncio +async def test_a_direct_sdk_session_with_no_recorded_identity_still_warms_unattributed(): + """No proxy auth object means no tenancy to preserve, so warming stays unattributed as before.""" + llm_router, redis = warming_rig(redis=FakeRedisCache()) + seed_session(redis, user_api_key=None, caller_scope="unscoped", warmth={"fast-claude": time.time()}) + await tick(llm_router, active=refresher(keys=FakeKeyDirectory({}))) + assert replayed_models(llm_router) == ["smart-claude"] + assert llm_router.completion_calls[0]["metadata"]["user_api_key_team_id"] is None diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py new file mode 100644 index 00000000000..eb17a64a7b9 --- /dev/null +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py @@ -0,0 +1,182 @@ +import json +import logging + +import pytest + +from litellm.router_strategy.complexity_router.cache_warming.store import ( + CacheWarmingStore, + _warn_redis_missing, + _warn_session_cap_reached, +) +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_RECORD_SCHEMA_VERSION, + CacheWarmingAttribution, + CacheWarmingRecord, +) + + +class FakeRedisCache: + def __init__(self, namespace: str | None = None) -> None: + self.namespace = namespace + self.data: dict[str, str] = {} + self.ttls: dict[str, int | None] = {} + self.hashes: dict[str, dict[str, str]] = {} + self.zsets: dict[str, dict[str, float]] = {} + self.expire_calls: list[str] = [] + + def _namespaced(self, key: str) -> str: + if self.namespace and not key.startswith(self.namespace): + return f"{self.namespace}:{key}" + return key + + async def async_get_cache(self, key: str, **kwargs: object) -> object: + raw = self.data.get(self._namespaced(key)) + return json.loads(raw) if raw is not None else None + + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: + namespaced = self._namespaced(key) + self.data[namespaced] = json.dumps(value) + ttl = kwargs.get("ttl") + self.ttls[namespaced] = ttl if isinstance(ttl, int) else None + + def async_register_script(self, script: str): + if 'redis.call("expire"' in script: + + async def compare_and_expire(keys: list, args: list) -> int: + key = self._namespaced(keys[0]) + raw = self.data.get(key) + if raw is not None and raw == str(args[0]): + self.ttls[key] = int(args[1]) + self.expire_calls.append(key) + return 1 + return 0 + + return compare_and_expire + if 'redis.call("del"' in script: + + async def compare_and_delete(keys: list, args: list) -> int: + key = self._namespaced(keys[0]) + raw = self.data.get(key) + if raw is not None and raw == str(args[0]): + del self.data[key] + return 1 + return 0 + + return compare_and_delete + if "HGET" in script: + + async def get_record(keys: list, args: list) -> str | None: + return self.hashes.get(self._namespaced(keys[0]), {}).get(str(args[0])) + + return get_record + if "HSET" in script: + + async def capture(keys: list, args: list) -> int: + sessions = self.hashes.setdefault(self._namespaced(keys[0]), {}) + index = self.zsets.setdefault(self._namespaced(keys[1]), {}) + member, record_json = str(args[0]), str(args[1]) + now, expires_at, max_sessions = float(args[2]), float(args[3]), int(args[4]) + for stale in [m for m, score in index.items() if score <= now]: + del index[stale] + sessions.pop(stale, None) + if member not in index and len(index) >= max_sessions: + return 0 + sessions[member] = record_json + index[member] = expires_at + return 1 + + return capture + + async def list_live(keys: list, args: list) -> list: + index = self.zsets.get(self._namespaced(keys[0]), {}) + now, limit = float(args[0]), int(args[1]) + live = sorted((score, member) for member, score in index.items() if score > now) + return [member.encode("utf-8") for _, member in live[:limit]] + + return list_live + + +def _record_json(**overrides: object) -> str: + base = CacheWarmingRecord( + schema_version=CACHE_WARMING_RECORD_SCHEMA_VERSION, + payload_compressed="blob", + payload_sha256="sha", + token_estimate=2048, + last_activity=1000.0, + served_model="sonnet", + attribution=CacheWarmingAttribution(user_api_key="hashed"), + auto_router_model_name="smart-router", + ).model_dump() + return json.dumps({**base, **overrides}) + + +def _store(redis: FakeRedisCache | None) -> CacheWarmingStore: + return CacheWarmingStore(redis_cache=redis, auto_router_model_name="smart-router") + + +async def _upsert(store: CacheWarmingStore, session_id: str = "s1", max_sessions: int = 100) -> None: + await store.upsert_session( + caller_scope="scope", + session_id=session_id, + payload_compressed="blob2", + payload_sha256="sha2", + token_estimate=4096, + served_model="sonnet", + attribution=CacheWarmingAttribution(), + ttl_seconds=1800, + max_sessions=max_sessions, + ) + + +def test_key_shapes_are_scoped_and_hash_tagged(): + """Every key a session owns carries the auto-router name and shares one Cluster hash tag with the others, + so two warming auto-routers on one Redis cannot read each other's warmth and a session's record, index + entry and warmth stamps stay on one node.""" + record = CacheWarmingStore.record_key("smart-router", "keyhash", "session-1") + assert record == "smart-router:keyhash:session-1" + other_router = CacheWarmingStore.record_key("other-router", "keyhash", "session-1") + assert CacheWarmingStore.warmth_key(record, "opus") != CacheWarmingStore.warmth_key(other_router, "opus") + store = _store(None) + slot = "{cache_warm:v1:smart-router}" + assert store.sessions_key() == f"{slot}:sessions" + assert store.index_key() == f"{slot}:index" + assert CacheWarmingStore.warmth_key(record, "opus").startswith(f"{slot}:") + + + + +@pytest.mark.asyncio +async def test_cap_enforced_atomically_with_the_record_write(): + _warn_session_cap_reached.cache_clear() + redis = FakeRedisCache() + store = _store(redis) + await _upsert(store, session_id="s1", max_sessions=2) + await _upsert(store, session_id="s2", max_sessions=2) + await _upsert(store, session_id="s3", max_sessions=2) + assert await store.get_record(store.record_key("smart-router", "scope", "s3")) is None + assert len(await store.list_session_keys(max_sessions=10)) == 2 + + + + + + + + + + +@pytest.mark.asyncio +async def test_get_record_returns_none_on_schema_version_mismatch(): + redis = FakeRedisCache() + store = _store(redis) + key = store.record_key("smart-router", "scope", "s1") + redis.hashes[store.sessions_key()] = {key: _record_json(schema_version=CACHE_WARMING_RECORD_SCHEMA_VERSION + 1)} + assert await store.get_record(key) is None + + + + + + + + diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_types.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_types.py new file mode 100644 index 00000000000..616d01b6c49 --- /dev/null +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_types.py @@ -0,0 +1,46 @@ +"""Freshness is one predicate, derived from the provider prompt-cache TTL.""" + +import pytest + +from litellm.router_strategy.complexity_router.cache_warming.types import ( + PROVIDER_PROMPT_CACHE_TTL_SECONDS, + is_cache_fresh, + needs_rewarming, +) + + +@pytest.mark.parametrize( + "age,fresh", + [(0, True), (PROVIDER_PROMPT_CACHE_TTL_SECONDS - 1, True), (PROVIDER_PROMPT_CACHE_TTL_SECONDS, False), (601, False)], +) +def test_freshness_is_the_provider_ttl_and_nothing_else(age, fresh): + """The router's warm-aware pick and the refresher's due-model calculation both read this, so a model can + never be preferred as warm by one while the other treats it as stale. An operator's idle_timeout_seconds + of 600 must not make a 400-second-old prefix look warm.""" + assert is_cache_fresh(1000.0 - age, 1000.0) is fresh + + +@pytest.mark.parametrize("refresh_interval,age,due", [(120, 119, False), (120, 120, True), (3000, 240, True)]) +def test_rewarming_never_waits_past_the_provider_ttl(refresh_interval, age, due): + """A short interval governs on its own; an interval longer than the TTL is capped by the TTL less the tick + slack, so it cannot open a window where the pick still believes a model is warm.""" + assert needs_rewarming(1000.0 - age, 1000.0, refresh_interval) is due + + +def test_attribution_covers_every_identity_field_the_proxy_stamps(): + """A replay is authorized against a principal rebuilt from attribution, so any identity id the proxy + stamps and the auth gates resolve by must be captured. If litellm adds one, this fails rather than + silently leaving that gate unbound for keyless callers.""" + from litellm.types.utils import StandardLoggingUserAPIKeyMetadata + + from litellm.router_strategy.complexity_router.cache_warming.types import CacheWarmingAttribution + + upstream_ids = { + field + for field in StandardLoggingUserAPIKeyMetadata.__annotations__ + if field.endswith("_id") and field.startswith("user_api_key_") + } + assert upstream_ids, "upstream identity metadata shape changed" + assert upstream_ids <= set(CacheWarmingAttribution.model_fields), ( + f"attribution is missing identity fields the proxy stamps: {upstream_ids - set(CacheWarmingAttribution.model_fields)}" + ) diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py new file mode 100644 index 00000000000..d56e6a85efd --- /dev/null +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py @@ -0,0 +1,330 @@ +"""Shared rig for the cache-warming refresher suites. + +Everything here is either the production object itself or the narrowest possible stand-in for a backend +the test process has no access to (Redis, the key database, the pod lock). +""" + +import asyncio +import json +import os +import time +from contextlib import contextmanager + +import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.proxy.utils import ProxyLogging +from litellm.router_strategy.complexity_router.cache_warming.refresher import CacheWarmingRefresher +from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_RECORD_SCHEMA_VERSION, + CacheWarmingAttribution, + CacheWarmingPayload, + CacheWarmingRecord, + compress_payload, +) +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter +from litellm.types.router import TaggedPreRoutingStrategy + +from tests.test_litellm.router_strategy.complexity_router.cache_warming.test_store import FakeRedisCache + +DEFAULT_MODEL_LIST = [ + {"model_name": "fast-claude", "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "sk-t"}}, + {"model_name": "smart-claude", "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-t"}}, +] +UNIFORM_POOL = [ + {"model_name": "uniform", "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-t"}}, + {"model_name": "uniform", "litellm_params": {"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"}}, +] +PRICED_TIERS = {"SIMPLE": ["claude-haiku-4-5"], "COMPLEX": ["claude-sonnet-4-5"]} +PRICED_MODEL_LIST = [ + {"model_name": "claude-haiku-4-5", "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "sk-t"}}, + {"model_name": "claude-sonnet-4-5", "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-t"}}, +] + + +class ReplayRouter(Router): + """The one router double: a real litellm.Router, so every lookup the request path makes on the way to a + replay is production code (model groups, aliases, pricing via get_model_group_info, tags, deployment + info, blocked deployments). Only the outbound dispatch is captured instead of sent.""" + + def __init__( + self, + model_list: list | None = None, + redis: FakeRedisCache | None = None, + enable_tag_filtering: bool = False, + replay_delay: float = 0.0, + ) -> None: + super().__init__( + model_list=[dict(entry) for entry in (model_list if model_list is not None else DEFAULT_MODEL_LIST)], + enable_tag_filtering=enable_tag_filtering, + ) + self.cache.redis_cache = redis + self.completion_calls: list[dict] = [] + self.anthropic_calls: list[dict] = [] + self.replay_delay = replay_delay + self.failing_message_marker: str | None = None + self.max_concurrent = 0 + self._in_flight = 0 + # Router assigns aanthropic_messages as an instance attribute (router.py:1211), which shadows a + # subclass method, so the capture is installed after super().__init__ or the replay hits the network + self.aanthropic_messages = self._capture_anthropic + + async def _capture_anthropic(self, **kwargs: object) -> None: + self.anthropic_calls.append(kwargs) + + async def acompletion(self, **kwargs: object): # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the overloads + marker = self.failing_message_marker + if marker is not None and marker in json.dumps(kwargs.get("messages"), default=str): + raise RuntimeError("provider down") + self._in_flight += 1 + self.max_concurrent = max(self.max_concurrent, self._in_flight) + await asyncio.sleep(self.replay_delay) + self._in_flight -= 1 + self.completion_calls.append(kwargs) + + +class FakeLeaseLock: + """acquire/extend/release shaped like RedisDistributedLock.""" + + def __init__(self, acquisition: object = None, extend_ok: bool = True) -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import ( + LockAcquisition, + ) + + self.acquisition = acquisition if acquisition is not None else LockAcquisition.ACQUIRED + self.extend_ok = extend_ok + self.acquire_calls: list[tuple[str, float]] = [] + self.extend_calls: list[str] = [] + self.release_calls: list[str] = [] + + async def acquire(self, key: str, token: str, ttl_seconds: float) -> object: + self.acquire_calls.append((key, ttl_seconds)) + return self.acquisition + + async def extend(self, key: str, token: str, ttl_seconds: float) -> bool: + self.extend_calls.append(key) + return self.extend_ok + + async def release(self, key: str, token: str) -> None: + self.release_calls.append(key) + + +class FakeKeyDirectory: + """get_key_object-shaped resolver: raises for unknown keys, like token_not_found_in_db. With no states + every looked-up key exists as unlimited, which is the shape most cases want.""" + + def __init__(self, states: dict | None = None, raise_all: bool = False) -> None: + self.states = states + self.raise_all = raise_all + self.lookups: list[str] = [] + + async def resolve(self, hashed_token: str, prisma_client: object, user_api_key_cache: object): + self.lookups.append(hashed_token) + if self.raise_all: + raise Exception(f"Authentication Error, Invalid proxy server token passed: {hashed_token}") + if self.states is None: + return key_state(token=hashed_token) + if hashed_token not in self.states: + raise Exception(f"Authentication Error, Invalid proxy server token passed: {hashed_token}") + return self.states[hashed_token] + + +def key_state(token: str = "hash-1", **fields: object): + from litellm.proxy._types import UserAPIKeyAuth + + return UserAPIKeyAuth(token=token, **fields) + + +def team(team_id: str, **fields: object): + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + return LiteLLM_TeamTableCachedObj(team_id=team_id, **fields) + + +def real_limiter(window_size: int | None = None) -> tuple[object, DualCache]: + """The production v3 limiter, standalone over an in-memory DualCache, so its own counters are the + assertion surface. window_size is read from the environment at construction.""" + from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache + + previous = os.environ.get("LITELLM_RATE_LIMIT_WINDOW_SIZE") + if window_size is not None: + os.environ["LITELLM_RATE_LIMIT_WINDOW_SIZE"] = str(window_size) + try: + counters = DualCache() + return (_PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(counters)), counters) + finally: + if window_size is not None: + if previous is None: + del os.environ["LITELLM_RATE_LIMIT_WINDOW_SIZE"] + else: + os.environ["LITELLM_RATE_LIMIT_WINDOW_SIZE"] = previous + + +@contextmanager +def registered_callbacks(*callbacks: object): + """litellm.callbacks is what ProxyLogging.pre_call_hook walks, so a hook is only reachable through the + entry point once it is registered there.""" + previous = litellm.callbacks + litellm.callbacks = list(callbacks) + try: + yield + finally: + litellm.callbacks = previous + + +def warming_rig( + redis: FakeRedisCache | None = None, + enable_tag_filtering: bool = False, + replay_delay: float = 0.0, + tiers: dict | None = None, + model_list: list | None = None, + **cache_warming_overrides: object, +) -> tuple[ReplayRouter, FakeRedisCache | None]: + llm_router = ReplayRouter( + model_list=model_list, redis=redis, enable_tag_filtering=enable_tag_filtering, replay_delay=replay_delay + ) + strategy = ComplexityRouter( + model_name="smart-router", + litellm_router_instance=llm_router, + complexity_router_config={ + "tiers": tiers if tiers is not None else {"SIMPLE": ["fast-claude"], "COMPLEX": ["smart-claude"]}, + "cache_warming": {"enabled": True, **cache_warming_overrides}, + }, + ) + llm_router.complexity_routers = {"smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=strategy)]} + return llm_router, redis + + +def priced_rig(redis: FakeRedisCache, **overrides: object) -> ReplayRouter: + return warming_rig(redis=redis, tiers=PRICED_TIERS, model_list=PRICED_MODEL_LIST, **overrides)[0] + + +def seed_session( + redis: FakeRedisCache, + session_id: str = "sess-1", + caller_scope: str = "hash-1", + served_model: str = "fast-claude", + last_activity: float | None = None, + warmth: dict | None = None, + user_api_key: str | None = "hash-1", + call_surface: str = "chat_completions", + content: str = "summarize the deployment policy", + tools: tuple | None = None, + tool_choice: object = None, + team_id: str | None = None, + user_id: str | None = None, + org_id: str | None = None, +) -> str: + payload = CacheWarmingPayload( + model=served_model, + messages=({"role": "user", "content": content},), + system="You are a policy assistant" if call_surface == "anthropic_messages" else None, + tools=tools, + tool_choice=tool_choice, + call_surface=call_surface, + ) + blob, sha = compress_payload(payload) + record = CacheWarmingRecord( + schema_version=CACHE_WARMING_RECORD_SCHEMA_VERSION, + payload_compressed=blob, + payload_sha256=sha, + token_estimate=2048, + last_activity=last_activity if last_activity is not None else time.time(), + served_model=served_model, + session_id=session_id, + attribution=CacheWarmingAttribution( + user_api_key=user_api_key, + user_api_key_team_id=team_id, + user_api_key_user_id=user_id, + user_api_key_org_id=org_id, + ), + auto_router_model_name="smart-router", + ) + store = CacheWarmingStore(redis_cache=redis, auto_router_model_name="smart-router") + record_key = CacheWarmingStore.record_key("smart-router", caller_scope, session_id) + redis.hashes.setdefault(store.sessions_key(), {})[record_key] = json.dumps(record.model_dump()) + redis.zsets.setdefault(store.index_key(), {})[record_key] = time.time() + 3600 + for model_group, stamp in (warmth or {}).items(): + redis.data[CacheWarmingStore.warmth_key(record_key, model_group)] = json.dumps(stamp) + return record_key + + +def warmth_stamp(redis: FakeRedisCache, record_key: str, model_group: str) -> float | None: + raw = redis.data.get(CacheWarmingStore.warmth_key(record_key, model_group)) + return json.loads(raw) if raw is not None else None + + +def proxy_logging_with_hooks(limiter: object | None = None) -> ProxyLogging: + """ProxyLogging with the hooks warming consults registered, as the proxy registers them, so a gate that + reads one is live rather than silently absent. The v3 limiter must be present because warming denies the + whole tick when any enforcement dependency is missing.""" + from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler + from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache + + logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + logging_obj.proxy_hook_mapping["max_iterations_limiter"] = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + logging_obj.proxy_hook_mapping["parallel_request_limiter"] = limiter or _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + return logging_obj + + +def refresher( + keys: FakeKeyDirectory | None = None, + lock: FakeLeaseLock | None = None, + proxy_logging: ProxyLogging | None = None, + limiter: object | None = None, + **kwargs: object, +) -> CacheWarmingRefresher: + directory = keys if keys is not None else FakeKeyDirectory() + lease_lock = lock if lock is not None else FakeLeaseLock() + logging_obj = proxy_logging if proxy_logging is not None else proxy_logging_with_hooks(limiter) + return CacheWarmingRefresher( + key_state_resolver=directory.resolve, + lock_factory=lambda _redis: lease_lock, + proxy_logging_resolver=lambda: logging_obj, + **kwargs, + ) + + +_DB = object() + + +async def tick( + llm_router, + prisma=_DB, + active: CacheWarmingRefresher | None = None, + keys: FakeKeyDirectory | None = None, + user_api_key_cache: DualCache | None = None, +) -> None: + """Default: a reachable DB whose directory knows every attributed key as unlimited. The key cache is a + real DualCache because the authorization and budget entry points read through it.""" + await (active if active is not None else refresher(keys=keys)).run_tick( + llm_router=llm_router, + prisma_client=object() if prisma is _DB else prisma, + # a proxy without a database still has an in-memory key cache, so the two dependencies stay + # independently observable rather than one masking the other + user_api_key_cache=user_api_key_cache or DualCache(), + ) + + +def replayed_models(llm_router: ReplayRouter) -> list: + return [call["model"] for call in llm_router.completion_calls] + + +def affinity_check(model_group_affinity_config: dict | None = None, session_mode: bool = True): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck + + return DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=session_mode, + model_group_affinity_config=model_group_affinity_config or {}, + ) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index e734d8ec876..d0458452c4a 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -4159,3 +4159,182 @@ def test_every_routing_decision_field_is_classified(): f"unclassified={declared - classified}, stale={classified - declared}" ) assert not (PROMPT_QUOTING_ROUTING_DECISION_FIELDS & DERIVED_ROUTING_DECISION_FIELDS) +class TestCacheWarmingConfig: + def test_cache_warming_config_defaults_disabled(self): + config = ComplexityRouterConfig() + assert config.cache_warming.enabled is False + assert config.cache_warming.refresh_interval_seconds == 270 + assert config.cache_warming.session_ttl_seconds == 3600 + assert config.cache_warming.idle_timeout_seconds == 600 + assert config.cache_warming.max_sessions == 1000 + assert config.cache_warming.warm_models is None + + def test_cache_warming_coerces_from_nested_yaml_dict(self): + config = ComplexityRouterConfig( + cache_warming={"enabled": True, "refresh_interval_seconds": 120, "warm_models": ["sonnet", "opus"]} + ) + assert config.cache_warming.enabled is True + assert config.cache_warming.refresh_interval_seconds == 120 + assert config.cache_warming.warm_models == ("sonnet", "opus") + + def test_cache_warming_and_adaptive_mutually_exclusive_raises(self): + with pytest.raises(ValueError, match="cache_warming and adaptive"): + ComplexityRouterConfig( + tiers={"SIMPLE": ["gpt-4o-mini"]}, + adaptive=True, + cache_warming={"enabled": True}, + ) + + def test_cache_warming_disabled_with_adaptive_is_allowed(self): + config = ComplexityRouterConfig(tiers={"SIMPLE": ["gpt-4o-mini"]}, adaptive=True) + assert config.cache_warming.enabled is False + + def test_cache_warming_rejects_nonpositive_intervals(self): + with pytest.raises(ValidationError): + ComplexityRouterConfig(cache_warming={"enabled": True, "refresh_interval_seconds": 0}) + + +class TestWarmAwarePick: + _POOL = ["fast-claude", "smart-claude", "cold-model"] + + @staticmethod + def _router(mock_router_instance, redis, **config_overrides): + from types import SimpleNamespace + + mock_router_instance.cache = SimpleNamespace(redis_cache=redis) + return ComplexityRouter( + model_name="warm-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": list(TestWarmAwarePick._POOL)}, + "session_affinity": False, + "cache_warming": {"enabled": True}, + **config_overrides, + }, + ) + + @staticmethod + def _seed(redis, warmth, last_activity=None, served_model="fast-claude"): + import json + import time as time_module + + from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore + from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_RECORD_SCHEMA_VERSION, + CacheWarmingAttribution, + CacheWarmingPayload, + CacheWarmingRecord, + compress_payload, + ) + + payload = CacheWarmingPayload( + model=served_model, + messages=({"role": "user", "content": "hello"},), + call_surface="chat_completions", + ) + blob, sha = compress_payload(payload) + record = CacheWarmingRecord( + schema_version=CACHE_WARMING_RECORD_SCHEMA_VERSION, + payload_compressed=blob, + payload_sha256=sha, + token_estimate=2048, + last_activity=last_activity if last_activity is not None else time_module.time(), + served_model=served_model, + attribution=CacheWarmingAttribution(user_api_key="hash-w"), + auto_router_model_name="warm-router", + ) + store = CacheWarmingStore(redis_cache=redis, auto_router_model_name="warm-router") + key = store.record_key("warm-router", "hash-w", "warm-sess") + redis.hashes.setdefault(store.sessions_key(), {})[key] = json.dumps(record.model_dump()) + for model_group, stamp in warmth.items(): + redis.data[CacheWarmingStore.warmth_key(key, model_group)] = json.dumps(stamp) + + @staticmethod + def _kwargs(): + return {"metadata": {"session_id": "warm-sess", "user_api_key_hash": "hash-w"}} + + @staticmethod + def _fresh_redis(): + from tests.test_litellm.router_strategy.complexity_router.cache_warming.test_store import FakeRedisCache + + return FakeRedisCache() + + @pytest.mark.asyncio + async def test_pick_restricted_to_warmed_members_and_served_model(self, mock_router_instance): + import time as time_module + + redis = self._fresh_redis() + self._seed(redis, warmth={"smart-claude": time_module.time()}, served_model="fast-claude") + router = self._router(mock_router_instance, redis) + picks = { + await router._pick_model_for_tier(ComplexityTier.SIMPLE, None, None, self._kwargs()) for _ in range(20) + } + assert "cold-model" not in picks + assert picks <= {"fast-claude", "smart-claude"} + + @pytest.mark.asyncio + async def test_stale_warm_entries_fall_back(self, mock_router_instance): + import time as time_module + + redis = self._fresh_redis() + stale = time_module.time() - (270 + 60 + 5) + self._seed(redis, warmth={"smart-claude": stale}, last_activity=time_module.time() - 601) + router = self._router(mock_router_instance, redis) + assert await router._warm_aware_pick(self._POOL, self._kwargs()) is None + + @pytest.mark.asyncio + async def test_served_model_counts_as_warm_only_while_session_active(self, mock_router_instance): + import time as time_module + + redis = self._fresh_redis() + self._seed(redis, warmth={}, served_model="fast-claude") + router = self._router(mock_router_instance, redis) + picks = { + await router._pick_model_for_tier(ComplexityTier.SIMPLE, None, None, self._kwargs()) for _ in range(20) + } + assert picks == {"fast-claude"} + + self._seed(redis, warmth={}, served_model="fast-claude", last_activity=time_module.time() - 601) + assert await router._warm_aware_pick(self._POOL, self._kwargs()) is None + + @pytest.mark.asyncio + async def test_falls_back_without_session_or_record_or_redis(self, mock_router_instance): + redis = self._fresh_redis() + router = self._router(mock_router_instance, redis) + assert await router._warm_aware_pick(self._POOL, {}) is None + assert await router._warm_aware_pick(self._POOL, self._kwargs()) is None + no_redis_router = self._router(MagicMock(), None) + assert await no_redis_router._warm_aware_pick(self._POOL, self._kwargs()) is None + pick = await router._pick_model_for_tier(ComplexityTier.SIMPLE, None, None, {}) + assert pick in self._POOL + + @pytest.mark.asyncio + async def test_disabled_or_single_member_pool_returns_none(self, mock_router_instance): + import time as time_module + + redis = self._fresh_redis() + self._seed(redis, warmth={"smart-claude": time_module.time()}) + router = self._router(mock_router_instance, redis) + assert await router._warm_aware_pick(["fast-claude"], self._kwargs()) is None + disabled = self._router(MagicMock(), redis, cache_warming={"enabled": False}) + assert await disabled._warm_aware_pick(self._POOL, self._kwargs()) is None + + @pytest.mark.asyncio + async def test_plugin_narrowed_candidates_get_warm_pick(self, mock_router_instance): + import time as time_module + + class ExcludeSmartClaude: + async def run(self, context): + context.candidate_models = [m for m in context.candidate_models if m != "smart-claude"] + return context + + redis = self._fresh_redis() + self._seed(redis, warmth={"smart-claude": time_module.time()}, served_model="fast-claude") + router = self._router(mock_router_instance, redis, plugins=[ExcludeSmartClaude()]) + picks = { + await router._pick_model_for_tier( + ComplexityTier.SIMPLE, None, None, self._kwargs() + ) + for _ in range(20) + } + assert picks == {"fast-claude"}