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): tighten complexity deployment affinity

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 04:10:26 +00:00
parent eaed7dc537
commit 384f90ecaa
3 changed files with 127 additions and 57 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

82
litellm/router.py
View file

@ -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."""

10
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
View file

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

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

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