scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(router): weaken optional affinity providers

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
tin 2026-08-06 05:11:27 +00:00
parent df1b93ea3d
commit f8f82c5909
3 changed files with 60 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
litellm/router.py
View file

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

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

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

38
tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py
View file

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