feat(router): independent, default-on deployment affinity for the auto-router

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

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

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

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(

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,

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(