diff --git a/litellm/constants.py b/litellm/constants.py index 30d3bb1f26e7..3b91f23fe397 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1322,6 +1322,7 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata" OLD_LITELLM_METADATA_FIELD: Final = "metadata" RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name" +SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 222005670122..fbf28e223c15 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -6,7 +6,7 @@ import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger -from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, @@ -425,6 +425,7 @@ def get_logging_caching_headers(request_data: dict) -> dict | None: "_guardrail_pipelines", "_pipeline_managed_guardrails", PRE_CALL_EXECUTED_GUARDRAILS_KEY, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 83ae59ef0504..c48fee96646d 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -18,6 +18,7 @@ INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -226,6 +227,7 @@ def parse_cache_control(cache_control): "applied_policies", "policy_sources", "routing_decision", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, "standard_logging_object", "proxy_server_request", diff --git a/litellm/router.py b/litellm/router.py index 1a76eb5d59bf..feaf69a44ae1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -46,6 +46,7 @@ DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function @@ -135,6 +136,7 @@ from litellm.router_utils.health_state_cache import DeploymentHealthCache from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, + warn_on_unknown_model_group_affinity_flags, ) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( build_io_token_rate_limit_headers, @@ -603,6 +605,10 @@ def __init__( # ``litellm.proxy.auth.auth_checks._is_model_cost_zero``. self._zero_cost_cache: dict[str, bool] = {} + self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds + self.model_group_affinity_config = model_group_affinity_config + warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) + if model_list is not None: # set_model_list will build indices automatically self.set_model_list(model_list) @@ -744,7 +750,6 @@ def __init__( litellm.failure_callback = [self.deployment_callback_on_failure] self.routing_strategy_args = routing_strategy_args self.provider_budget_config = provider_budget_config - self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds self.router_budget_logger: RouterBudgetLimiting | None = None if RouterBudgetLimiting.should_init_router_budget_limiter( model_list=model_list, provider_budget_config=self.provider_budget_config @@ -766,7 +771,6 @@ def __init__( ) self.model_group_retry_policy: dict[str, RetryPolicy] | None = model_group_retry_policy - self.model_group_affinity_config: dict[str, list[str]] | None = model_group_affinity_config self.allowed_fails_policy: AllowedFailsPolicy | None = None if allowed_fails_policy is not None: @@ -789,21 +793,8 @@ def __init__( # If model_group_affinity_config is set but no global affinity checks were # enabled, we still need the DeploymentAffinityCheck callback (with global # flags all False) so per-group config can activate affinity per model group. - if self.model_group_affinity_config and not any( - isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or []) - ): - if self.optional_callbacks is None: - self.optional_callbacks = [] - affinity_callback: Final = DeploymentAffinityCheck( - cache=self.cache, - ttl_seconds=self.deployment_affinity_ttl_seconds, - enable_user_key_affinity=False, - enable_responses_api_affinity=False, - enable_session_id_affinity=False, - model_group_affinity_config=self.model_group_affinity_config, - ) - self.optional_callbacks.append(affinity_callback) - litellm.logging_callback_manager.add_litellm_callback(affinity_callback) + if self.model_group_affinity_config: + self._ensure_deployment_affinity_callback() if self.alerting_config is not None: self._initialize_alerting() @@ -1662,6 +1653,28 @@ def _move_before_deployment_affinity( _move_before_deployment_affinity(self.optional_callbacks, ec_callback) _move_before_deployment_affinity(litellm.callbacks, ec_callback) + def _ensure_deployment_affinity_callback(self) -> None: + """Register the DeploymentAffinityCheck callback (global flags all False) if absent. + + Needed when nothing enabled a global affinity flag but affinity can still + activate per request: per-group `model_group_affinity_config` entries, or the + session-affinity marker a complexity router stamps at pre-routing time. + """ + if any(isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])): + return + if self.optional_callbacks is None: + self.optional_callbacks = [] + affinity_callback: Final = DeploymentAffinityCheck( + cache=self.cache, + ttl_seconds=self.deployment_affinity_ttl_seconds, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=False, + model_group_affinity_config=self.model_group_affinity_config, + ) + self.optional_callbacks.append(affinity_callback) + litellm.logging_callback_manager.add_litellm_callback(affinity_callback) + def add_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None): if optional_pre_call_checks is None: return @@ -7683,6 +7696,8 @@ def init_complexity_router_deployment(self, deployment: Deployment): strategy=complexity_router, strategy_label="Complexity-router", ) + if complexity_router._uses_deployment_pin: + self._ensure_deployment_affinity_callback() def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """True when this deployment opts in via the `auto_router/adaptive_router` model prefix.""" @@ -11190,6 +11205,9 @@ async def async_pre_routing_hook( router_strategy: Final = self._select_pre_routing_strategy(model=model, request_kwargs=request_kwargs) if router_strategy is None: self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) + self._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None + ) return None pre_routing_hook_response: Final = await router_strategy.async_pre_routing_hook( @@ -11203,6 +11221,11 @@ async def async_pre_routing_hook( request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), ) + self._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, + key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + value=(pre_routing_hook_response.session_affinity_ttl_seconds if pre_routing_hook_response else None), + ) # `model` (the alias, e.g. "smart-router") is never the deployment actually # called - apply the alias's own litellm_params (besides `model` itself, @@ -11234,21 +11257,40 @@ def _record_routing_decision( to the deployment that actually served the request. Every attempt therefore writes or clears, never just writes. """ - if routing_decision is None: + Router._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, + key="routing_decision", + value=( + None + if routing_decision is None + else Router._redact_prompt_text_if_needed( + request_kwargs=request_kwargs, routing_decision=routing_decision + ) + ), + ) + + @staticmethod + def _stamp_or_clear_metadata_key(request_kwargs: dict, key: str, value: object | None) -> None: + """Write a proxy-internal metadata key for THIS routing attempt, or clear it. + + Fallbacks and retries re-enter the pre-routing hook with the same + `request_kwargs`, so every attempt must write or clear, never just write; + a value left behind by an earlier attempt would be attributed to this one. + `get_or_create_metadata_bucket` is the single owner of "which dict holds + proxy-internal metadata": it picks `litellm_metadata` when present (so the + value never lands in the `metadata` dict that routes like /v1/messages + forward to the provider) and replaces a non-dict value rather than silently + skipping the write. Clearing pops from BOTH buckets so a request whose + bucket resolution changed between attempts cannot resurrect a stale value. + """ + if value is None: for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")): if isinstance(bucket, dict): - bucket.pop("routing_decision", None) + bucket.pop(key, None) return - # `get_or_create_metadata_bucket` is the single owner of "which dict holds - # proxy-internal metadata": it picks `litellm_metadata` when present (so the - # decision never lands in the `metadata` dict that routes like /v1/messages - # forward to the provider) and replaces a non-dict value rather than silently - # skipping the write. _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket["routing_decision"] = Router._redact_prompt_text_if_needed( - request_kwargs=request_kwargs, routing_decision=routing_decision - ) + metadata_bucket[key] = value @staticmethod def _redact_prompt_text_if_needed( diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index f6ced0bb9d2c..32d252f3f680 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1651,6 +1651,28 @@ def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) caller_scope: Final = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" + @property + def _uses_tier_pin(self) -> bool: + return bool(self.config.session_affinity and not self.config.plugins) + + @property + def _uses_deployment_pin(self) -> bool: + """session_affinity implies the deployment pin: a session frozen onto one model + group but load-balanced across its deployments would still go cache-cold, which + is the exact failure both flags exist to prevent.""" + return bool((self.config.deployment_affinity or self.config.session_affinity) and not self.config.plugins) + + def _with_session_deployment_affinity( + self, response: PreRoutingHookResponse | None + ) -> PreRoutingHookResponse | None: + if response is None or not self._uses_deployment_pin: + return response + return response.model_copy( + update={ # mutable-ok: model_copy types update as a plain dict + "session_affinity_ttl_seconds": self.config.session_affinity_ttl_seconds + } + ) + async def async_pre_routing_hook( self, model: str, @@ -1685,7 +1707,7 @@ async def async_pre_routing_hook( resolved_messages: Final = self._resolve_messages(messages, request_kwargs) conversation_continuing: Final = _conversation_is_continuing(resolved_messages) - use_session_affinity: Final = self.config.session_affinity and not self.config.plugins + use_session_affinity: Final = self._uses_tier_pin session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) 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 @@ -1724,17 +1746,19 @@ async def async_pre_routing_hook( "ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model ) has_original_messages: Final = messages is not None and len(messages) > 0 - return PreRoutingHookResponse( - model=routed_model, - messages=messages if has_original_messages else None, - routing_decision=self._build_routing_decision( - routed_model=routed_model, - cause=cause, - tier=self._tier_for_model(routed_model), - escalation_keyword=pin_escalation_keyword, - escalated=escalated, - conversation_continuing=conversation_continuing, - ), + return self._with_session_deployment_affinity( + PreRoutingHookResponse( + model=routed_model, + messages=messages if has_original_messages else None, + routing_decision=self._build_routing_decision( + routed_model=routed_model, + cause=cause, + tier=self._tier_for_model(routed_model), + escalation_keyword=pin_escalation_keyword, + escalated=escalated, + conversation_continuing=conversation_continuing, + ), + ) ) response: Final = await self._classify_and_route( @@ -1752,7 +1776,7 @@ async def async_pre_routing_hook( value=response.model, ttl=self.config.session_affinity_ttl_seconds, ) - return response + return self._with_session_deployment_affinity(response) async def _classify_and_route( self, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 69609a973b02..0999af66fd8d 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -508,13 +508,39 @@ class ComplexityRouterConfig(BaseModel): "session's first turn and reuse it for every later turn, skipping re-classification. " "Off by default so every turn is classified on its own merits and routed to the cheapest " "adequate tier. Set True to keep a multi-turn session on one model, which preserves " - "provider prompt caches and avoids cross-model conversation-history errors." + "provider prompt caches and avoids cross-model conversation-history errors. Always " + "implies the deployment pin regardless of deployment_affinity: the session sticks to " + "one deployment of the pinned model, since freezing the model while re-shuffling its " + "deployments would still go cache-cold." + ), + ) + deployment_affinity: bool = Field( + default=True, + description=( + "When True and a session_id is resolvable on the request, pin the deployment chosen " + "inside each routed model group and reuse it whenever the session returns to that " + "group, without pinning which group the session routes to. Independent of " + "session_affinity, which pins the model group instead (and always carries this " + "deployment pin with it): with session_affinity off, " + "every turn is still classified on its own merits while a session that escalates to a " + "stronger tier and comes back still lands on the deployment it used before, which is " + "what keeps a provider prompt cache warm. Pins are held per model group, so switching " + "tiers does not disturb the pin left behind in the previous group. On by default " + "because re-shuffling a conversation across deployments of the same model discards " + "that cache for no benefit; set False to keep every turn load-balanced across the " + "group, which is what a deployment set with tight per-deployment rate limits wants. " + "Inert when no session_id is resolvable, since there is nothing to key a pin on, and " + "suppressed when plugins are configured, for the same reason session_affinity is." ), ) session_affinity_ttl_seconds: int = Field( default=3600, gt=0, - description="TTL for the session affinity pin; refreshed on every cache hit", + description=( + "TTL for the session affinity pin; refreshed on every cache hit. Bounds both the " + "session_affinity model pin and the deployment_affinity deployment pin, so it measures " + "idle time for the session's routing decisions rather than total session length" + ), ) plugins: list[RoutingPlugin] | None = Field( diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index b4408d17ffb2..7fb90ab89de0 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -13,12 +13,15 @@ """ import hashlib +import json +from collections.abc import Mapping, Sequence from typing import Any, Final, cast from typing_extensions import TypedDict from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues @@ -29,6 +32,47 @@ class DeploymentAffinityCacheValue(TypedDict): model_id: str +VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( + { + "deployment_affinity", + "responses_api_deployment_check", + "session_affinity", + "encrypted_content_affinity", + } +) + + +def warn_on_unknown_model_group_affinity_flags(model_group_affinity_config: Mapping[str, Sequence[str]] | None) -> None: + """`model_group_affinity_config` is one Router-level config consumed by two callbacks: + DeploymentAffinityCheck acts on three of the flags and EncryptedContentAffinityCheck + on the fourth, so typo detection lives here at the schema, not inside either consumer. + """ + if model_group_affinity_config is None: + return + for group, flags in model_group_affinity_config.items(): + unknown = set(flags) - VALID_MODEL_GROUP_AFFINITY_FLAGS + if unknown: + verbose_router_logger.warning( + "model_group_affinity_config: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s", + unknown, + group, + VALID_MODEL_GROUP_AFFINITY_FLAGS, + ) + + +_CLAIM_PIN_SCRIPT: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if current == ARGV[1] then + redis.call('EXPIRE', KEYS[1], ARGV[2]) +end +return current +""" + + class DeploymentAffinityCheck(CustomLogger): """ Router deployment affinity callback. @@ -38,14 +82,6 @@ class DeploymentAffinityCheck(CustomLogger): """ CACHE_KEY_PREFIX = "deployment_affinity:v1" - VALID_FLAGS = frozenset( - { - "deployment_affinity", - "responses_api_deployment_check", - "session_affinity", - "encrypted_content_affinity", - } - ) def __init__( self, @@ -63,15 +99,6 @@ def __init__( self.enable_responses_api_affinity = enable_responses_api_affinity self.enable_session_id_affinity = enable_session_id_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} - for group, flags in self.model_group_affinity_config.items(): - unknown = set(flags) - self.VALID_FLAGS - if unknown: - verbose_router_logger.warning( - "DeploymentAffinityCheck: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s", - unknown, - group, - self.VALID_FLAGS, - ) def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -218,8 +245,13 @@ def get_affinity_cache_key(cls, model_group: str, user_key: str) -> str: return f"{cls.CACHE_KEY_PREFIX}:{model_group}:{hashed_user_key}" @classmethod - 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}" + def get_session_affinity_cache_key(cls, model_group: str, session_id: str, user_key: str | None) -> str: + """Session pins are scoped by the caller's hashed API key so two callers reusing + the same client-supplied session_id cannot read or steer each other's pin. + `"unscoped"` covers direct Router usage with no authenticated caller, matching + the complexity router's own session pin key.""" + hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped" + return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" @staticmethod def _get_user_key_from_metadata_dict(metadata: dict) -> str | None: @@ -278,6 +310,97 @@ def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: return session_id return None + @staticmethod + def _get_marker_session_affinity_ttl(request_kwargs: dict) -> int | None: + """TTL from the session-affinity marker the Router stamps at pre-routing time + when an auto-router routed this request with session_affinity enabled. + Marker presence enables session pinning for this request only; anything that + is not a positive int is treated as absent.""" + for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): + ttl = metadata.get(SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY) + if isinstance(ttl, int) and not isinstance(ttl, bool) and ttl > 0: + return ttl + return None + + @staticmethod + def _pinned_model_id(stored: object) -> str | None: + """Deployment id held by a stored pin, for both the dict shape this writes and the + bare string older writers left behind. None when the value is neither.""" + if isinstance(stored, dict): + model_id: Final = stored.get("model_id") + return str(model_id) if model_id is not None else None + if isinstance(stored, str): + return stored + return None + + def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: + """The one owner of authoritative local pin writes: a plain set keeps a live + key's original expiry (`allow_ttl_override`), so the entry is replaced to make + the TTL real. Every local pin write goes through here so the redis-winner sync + and the pod-local claim can never disagree about expiry again.""" + self.cache.in_memory_cache.delete_cache(cache_key) + self.cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds) + + async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: + """First-writer-wins pin write: store `pin_value` only when the key is absent and + return the deployment id the key holds afterwards, so a caller learns whether it won + by comparing against its own id, and None when the stored value is one no reader can + interpret. Concurrent claimers converge on the + first write instead of the last. Re-claiming with the stored value refreshes its + TTL, the same keepalive the complexity router's model pin documents: an active + session must not lose its pin mid-conversation just because it outlives the + original write, so `session_affinity_ttl_seconds` bounds idle time, not total + session length. On Redis one Lua script does the get-or-set-or-refresh + atomically (same registration seam the rate limiters use) and the in-memory + tier is synchronized to the winner; without Redis, and whenever Redis is + unreachable, the pod-local check-and-set below stands in and is atomic because it + runs synchronously on the event loop. Degrading to a pod-local claim rather than + propagating the fault is what keeps same-pod stickiness through a Redis blip: the + caller only logs this result, so an escaping error would leave the session with no + pin at all and reshuffle every turn for the outage, which is worse than losing + cross-pod agreement. The redis tier is + resolved per call because the proxy attaches it after Router construction + (`Router._update_redis_cache`); the compiled script is cached per event loop + underneath the registration seam. + """ + redis_cache: Final = self.cache.redis_cache + if redis_cache is not None: + try: + claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) + raw: Final = await claim_script(keys=(cache_key,), args=(json.dumps(pin_value), int(ttl_seconds))) + decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw + if not isinstance(decoded, str): + return pin_value["model_id"] + try: + winner: object = json.loads(decoded) + except json.JSONDecodeError: + winner = decoded + self._set_local_pin(cache_key=cache_key, value=winner, ttl_seconds=ttl_seconds) + return self._pinned_model_id(winner) + except Exception as e: # noqa: BLE001 # any Redis/Lua failure degrades to the pod-local claim, never unpins + verbose_router_logger.debug( + "DeploymentAffinityCheck: redis pin claim failed, falling back to pod-local claim. error=%s", e + ) + + return self._claim_pin_in_memory(cache_key=cache_key, pin_value=pin_value, ttl_seconds=ttl_seconds) + + def _claim_pin_in_memory( + self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int + ) -> str | None: + """Pod-local half of the claim, used when no Redis tier is attached and as the + fallback when the Redis claim fails. Mirrors the Lua script exactly, including + the keepalive: re-claiming with the stored value slides the idle window through + `_set_local_pin`. Both branches stay synchronous, hence atomic on the event + loop.""" + existing: Final = self.cache.in_memory_cache.get_cache(cache_key) + if existing is not None: + existing_model_id: Final = self._pinned_model_id(existing) + if existing_model_id == pin_value["model_id"]: + self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) + return existing_model_id + self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) + return pin_value["model_id"] + @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: for deployment in healthy_deployments: @@ -334,12 +457,21 @@ async def async_filter_deployments( if stable_model_map_key is None: return typed_healthy_deployments + session_affinity_active: Final = ( + enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None + ) + user_key: Final = ( + self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) + if (session_affinity_active or enable_user_key) + else None + ) + # 2) Session-id -> deployment affinity - if enable_session_id: + if session_affinity_active: session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs=request_kwargs) if session_id is not None: session_cache_key: Final = self.get_session_affinity_cache_key( - model_group=stable_model_map_key, session_id=session_id + model_group=stable_model_map_key, session_id=session_id, user_key=user_key ) session_cache_result: Final = await self.cache.async_get_cache(key=session_cache_key) @@ -371,7 +503,6 @@ async def async_filter_deployments( if not enable_user_key: return typed_healthy_deployments - user_key: Final = self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if user_key is None: return typed_healthy_deployments @@ -438,18 +569,22 @@ async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type enable_session_id, ) = self._get_effective_flags(deployment_model_name) - if not enable_user_key and not enable_session_id: - return None + marker_session_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs=kwargs) + session_affinity_active: Final = enable_session_id or marker_session_ttl is not None - user_key = None - if enable_user_key: - user_key = self._get_user_key_from_request_kwargs(request_kwargs=kwargs) + if not enable_user_key and not session_affinity_active: + return None - session_id = None - if enable_session_id: - session_id = self._get_session_id_from_request_kwargs(request_kwargs=kwargs) + user_key: Final = ( + self._get_user_key_from_request_kwargs(request_kwargs=kwargs) + if (enable_user_key or session_affinity_active) + else None + ) + session_id: Final = ( + self._get_session_id_from_request_kwargs(request_kwargs=kwargs) if session_affinity_active else None + ) - if user_key is None and session_id is None: + if not ((enable_user_key and user_key is not None) or session_id is not None): return None model_info = kwargs.get("model_info") @@ -473,22 +608,31 @@ async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type verbose_router_logger.warning("DeploymentAffinityCheck: model_id missing; skipping affinity cache update.") return None - if user_key is not None: + pin_value: Final = DeploymentAffinityCacheValue(model_id=str(model_id)) + + if enable_user_key and user_key is not None: try: cache_key: Final = self.get_affinity_cache_key(model_group=deployment_model_name, user_key=user_key) - await self.cache.async_set_cache( - cache_key, - DeploymentAffinityCacheValue(model_id=str(model_id)), - ttl=self.ttl_seconds, - ) - - verbose_router_logger.debug( - "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s", - deployment_model_name, - model_id, - self.ttl_seconds, - self._shorten_for_logs(user_key), + claimed_user_pin: Final = await self._claim_pin( + cache_key=cache_key, + pin_value=pin_value, + ttl_seconds=self.ttl_seconds, ) + if claimed_user_pin == pin_value["model_id"]: + verbose_router_logger.debug( + "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s", + deployment_model_name, + model_id, + self.ttl_seconds, + self._shorten_for_logs(user_key), + ) + else: + verbose_router_logger.debug( + "DeploymentAffinityCheck: affinity pin already claimed model_map_key=%s existing=%s ours=%s", + deployment_model_name, + claimed_user_pin, + model_id, + ) except Exception as e: # Non-blocking: affinity is a best-effort optimization. verbose_router_logger.debug( @@ -500,21 +644,31 @@ async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type # Also persist Session-ID affinity if enabled and session-id is provided if session_id is not None: try: + session_affinity_ttl: Final = marker_session_ttl if marker_session_ttl is not None else self.ttl_seconds session_cache_key: Final = self.get_session_affinity_cache_key( - model_group=deployment_model_name, session_id=session_id + model_group=deployment_model_name, session_id=session_id, user_key=user_key ) - await self.cache.async_set_cache( - session_cache_key, - DeploymentAffinityCacheValue(model_id=str(model_id)), - ttl=self.ttl_seconds, - ) - verbose_router_logger.debug( - "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s", - deployment_model_name, - model_id, - self.ttl_seconds, - session_id, + claimed_session_pin: Final = await self._claim_pin( + cache_key=session_cache_key, + pin_value=pin_value, + ttl_seconds=session_affinity_ttl, ) + if claimed_session_pin == pin_value["model_id"]: + verbose_router_logger.debug( + "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s", + deployment_model_name, + model_id, + session_affinity_ttl, + session_id, + ) + else: + verbose_router_logger.debug( + "DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s", + deployment_model_name, + claimed_session_pin, + model_id, + session_id, + ) except Exception as e: verbose_router_logger.debug( "DeploymentAffinityCheck: failed to set session affinity cache. model_map_key=%s error=%s", diff --git a/litellm/types/router.py b/litellm/types/router.py index 4280da08cbbf..e166d8447357 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -816,6 +816,7 @@ class PreRoutingHookResponse(BaseModel): model: str messages: list[dict[str, Any]] | None routing_decision: StandardLoggingRoutingDecision | None = None + session_affinity_ttl_seconds: int | None = None _PreRoutingStrategyT_co = TypeVar("_PreRoutingStrategyT_co", covariant=True) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 6d6fd2e55079..22f9e6bb67a4 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -676,6 +676,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "applied_policies": ["spoofed-policy"], "policy_sources": {"spoofed-policy": "request"}, "routing_decision": {"cause": "forged", "routed_model": "spoofed"}, + "_session_deployment_affinity_ttl": 999999, "internal_call_origin": "autorouter_classifier", "_guardrail_pipelines": [{"name": "spoofed"}], "_pipeline_managed_guardrails": ["evaded"], @@ -719,6 +720,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "applied_policies", "policy_sources", "routing_decision", + "_session_deployment_affinity_ttl", "internal_call_origin", "_guardrail_pipelines", "_pipeline_managed_guardrails", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 356556f35631..94b6b68855ba 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3414,6 +3414,103 @@ def session_affinity_config(self, basic_config) -> Dict: def _request_kwargs(session_id: str) -> Dict: return {"metadata": {"session_id": session_id}} + @pytest.mark.asyncio + async def test_hook_response_carries_session_affinity_ttl_on_classify_and_pin_paths( + self, mock_router_instance, session_affinity_config + ): + """The hook response's session_affinity_ttl_seconds is what the Router stamps as + the deployment-affinity marker, so both the classify path (turn 1) and the + session-pin path (turn 2) must carry the configured TTL.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**session_affinity_config, "session_affinity_ttl_seconds": 321}, + ) + request_kwargs = self._request_kwargs("marker-session") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.session_affinity_ttl_seconds == 321 + assert second.session_affinity_ttl_seconds == 321 + + @pytest.mark.parametrize( + "session_affinity,deployment_affinity,plugins,tier_pinned,deployment_pinned", + [ + (False, False, False, False, False), + (False, True, False, False, True), + (True, False, False, True, True), + (True, True, False, True, True), + (False, True, True, False, False), + (True, True, True, False, False), + ], + ) + @pytest.mark.asyncio + async def test_tier_pin_and_deployment_pin_are_independently_gated( + self, + mock_router_instance, + basic_config, + session_affinity, + deployment_affinity, + plugins, + tier_pinned, + deployment_pinned, + ): + """deployment_affinity pins the deployment inside each routed group without pinning which + group the session routes to, so with session_affinity off the tier must still reclassify + on every turn while the marker the Router stamps is still emitted. Turn 1 classifies + REASONING and turn 2 SIMPLE, so a reclassified turn 2 moves model while a tier-pinned one + does not. plugins suppress both pins, since a stale pin would bypass the plugin pipeline.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_affinity": session_affinity, + "deployment_affinity": deployment_affinity, + **({"plugins": [_DummyPlugin()]} if plugins else {}), + }, + ) + request_kwargs = self._request_kwargs("matrix-session") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" + assert second.model == ("o1-preview" if tier_pinned else "gpt-4o-mini") + assert (first.session_affinity_ttl_seconds is not None) is deployment_pinned + assert (second.session_affinity_ttl_seconds is not None) is deployment_pinned + + @pytest.mark.asyncio + async def test_hook_response_has_no_session_affinity_ttl_when_disabled_or_plugins( + self, mock_router_instance, basic_config, session_affinity_config + ): + mock_router_instance.cache = DualCache() + disabled_router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "deployment_affinity": False}, + ) + plugin_router = ComplexityRouter( + model_name="test-router-plugins", + litellm_router_instance=mock_router_instance, + complexity_router_config={**session_affinity_config, "plugins": [_DummyPlugin()]}, + ) + disabled = await disabled_router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-off"), messages=self.SIMPLE_MESSAGE + ) + with_plugins = await plugin_router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-plugins"), messages=self.SIMPLE_MESSAGE + ) + assert disabled.session_affinity_ttl_seconds is None + assert with_plugins.session_affinity_ttl_seconds is None + @pytest.mark.asyncio async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): """Regression: session_affinity defaults to False, so a shared session_id must NOT diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index e29adda33281..428eb0ceafdd 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -465,8 +465,7 @@ async def test_async_pre_call_hook_uses_model_map_key_scope(): Deployment affinity caching uses (user_api_key_hash, model_map_key) -> model_id. """ - cache = AsyncMock() - cache.async_set_cache = AsyncMock() + cache = DualCache() callback = DeploymentAffinityCheck( cache=cache, @@ -489,11 +488,7 @@ async def test_async_pre_call_hook_uses_model_map_key_scope(): model_group="claude-sonnet-4-5@20250929", user_key="user-key-abc", ) - cache.async_set_cache.assert_called_once_with( - expected_cache_key, - {"model_id": "model-id-123"}, - ttl=123, - ) + assert await cache.async_get_cache(key=expected_cache_key) == {"model_id": "model-id-123"} @pytest.mark.asyncio diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index 0bcb0247aad8..4053e6d118ba 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,6 +1,6 @@ import os import sys -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -10,6 +10,7 @@ import litellm from litellm.caching.dual_cache import DualCache +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -163,7 +164,7 @@ async def test_async_session_id_affinity_priority_over_user_key(): await callback.cache.async_set_cache( DeploymentAffinityCheck.get_session_affinity_cache_key( - "model_group", "session1" + "model_group", "session1", user_key="user1" ), {"model_id": "deployment-2"}, ) @@ -180,3 +181,439 @@ async def test_async_session_id_affinity_priority_over_user_key(): assert len(filtered) == 1 assert filtered[0]["model_info"]["id"] == "deployment-2" + + +MOCK_RESPONSES_API_RESPONSE = { + "id": "resp_mock-resp-456", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "azure/computer-use-preview", + "output": [], + "usage": { + "input_tokens": 5, + "output_tokens": 10, + "total_tokens": 15, + "output_tokens_details": {"reasoning_tokens": 0}, + }, +} + + +def _smart_router(session_affinity=True, ttl_seconds=777, deployment_affinity=True): + return litellm.Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "target-group", + "complexity_router_config": { + "session_affinity": session_affinity, + "deployment_affinity": deployment_affinity, + "session_affinity_ttl_seconds": ttl_seconds, + "tiers": { + "SIMPLE": "target-group", + "MEDIUM": "target-group", + "COMPLEX": "target-group", + "REASONING": "target-group", + }, + }, + }, + }, + { + "model_name": "target-group", + "litellm_params": { + "model": "azure/computer-use-preview-1", + "api_key": "mock-api-key-1", + "api_version": "mock-api-version", + "api_base": "https://mock-endpoint-1.openai.azure.com", + }, + "model_info": {"id": "deployment-1", "base_model": "computer-use-preview"}, + }, + { + "model_name": "target-group", + "litellm_params": { + "model": "azure/computer-use-preview-2", + "api_key": "mock-api-key-2", + "api_version": "mock-api-version-2", + "api_base": "https://mock-endpoint-2.openai.azure.com", + }, + "model_info": {"id": "deployment-2", "base_model": "computer-use-preview"}, + }, + ], + ) + + +def _session_pin_key(session_id, user_key): + return DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="target-group", session_id=session_id, user_key=user_key + ) + + +def _cleanup_router_callbacks(router): + for callback in router.optional_callbacks or []: + litellm.logging_callback_manager.remove_callback_from_all_lists(callback) + + +async def _one_turn(router, model, session_id, key_hash): + """One request with the shuffle forced to deployment-1, so any other landing + deployment can only come from a pin read.""" + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=lambda seq: seq[0], + ), + ): + mock_post.return_value = MockResponse(MOCK_RESPONSES_API_RESPONSE, 200) + response = await router.aresponses( + model=model, + input=f"turn for {session_id} {key_hash}", + litellm_metadata={"session_id": session_id, "user_api_key_hash": key_hash}, + ) + return response._hidden_params["model_id"] + + +@pytest.mark.asyncio +async def test_auto_router_session_affinity_writes_scoped_pin_and_follows_it(): + """Turn 1 persists a key-scoped deployment pin; a pin seeded to the deployment + the shuffle would never pick is then followed, proving the read path.""" + router = _smart_router() + try: + served = await _one_turn(router, "smart-router", "write-session", "key-1") + assert await router.cache.async_get_cache(key=_session_pin_key("write-session", "key-1")) == { + "model_id": served + } + assert await router.cache.async_get_cache(key=_session_pin_key("write-session", None)) is None + + await router.cache.async_set_cache( + key=_session_pin_key("read-session", "key-1"), value={"model_id": "deployment-2"} + ) + assert await _one_turn(router, "smart-router", "read-session", "key-1") == "deployment-2" + finally: + _cleanup_router_callbacks(router) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,key_hash", + [ + ("target-group", "key-1"), + ("smart-router", "key-2"), + ], + ids=["direct-group-call", "different-api-key"], +) +async def test_seeded_session_pin_is_invisible_outside_its_scope(model, key_hash): + """The pin binds (auto-routed request, api key, session): a direct call to the + group and a different key reusing the session id must both ignore it.""" + router = _smart_router() + try: + await router.cache.async_set_cache( + key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"} + ) + assert await _one_turn(router, model, "scoped-session", key_hash) == "deployment-1" + finally: + _cleanup_router_callbacks(router) + + +@pytest.mark.asyncio +async def test_marker_write_uses_marker_ttl_and_writes_only_the_session_pin(): + """The write hook honors the marker's TTL over the callback default and writes + no user-key entry when only session affinity is active.""" + import time as time_module + + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + await callback.async_pre_call_deployment_hook( + kwargs={ + "model_info": {"id": "deployment-1"}, + "metadata": { + "deployment_model_name": "target-group", + "session_id": "ttl-session", + "user_api_key_hash": "key-1", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777, + }, + }, + call_type=None, + ) + + session_key = _session_pin_key("ttl-session", "key-1") + assert cache.in_memory_cache.cache_dict == {session_key: {"model_id": "deployment-1"}} + assert cache.in_memory_cache.ttl_dict[session_key] == pytest.approx(time_module.time() + 777, abs=5) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_marker", ["777", True, -5, 0, None]) +async def test_malformed_marker_values_do_not_enable_session_affinity(bad_marker): + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + await callback.async_pre_call_deployment_hook( + kwargs={ + "model_info": {"id": "deployment-1"}, + "metadata": { + "deployment_model_name": "target-group", + "session_id": "bad-marker-session", + "user_api_key_hash": "key-1", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: bad_marker, + }, + }, + call_type=None, + ) + + assert cache.in_memory_cache.cache_dict == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enable_user_key", [False, True], ids=["session-pin", "user-key-pin"]) +async def test_concurrent_first_requests_never_flip_a_claimed_pin(enable_user_key): + """Two overlapping first requests select different deployments before either + write lands. Pins are first-writer-wins claims, so the second write must leave + the stored pin unchanged instead of flipping it.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=enable_user_key, + enable_responses_api_affinity=False, + ) + + def racing_kwargs(deployment_id): + metadata = {"deployment_model_name": "target-group", "user_api_key_hash": "key-1"} + if not enable_user_key: + metadata["session_id"] = "racing-session" + metadata[SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] = 777 + return {"model_info": {"id": deployment_id}, "metadata": metadata} + + await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-1"), call_type=None) + await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-2"), call_type=None) + + pinned_key = ( + DeploymentAffinityCheck.get_affinity_cache_key(model_group="target-group", user_key="key-1") + if enable_user_key + else _session_pin_key("racing-session", "key-1") + ) + assert await cache.async_get_cache(key=pinned_key) == {"model_id": "deployment-1"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stored_pin", + [{"model_id": "deployment-1"}, "deployment-1"], + ids=["dict-pin", "legacy-string-pin"], +) +async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployment(stored_pin): + """The pod-local claim mirrors the Lua keepalive: the winning deployment's + re-claim extends the pin's expiry, a losing deployment's claim touches neither + the value nor the expiry, so no-Redis setups keep stickiness across an active + session and ttl bounds idle time there too. Sameness is judged on the pinned + model id, so a legacy string pin written by the Redis branch slides the same.""" + import time as time_module + + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + pin_key = _session_pin_key("slide-session", "key-1") + cache.in_memory_cache.set_cache(pin_key, stored_pin, ttl=10) + first_expiry = cache.in_memory_cache.ttl_dict[pin_key] + + reclaimed = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-1"}, ttl_seconds=777) + assert reclaimed == "deployment-1" + assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) + assert cache.in_memory_cache.ttl_dict[pin_key] > first_expiry + + lost = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-2"}, ttl_seconds=10) + assert lost == "deployment-1" + assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) + + +@pytest.mark.asyncio +async def test_claim_pin_uses_redis_attached_after_construction(): + """The proxy attaches Redis via Router._update_redis_cache after the Router (and + this callback) are built. The claim must resolve the redis tier per call, or pins + silently stay pod-local and cross-pod first-writer-wins is lost.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + captured = {} + + async def fake_runner(keys, args, client=None): + captured["keys"] = keys + captured["args"] = args + return b'{"model_id": "other-pod-winner"}' + + late_redis = MagicMock() + late_redis.async_register_script = MagicMock(return_value=fake_runner) + cache.redis_cache = late_redis + + import time as time_module + + pin_key = _session_pin_key("late-redis-session", "key-1") + cache.in_memory_cache.set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10) + + claimed = await callback._claim_pin( + cache_key=pin_key, + pin_value={"model_id": "our-deployment"}, + ttl_seconds=777, + ) + + assert claimed == "other-pod-winner" + assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) + assert captured["keys"] == (pin_key,) + assert captured["args"] == ('{"model_id": "our-deployment"}', 777) + assert cache.in_memory_cache.get_cache(_session_pin_key("late-redis-session", "key-1")) == { + "model_id": "other-pod-winner" + } + + +@pytest.mark.asyncio +async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): + """A Redis outage must cost cross-pod agreement, never same-pod stickiness. The write + hook only logs this result, so an escaping error would leave the session unpinned and + reshuffle every turn for the whole outage. DualCache's write path, which this claim + replaced, wrote the in-memory tier before ever touching Redis.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + async def exploding_runner(keys, args, client=None): + raise ConnectionError("redis is down") + + down_redis = MagicMock() + down_redis.async_register_script = MagicMock(return_value=exploding_runner) + cache.redis_cache = down_redis + + key = _session_pin_key("outage-session", "key-1") + claimed = await callback._claim_pin(cache_key=key, pin_value={"model_id": "our-deployment"}, ttl_seconds=777) + + assert claimed == "our-deployment" + assert cache.in_memory_cache.get_cache(key) == {"model_id": "our-deployment"} + + second = await callback._claim_pin(cache_key=key, pin_value={"model_id": "another-deployment"}, ttl_seconds=777) + assert second == "our-deployment" + + +@pytest.mark.asyncio +async def test_marker_session_affinity_read_and_write_agree_for_wildcard_groups(): + """Wildcard deployments keep the literal pattern as model_name on both the read + path and the write path, so the marker-gated pin round-trips through one key.""" + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + request_kwargs = { + "model_info": {"id": "wild-deployment-2"}, + "metadata": { + "deployment_model_name": "openai/*", + "session_id": "wild-session", + "user_api_key_hash": "key-1", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777, + }, + } + + await callback.async_pre_call_deployment_hook(kwargs=request_kwargs, call_type=None) + filtered = await callback.async_filter_deployments( + model="openai/gpt-4o", + healthy_deployments=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": f"wild-deployment-{i}"}, + } + for i in (1, 2) + ], + messages=[], + request_kwargs=request_kwargs, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["wild-deployment-2"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,session_affinity,deployment_affinity,expect_marker", + [ + ("smart-router", False, True, True), + ("smart-router", True, False, True), + ("smart-router", False, False, False), + ("target-group", False, True, False), + ], + ids=[ + "deployment-affinity-stamps", + "session-affinity-implies-deployment-pin", + "both-off-no-stamp", + "non-auto-routed-clears", + ], +) +async def test_pre_routing_hook_stamps_or_clears_the_marker_per_attempt( + model, session_affinity, deployment_affinity, expect_marker +): + """Every routing attempt writes or clears the marker, so a fallback from an + auto-routed group to a plain group cannot carry a stale marker. session_affinity + implies the deployment pin: a session frozen onto one group must not re-shuffle + across that group's deployments.""" + router = _smart_router(session_affinity=session_affinity, deployment_affinity=deployment_affinity) + try: + request_kwargs = { + "metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111}, + "litellm_metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111}, + } + await router.async_pre_routing_hook( + model=model, + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "Hello"}], + ) + if expect_marker: + assert request_kwargs["litellm_metadata"][SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] == 777 + else: + assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["metadata"] + assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["litellm_metadata"] + finally: + _cleanup_router_callbacks(router) + + +def test_complexity_router_with_deployment_affinity_registers_affinity_callback(): + enabled = _smart_router() + session_only = _smart_router(session_affinity=True, deployment_affinity=False) + disabled = _smart_router(session_affinity=False, deployment_affinity=False) + try: + assert [ + (cb.enable_user_key_affinity, cb.enable_responses_api_affinity, cb.enable_session_id_affinity) + for cb in enabled.optional_callbacks or [] + if isinstance(cb, DeploymentAffinityCheck) + ] == [(False, False, False)] + assert any(isinstance(cb, DeploymentAffinityCheck) for cb in session_only.optional_callbacks or []) + assert not any(isinstance(cb, DeploymentAffinityCheck) for cb in disabled.optional_callbacks or []) + finally: + _cleanup_router_callbacks(enabled) + _cleanup_router_callbacks(session_only) + _cleanup_router_callbacks(disabled) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index da2d78edb738..67fa827a8e47 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7550,3 +7550,68 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): assert capture.messages, "the fallback failure path did not log at ERROR" assert huge_message not in "".join(capture.messages) assert max(len(message) for message in capture.messages) < 5_000 +def test_stamp_or_clear_metadata_key_writes_and_clears_both_buckets(): + request_kwargs = {"metadata": {}} + litellm.Router._stamp_or_clear_metadata_key(request_kwargs=request_kwargs, key="probe", value=7) + assert request_kwargs["metadata"]["probe"] == 7 + + stale_kwargs = {"metadata": {"probe": 7}, "litellm_metadata": {"probe": 7}} + litellm.Router._stamp_or_clear_metadata_key(request_kwargs=stale_kwargs, key="probe", value=None) + assert "probe" not in stale_kwargs["metadata"] + assert "probe" not in stale_kwargs["litellm_metadata"] + + +@pytest.mark.parametrize( + "complexity_router_config,expect_callback", + [ + ({"tiers": {"SIMPLE": "gpt-4o"}}, True), + ({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False}, False), + ({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False, "session_affinity": True}, True), + ], +) +def test_complexity_router_registers_affinity_callback_for_deployment_pin(complexity_router_config, expect_callback): + """The marker the complexity router stamps is inert unless a DeploymentAffinityCheck is + registered to read it, so deployment_affinity has to pull the callback in, and its default-on + means a bare config registers one. Opting out must skip the callback entirely rather than + register a filter that can never fire, including when session_affinity is on, since the two + pins are independent.""" + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + { + "model_name": "my-complexity-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": complexity_router_config, + }, + }, + ] + ) + try: + registered = any(isinstance(cb, DeploymentAffinityCheck) for cb in router.optional_callbacks or []) + assert registered is expect_callback + finally: + for cb in router.optional_callbacks or []: + litellm.logging_callback_manager.remove_callback_from_all_lists(cb) + + +def test_ensure_deployment_affinity_callback_is_idempotent(): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = litellm.Router(model_list=[]) + try: + router._ensure_deployment_affinity_callback() + router._ensure_deployment_affinity_callback() + affinity_callbacks = [ + cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) + ] + assert len(affinity_callbacks) == 1 + finally: + for cb in router.optional_callbacks or []: + litellm.logging_callback_manager.remove_callback_from_all_lists(cb) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3a670bc7345e..d621e85f09be 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16760 + "limit": 16758 }, "LIT011": { "limit": 5598 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a75c23da1cf8..b361a16a0e6e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31695,6 +31695,12 @@ export interface components { * @description Default model to use if tier cannot be determined */ default_model?: string | null; + /** + * Deployment Affinity + * @description When True and a session_id is resolvable on the request, pin the deployment chosen inside each routed model group and reuse it whenever the session returns to that group, without pinning which group the session routes to. Independent of session_affinity, which pins the model group instead (and always carries this deployment pin with it): with session_affinity off, every turn is still classified on its own merits while a session that escalates to a stronger tier and comes back still lands on the deployment it used before, which is what keeps a provider prompt cache warm. Pins are held per model group, so switching tiers does not disturb the pin left behind in the previous group. On by default because re-shuffling a conversation across deployments of the same model discards that cache for no benefit; set False to keep every turn load-balanced across the group, which is what a deployment set with tight per-deployment rate limits wants. Inert when no session_id is resolvable, since there is nothing to key a pin on, and suppressed when plugins are configured, for the same reason session_affinity is. + * @default true + */ + deployment_affinity: boolean; /** * Dimension Weights * @description Weights for each scoring dimension @@ -31752,13 +31758,13 @@ export interface components { semantic_keyword_matching: boolean; /** * Session Affinity - * @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors. + * @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors. Always implies the deployment pin regardless of deployment_affinity: the session sticks to one deployment of the pinned model, since freezing the model while re-shuffling its deployments would still go cache-cold. * @default false */ session_affinity: boolean; /** * Session Affinity Ttl Seconds - * @description TTL for the session affinity pin; refreshed on every cache hit + * @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity deployment pin, so it measures idle time for the session's routing decisions rather than total session length * @default 3600 */ session_affinity_ttl_seconds: number;