From 384f90ecaaccd30df078aacca796d42605455dc7 Mon Sep 17 00:00:00 2001 From: tin Date: Thu, 6 Aug 2026 04:10:26 +0000 Subject: [PATCH] fix(router): tighten complexity deployment affinity Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 82 ++++++++++------- .../deployment_affinity_check.py | 10 +- .../test_session_id_affinity.py | 92 +++++++++++++++---- 3 files changed, 127 insertions(+), 57 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b8f980df953..81c80573c66 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -20,8 +20,8 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence -from functools import lru_cache +from collections.abc import AsyncGenerator, Callable, Generator, Iterator, Mapping, Sequence +from functools import lru_cache, reduce from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeVar, Union, cast @@ -254,6 +254,35 @@ else: PreRoutingHookResponse = Any +def _iter_complexity_router_session_affinity_groups( + complexity_router: "ComplexityRouter", +) -> Iterator[tuple[str, int]]: + config: Final = complexity_router.config + if not config.session_affinity or config.plugins: + return + + for tier_value in config.tiers.values(): + for model_group in tier_value if isinstance(tier_value, list) else (tier_value,): + yield model_group, config.session_affinity_ttl_seconds + + if config.default_model is not None: + yield config.default_model, config.session_affinity_ttl_seconds + + +def _merge_minimum_session_affinity_ttl( + group_ttls: Mapping[str, int], + group_ttl: tuple[str, int], +) -> Mapping[str, int]: + model_group, ttl = group_ttl + existing_ttl: Final[int | None] = group_ttls.get(model_group) + return MappingProxyType( + { + **group_ttls, + model_group: ttl if existing_ttl is None else min(existing_ttl, ttl), + } + ) + + def _cost_value_as_float(value: str | float | None) -> float | None: if value is None: return None @@ -7622,42 +7651,24 @@ class Router: return classify_strategy_router_model(litellm_params.model) == "complexity" def _get_complexity_router_session_affinity_group_ttls(self) -> Mapping[str, int]: - entries: Final = tuple( - (model_group, tagged_strategy.strategy.config.session_affinity_ttl_seconds) - for strategies in self.complexity_routers.values() - for tagged_strategy in strategies - if tagged_strategy.strategy.config.session_affinity and not tagged_strategy.strategy.config.plugins - for configured_models in ( - tuple( - model - for tier_value in tagged_strategy.strategy.config.tiers.values() - for model in (tier_value if isinstance(tier_value, list) else (tier_value,)) - ), - ) - for model_group in configured_models - ) - entries_with_defaults: Final = entries + tuple( + return reduce( + _merge_minimum_session_affinity_ttl, ( - tagged_strategy.strategy.config.default_model, - tagged_strategy.strategy.config.session_affinity_ttl_seconds, - ) - for strategies in self.complexity_routers.values() - for tagged_strategy in strategies - if tagged_strategy.strategy.config.session_affinity - and not tagged_strategy.strategy.config.plugins - and tagged_strategy.strategy.config.default_model is not None - ) - groups: Final = frozenset(model_group for model_group, _ in entries_with_defaults) - return MappingProxyType( - { - model_group: min( - ttl for candidate_group, ttl in entries_with_defaults if candidate_group == model_group - ) - for model_group in groups - } + group_ttl + for strategies in self.complexity_routers.values() + for tagged_strategy in strategies + for group_ttl in _iter_complexity_router_session_affinity_groups(tagged_strategy.strategy) + ), + MappingProxyType({}), ) def _ensure_deployment_affinity_check(self) -> None: + """ + Ensure deployment affinity exists when explicit group settings or complexity routing require it. + + Explicit model-group settings can enable affinity while global flags remain false, so those settings + also require the callback to exist. Otherwise the callback stays absent when nothing needs pinning. + """ if self.optional_callbacks is None: self.optional_callbacks = [] @@ -7728,7 +7739,8 @@ class Router: strategy=complexity_router, strategy_label="Complexity-router", ) - self._ensure_deployment_affinity_check() + if self._get_complexity_router_session_affinity_group_ttls(): + self._ensure_deployment_affinity_check() 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.""" 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 b37a4fdce1f..30071ae4a57 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -70,8 +70,7 @@ class DeploymentAffinityCheck(CustomLogger): 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", + "DeploymentAffinityCheck: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s", unknown, group, self.VALID_FLAGS, @@ -464,7 +463,7 @@ class DeploymentAffinityCheck(CustomLogger): self._get_session_id_from_request_kwargs(request_kwargs=kwargs) if enable_session_id else None ) - if user_key is None and session_id is None: + if (enable_user_key and user_key is None) and (enable_session_id and session_id is None): return None model_info = kwargs.get("model_info") @@ -488,7 +487,7 @@ class DeploymentAffinityCheck(CustomLogger): verbose_router_logger.warning("DeploymentAffinityCheck: model_id missing; skipping affinity cache update.") return None - if user_key is not None: + 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( @@ -525,8 +524,7 @@ class DeploymentAffinityCheck(CustomLogger): ttl=session_affinity_ttl, ) verbose_router_logger.debug( - "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s " - "ttl=%s session_id=%s", + "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s", deployment_model_name, model_id, session_affinity_ttl, 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 11379ad2736..3ad0a79af67 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 asyncio import os import sys -import asyncio from unittest.mock import AsyncMock, patch import pytest @@ -99,7 +99,7 @@ def _responses_mock() -> MockResponse: ) -def _deterministic_choice(): +def _first_then_last_choice(): choice_calls = {"count": 0} def choose(sequence): @@ -212,6 +212,32 @@ async def test_async_session_id_affinity_routes_to_same_deployment(): @pytest.mark.asyncio async def test_complexity_router_session_affinity_pins_deployment_and_scopes_api_key(): + baseline_router = _complexity_router(session_affinity=False) + + 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=_first_then_last_choice(), + ), + ): + mock_post.return_value = _responses_mock() + baseline_first = await baseline_router.aresponses( + model="smart-router", + input="Hello", + metadata={"session_id": "shared-session", "user_api_key_hash": "key-1"}, + ) + baseline_second = await baseline_router.aresponses( + model="smart-router", + input="Follow-up", + metadata={"session_id": "shared-session", "user_api_key_hash": "key-1"}, + ) + + assert baseline_first._hidden_params["model_id"] != baseline_second._hidden_params["model_id"] + router = _complexity_router() with ( @@ -221,7 +247,7 @@ async def test_complexity_router_session_affinity_pins_deployment_and_scopes_api ) as mock_post, patch( "litellm.router_strategy.simple_shuffle.random.choice", - side_effect=_deterministic_choice(), + side_effect=_first_then_last_choice(), ), ): mock_post.return_value = _responses_mock() @@ -262,7 +288,7 @@ async def test_complexity_router_affinity_falls_back_when_pinned_deployment_is_i ) as mock_post, patch( "litellm.router_strategy.simple_shuffle.random.choice", - side_effect=_deterministic_choice(), + side_effect=_first_then_last_choice(), ), ): mock_post.return_value = _responses_mock() @@ -313,9 +339,11 @@ async def test_complexity_router_session_affinity_uses_router_configured_ttl(): session_id="ttl-session", user_key="key-1", ) - assert (session_cache_key, {"model_id": "deployment-1"}) in [ - (call.args[0], call.args[1]) for call in cache.async_set_cache.call_args_list - ] + assert cache.async_set_cache.call_count == 1 + assert cache.async_set_cache.call_args.args[:2] == ( + session_cache_key, + {"model_id": "deployment-1"}, + ) assert any(call.kwargs.get("ttl") == 17 for call in cache.async_set_cache.call_args_list) @@ -330,7 +358,7 @@ async def test_complexity_router_session_affinity_expires_and_reselects(): ) as mock_post, patch( "litellm.router_strategy.simple_shuffle.random.choice", - side_effect=_deterministic_choice(), + side_effect=_first_then_last_choice(), ), ): mock_post.return_value = _responses_mock() @@ -349,7 +377,7 @@ async def test_complexity_router_session_affinity_expires_and_reselects(): assert second_response._hidden_params["model_id"] != first_response._hidden_params["model_id"] -def test_complexity_router_registers_model_pool_groups_and_respects_disabled_affinity(): +def test_complexity_router_registers_model_pool_groups(): router = _complexity_router( session_affinity=True, tiers={"SIMPLE": ["target-group", "pool-group"], "MEDIUM": "target-group"}, @@ -358,13 +386,11 @@ def test_complexity_router_registers_model_pool_groups_and_respects_disabled_aff "target-group": 7, "pool-group": 7, } - disabled_router = _complexity_router(session_affinity=False) - callback = next( - callback - for callback in disabled_router.optional_callbacks or [] - if isinstance(callback, DeploymentAffinityCheck) - ) - assert callback._get_effective_flags("target-group")[2] is False + + +def test_complexity_router_without_session_affinity_does_not_register_callback(): + router = _complexity_router(session_affinity=False) + assert not any(isinstance(callback, DeploymentAffinityCheck) for callback in router.optional_callbacks or []) def test_complexity_router_plugins_do_not_enable_deployment_affinity(): @@ -374,6 +400,40 @@ def test_complexity_router_plugins_do_not_enable_deployment_affinity(): router = _complexity_router(plugins=[NoOpPlugin()]) assert dict(router._get_complexity_router_session_affinity_group_ttls()) == {} + assert not any(isinstance(callback, DeploymentAffinityCheck) for callback in router.optional_callbacks or []) + + +@pytest.mark.asyncio +async def test_session_affinity_without_stable_model_map_key_falls_back_to_normal_selection(): + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + session_affinity_group_ttls=lambda: {"mixed-group": 60}, + ) + healthy_deployments = [ + { + "model_name": "mixed-group", + "litellm_params": {"model": "azure/deployment-1"}, + "model_info": {"id": "deployment-1"}, + }, + { + "model_name": "mixed-group", + "litellm_params": {"model": "bedrock/deployment-2"}, + "model_info": {"id": "deployment-2"}, + }, + ] + + filtered_deployments = await callback.async_filter_deployments( + model="mixed-group", + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": {"session_id": "mixed-session"}}, + parent_otel_span=None, + ) + + assert filtered_deployments == healthy_deployments @pytest.mark.asyncio