scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

feat(complexity_router): multi-model prompt-cache warming for the auto-router

Browse source 
Provider prompt caches are per-model, so every mid-session tier switch the
complexity auto-router makes lands on a cold cache and pays the full cache write
again. Opt-in cache_warming captures each session's latest payload at the routing
decision and a leader-elected background refresher replays it with max_tokens=1
against every cacheable tier model just under the provider cache TTL, so the
switch is a pure cache read.

A replay is a request, so it is admitted through the request path's own entry
points rather than beside them. For each replay the refresher assembles a request
body, reserves budget through the same wrapper auth calls right after
common_checks, stamps identity with the proxy's own stamper, applies every
key-level, team-level and project-level control, applies the key and team scoped
dynamic logging settings, runs ProxyLogging.pre_call_hook, applies the
fully-blocked-model check, and hands the dict that hook returns to
Router.acompletion or Router.aanthropic_messages. post_call_failure_hook runs on
every rejection and every dispatch failure, so the parallel request slot, the
reserved TPM tokens and the budget reservation all come back. Warming therefore
inherits both halves of every contract it touches (the limiter's descriptors
across every scope with its own configured window, its RPM and max-parallel
check, its upfront reservation and the stash its success callback reconciles
from; the key, team, user, end-user, organization and tag budget counters; every
configured guardrail and pipeline, including the ones defined on the deployment)
instead of reimplementing them. That deletes nine functions and the admission
block they served.

Warming writes no spend logs of its own, so the replay rows are the only record
of warming cost that will exist. They carry the customer's own tags and
spend_logs_metadata with the litellm_cache_warming tag alongside rather than
instead, and they fan out to the key and team scoped loggers, so warming is
included in per-tag chargeback and filterable out of it. Two Request-free blocks
of add_litellm_data_to_request are extracted verbatim as
LiteLLMProxyRequestSetup.add_key_team_project_metadata and
apply_dynamic_logging_settings so both callers share them; the move is
statement-for-statement identical, with no behavior change on the request path.
Ordering there is load-bearing: add_key_level_controls resets data["cache"] and
refills it from key metadata, so it runs after the body is built and a key's own
cache controls override warming's response-cache bypass exactly as they override
a caller's.

Blocked and expired keys are still checked locally because common_checks
dereferences the FastAPI Request; extracting a Request-free core so its other
gates bind on a replay too is a follow-up. Sessions on a key that declares
max_iterations are skipped entirely, because that limiter counts every request on
a session_id and cannot be consulted without incrementing it.

Metadata precedence (litellm_metadata before metadata, stringified) had three
implementations; core_helpers.iter_request_metadata_dicts and
get_request_metadata_field are now the single owner and DeploymentAffinityCheck
deletes its four private copies to delegate to them.

Resolves LIT-4865
This commit is contained in:
Tin Chi Lo 2026-07-29 16:29:46 -07:00
parent fb79a4ee3b
commit 3d2e18f41d
26 changed files with 3090 additions and 175 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/constants.py
View file

@ -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))

4
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -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

24
litellm/litellm_core_utils/core_helpers.py
View file

@ -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

18
litellm/proxy/auth/auth_checks.py
View file

@ -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)

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

@ -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

212
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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(

25
litellm/proxy/proxy_server.py
View file

@ -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.

42
litellm/router_strategy/complexity_router/README.md
View file

@ -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

33
litellm/router_strategy/complexity_router/cache_warming/__init__.py Normal file
View file

@ -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",
]

224
litellm/router_strategy/complexity_router/cache_warming/capture.py Normal file
View file

@ -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,
)

22
litellm/router_strategy/complexity_router/cache_warming/eligibility.py Normal file
View file

@ -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)

809
litellm/router_strategy/complexity_router/cache_warming/refresher.py Normal file
View file

@ -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}

222
litellm/router_strategy/complexity_router/cache_warming/store.py Normal file
View file

@ -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)

91
litellm/router_strategy/complexity_router/cache_warming/types.py Normal file
View file

@ -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)

89
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -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(

62
litellm/router_strategy/complexity_router/config.py
View file

@ -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()

61
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
View file

@ -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.

72
tests/test_litellm/proxy/proxy_server/test_background_health.py
View file

@ -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
# ---------------------------------------------------------------------------

0
tests/test_litellm/router_strategy/complexity_router/__init__.py Normal file
View file

0
tests/test_litellm/router_strategy/complexity_router/cache_warming/__init__.py Normal file
View file

176
tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py Normal file
View file

@ -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"

337
tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py Normal file
View file

@ -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

182
tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py Normal file
View file

@ -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

46
tests/test_litellm/router_strategy/complexity_router/cache_warming/test_types.py Normal file
View file

@ -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)}"
)

330
tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py Normal file
View file

@ -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 {},
)

179
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -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"}