From f8f82c5909cc5039c2c5e345e4e4e2258456c162 Mon Sep 17 00:00:00 2001 From: tin Date: Thu, 6 Aug 2026 05:11:27 +0000 Subject: [PATCH] fix(router): weaken optional affinity providers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 32 ++++++++++------ .../router_strategy/test_complexity_router.py | 2 + .../test_session_id_affinity.py | 38 +++++++++++++++++++ 3 files changed, 60 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 67af54b8047..2f7d7271d03 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1676,6 +1676,9 @@ class Router: break if existing_affinity_callback is not None: + existing_session_affinity_group_ttls_provider: Final = ( + self._get_complexity_router_session_affinity_group_ttls_provider() + ) existing_affinity_callback.enable_user_key_affinity = ( existing_affinity_callback.enable_user_key_affinity or enable_user_key_affinity ) @@ -1686,12 +1689,13 @@ class Router: existing_affinity_callback.enable_session_id_affinity or enable_session_id_affinity ) existing_affinity_callback.ttl_seconds = self.deployment_affinity_ttl_seconds - existing_affinity_callback.session_affinity_group_ttls = ( - self._get_complexity_router_session_affinity_group_ttls - ) + existing_affinity_callback.session_affinity_group_ttls = existing_session_affinity_group_ttls_provider if self.model_group_affinity_config: existing_affinity_callback.model_group_affinity_config = self.model_group_affinity_config else: + new_session_affinity_group_ttls_provider: Final = ( + self._get_complexity_router_session_affinity_group_ttls_provider() + ) affinity_callback: Final = DeploymentAffinityCheck( cache=self.cache, ttl_seconds=self.deployment_affinity_ttl_seconds, @@ -1699,7 +1703,7 @@ class Router: enable_responses_api_affinity=enable_responses_api_affinity, enable_session_id_affinity=enable_session_id_affinity, model_group_affinity_config=self.model_group_affinity_config, - session_affinity_group_ttls=self._get_complexity_router_session_affinity_group_ttls, + session_affinity_group_ttls=new_session_affinity_group_ttls_provider, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -7651,6 +7655,17 @@ class Router: } ) + def _get_complexity_router_session_affinity_group_ttls_provider(self) -> Callable[[], Mapping[str, int]]: + router_ref: Final = weakref.ref(self) + + def session_affinity_group_ttls_provider() -> Mapping[str, int]: + router: Final = router_ref() + if router is None: + return MappingProxyType({}) + return router._get_complexity_router_session_affinity_group_ttls() + + return session_affinity_group_ttls_provider + def _ensure_deployment_affinity_check(self) -> None: """ Ensure deployment affinity exists when explicit group settings or complexity routing require it. @@ -7661,14 +7676,7 @@ class Router: if self.optional_callbacks is None: self.optional_callbacks = [] - router_ref: Final = weakref.ref(self) - - def session_affinity_group_ttls_provider() -> Mapping[str, int]: - router: Final = router_ref() - if router is None: - return MappingProxyType({}) - return router._get_complexity_router_session_affinity_group_ttls() - + session_affinity_group_ttls_provider: Final = self._get_complexity_router_session_affinity_group_ttls_provider() existing_affinity_callback: DeploymentAffinityCheck | None = next( (callback for callback in self.optional_callbacks if isinstance(callback, DeploymentAffinityCheck)), None, diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 2a269affe32..9db9769a576 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -98,6 +98,8 @@ def test_router_complexity_session_affinity_provider_registers_callback(): assert tuple(_iter_complexity_router_session_affinity_groups(complexity_router)) == (("target-group", 17),) assert dict(parent_router._get_complexity_router_session_affinity_group_ttls()) == {"target-group": 17} + provider = parent_router._get_complexity_router_session_affinity_group_ttls_provider() + assert dict(provider()) == {"target-group": 17} parent_router._ensure_deployment_affinity_check() 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 abf9412c9fd..d634566d871 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 @@ -3,6 +3,7 @@ import gc import os import sys import weakref +from collections.abc import Callable from unittest.mock import AsyncMock, patch import pytest @@ -448,6 +449,43 @@ def test_complexity_router_affinity_provider_does_not_keep_router_alive(): litellm.logging_callback_manager.remove_callback_from_all_lists(affinity_callback) +def _assert_router_affinity_callback_does_not_keep_router_alive( + router_factory: Callable[[], litellm.Router], +) -> None: + router = router_factory() + affinity_callback = next( + callback for callback in router.optional_callbacks or [] if isinstance(callback, DeploymentAffinityCheck) + ) + router.optional_callbacks = [] + router.discard() + router_ref = weakref.ref(router) + + try: + del router + gc.collect() + + assert router_ref() is None + assert affinity_callback.session_affinity_group_ttls is not None + assert affinity_callback.session_affinity_group_ttls() == {} + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(affinity_callback) + + +def test_optional_pre_call_checks_affinity_provider_does_not_keep_new_router_alive(): + _assert_router_affinity_callback_does_not_keep_router_alive( + lambda: litellm.Router(model_list=[], optional_pre_call_checks=["session_affinity"]) + ) + + +def test_optional_pre_call_checks_affinity_provider_does_not_keep_existing_router_alive(): + def build_router() -> litellm.Router: + router = _complexity_router() + router.add_optional_pre_call_checks(["session_affinity"]) + return router + + _assert_router_affinity_callback_does_not_keep_router_alive(build_router) + + @pytest.mark.asyncio async def test_session_affinity_without_stable_model_map_key_falls_back_to_normal_selection(): callback = DeploymentAffinityCheck(