From 5f29454ab783b23a3b01238c15f5cae4f3aaa9c3 Mon Sep 17 00:00:00 2001 From: Akshat Shah Date: Sat, 10 Oct 2026 23:22:57 +0000 Subject: [PATCH 1/6] feat(router): yield session-affinity pin when pinned deployment is saturated Adds an opt-in router setting, session_affinity_yield_on_saturated_deployment. When a session is pinned to a deployment that is at its max_parallel_requests limit and another deployment in the model group has capacity, the turn routes to the idle deployment instead of queueing or 429ing behind the pin. When every deployment is full the filter keeps today's behavior. The pin itself is kept so the session returns to its warm deployment once it frees up. Fixes #45868 --- litellm/router.py | 9 ++ .../deployment_affinity_check.py | 79 +++++++++++ .../test_session_id_affinity.py | 125 ++++++++++++++++++ 3 files changed, 213 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index fd16c289b0d..99be9e13978 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -966,6 +966,7 @@ class Router: router_general_settings: RouterGeneralSettings | None = RouterGeneralSettings(), deployment_affinity_ttl_seconds: int = 3600, model_group_affinity_config: dict[str, list[str]] | None = None, + session_affinity_yield_on_saturated_deployment: bool = False, ignore_invalid_deployments: bool = False, enable_health_check_routing: bool = False, health_check_staleness_threshold: int | None = None, @@ -1011,6 +1012,7 @@ class Router: alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None. provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None. deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. + session_affinity_yield_on_saturated_deployment (bool): When True, a session pin yields to another deployment in the model group that still has max_parallel_requests capacity once the pinned deployment is saturated. Defaults to False. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). @@ -1177,6 +1179,7 @@ class Router: self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds self.model_group_affinity_config = model_group_affinity_config + self.session_affinity_yield_on_saturated_deployment = session_affinity_yield_on_saturated_deployment warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) if model_list is not None: @@ -2458,6 +2461,7 @@ class Router: enable_session_id_affinity=False, model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, + enable_session_affinity_yield_on_saturated_deployment=self.session_affinity_yield_on_saturated_deployment, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -2492,6 +2496,10 @@ class Router: existing_affinity_callback.enable_session_id_affinity = ( existing_affinity_callback.enable_session_id_affinity or enable_session_id_affinity ) + existing_affinity_callback.enable_session_affinity_yield_on_saturated_deployment = ( + existing_affinity_callback.enable_session_affinity_yield_on_saturated_deployment + or self.session_affinity_yield_on_saturated_deployment + ) existing_affinity_callback.ttl_seconds = self.deployment_affinity_ttl_seconds if self.model_group_affinity_config: existing_affinity_callback.model_group_affinity_config = self.model_group_affinity_config @@ -2504,6 +2512,7 @@ class Router: enable_session_id_affinity=enable_session_id_affinity, model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, + enable_session_affinity_yield_on_saturated_deployment=self.session_affinity_yield_on_saturated_deployment, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) 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 432fe11dc2e..dcf1d4d71e7 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -86,6 +86,7 @@ class DeploymentAffinityCheck(CustomLogger): enable_session_id_affinity: bool = False, model_group_affinity_config: dict[str, list[str]] | None = None, is_priority_group: Callable[[str], bool] | None = None, + enable_session_affinity_yield_on_saturated_deployment: bool = False, ): super().__init__() self.cache = cache @@ -95,6 +96,9 @@ class DeploymentAffinityCheck(CustomLogger): self.enable_session_id_affinity = enable_session_id_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} self.is_priority_group = is_priority_group + self.enable_session_affinity_yield_on_saturated_deployment = ( + enable_session_affinity_yield_on_saturated_deployment + ) def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -330,6 +334,75 @@ class DeploymentAffinityCheck(CustomLogger): def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) + @staticmethod + def _max_parallel_requests_cache_key(model_id: str) -> str: + return f"{model_id}_max_parallel_requests_client" + + @staticmethod + def _deployment_model_id(deployment: dict) -> str | None: + model_info: Final = deployment.get("model_info") + if not isinstance(model_info, dict): + return None + model_id: Final = model_info.get("id") + return str(model_id) if model_id is not None else None + + async def _deployment_has_capacity(self, model_id: str) -> bool: + """Best-effort capacity probe for max_parallel_requests slots. + + True when the deployment has no slot limit registered (unbounded) or + fewer requests in flight than its limit. Fail-open: any cache hiccup or + unreadable limit value counts as capacity, so this never blocks routing. + """ + try: + limit: Final = await self.cache.async_get_cache(key=self._max_parallel_requests_cache_key(model_id)) + except Exception as e: # noqa: BLE001 # affinity is best-effort; never break routing + verbose_router_logger.debug( + "DeploymentAffinityCheck: failed to read max_parallel_requests limit. error=%s", + e, + ) + return True + if limit is None: + return True + in_flight: Final = getattr(limit, "in_flight", None) + max_parallel_requests: Final = getattr(limit, "max_parallel_requests", None) + if not isinstance(in_flight, int) or not isinstance(max_parallel_requests, int): + return True + return in_flight < max_parallel_requests + + async def _yield_session_pin_if_saturated( + self, + session_deployment: dict, + healthy_deployments: list[dict], + ) -> list[dict] | None: + """Idle deployments when the pinned one is saturated, else None. + + Opt-in via session_affinity_yield_on_saturated_deployment. A turn yields + only when the pinned deployment is at its max_parallel_requests limit + and at least one sibling in the model group has capacity; when every + deployment is full it returns None so the caller keeps today's behavior + (wait behind the pin or 429). The pin itself is kept, so the session + returns to its warm deployment once it frees up. + """ + if not self.enable_session_affinity_yield_on_saturated_deployment: + return None + pinned_model_id: Final = self._deployment_model_id(session_deployment) + if pinned_model_id is None or await self._deployment_has_capacity(pinned_model_id): + return None + idle_deployments: Final = [] + for deployment in healthy_deployments: + deployment_model_id: Final = self._deployment_model_id(deployment) + if deployment_model_id is None or deployment_model_id == pinned_model_id: + continue + if await self._deployment_has_capacity(deployment_model_id): + idle_deployments.append(deployment) + if not idle_deployments: + return None + verbose_router_logger.debug( + "DeploymentAffinityCheck: session pin yields, pinned deployment=%s is saturated", + pinned_model_id, + ) + return idle_deployments + async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds) return self._pinned_model_id(winner) @@ -432,6 +505,12 @@ class DeploymentAffinityCheck(CustomLogger): model_id=session_model_id, ) if session_deployment is not None: + yielded_deployments: Final = await self._yield_session_pin_if_saturated( + session_deployment=session_deployment, + healthy_deployments=typed_healthy_deployments, + ) + if yielded_deployments is not None: + return yielded_deployments verbose_router_logger.debug( "DeploymentAffinityCheck: session-id affinity hit -> deployment=%s session_id=%s", session_model_id, diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 9bbaed0ae1b..8d7d2582ba8 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -10,6 +10,7 @@ from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY +from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -771,3 +772,127 @@ def test_complexity_router_with_deployment_affinity_registers_affinity_callback( _cleanup_router_callbacks(enabled) _cleanup_router_callbacks(session_only) _cleanup_router_callbacks(disabled) + + +def _saturate_deployment(cache: DualCache, model_id: str, model_group: str, max_parallel_requests: int) -> None: + limit: Final = MaxParallelRequestsLimit( + max_parallel_requests=max_parallel_requests, + model_id=model_id, + model_group=model_group, + ) + for _ in range(max_parallel_requests): + limit.acquire() + cache.set_cache(key=f"{model_id}_max_parallel_requests_client", value=limit, local_only=True) + + +@pytest.mark.asyncio +async def test_session_affinity_yield_saturated_pin_routes_to_idle_deployment(): + """A session pin yields when its deployment is at max_parallel_requests and a sibling has room. + + Regression test for https://github.com/BerriAI/litellm/issues/45868. + """ + cache: Final = DualCache() + callback: Final = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + enable_session_affinity_yield_on_saturated_deployment=True, + ) + deployments: Final = [ + { + "model_name": "group", + "litellm_params": {"model": "openai/model-a"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/model-b"}, + "model_info": {"id": "dep-b"}, + }, + ] + session_kwargs: Final = {"metadata": {"session_id": "s1"}} + session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + await cache.async_set_cache(session_cache_key, {"model_id": "dep-b"}) + + _saturate_deployment(cache, model_id="dep-b", model_group="group", max_parallel_requests=3) + + filtered: Final = await callback.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=[], + request_kwargs=session_kwargs, + ) + assert [d["model_info"]["id"] for d in filtered] == ["dep-a"], ( + "with dep-b saturated the turn must yield to the idle sibling instead of queueing or 429ing behind the pin" + ) + assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-b"}, ( + "the pin stays on dep-b, so the session returns to its warm deployment once it frees up" + ) + + +@pytest.mark.asyncio +async def test_session_affinity_yield_keeps_todays_behavior_when_all_saturated_or_disabled(): + """Yielding only happens when another deployment actually has capacity. + + Regression test for https://github.com/BerriAI/litellm/issues/45868. + """ + deployments: Final = [ + { + "model_name": "group", + "litellm_params": {"model": "openai/model-a"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/model-b"}, + "model_info": {"id": "dep-b"}, + }, + ] + session_kwargs: Final = {"metadata": {"session_id": "s1"}} + + cache_all_full: Final = DualCache() + callback_yield: Final = DeploymentAffinityCheck( + cache=cache_all_full, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + enable_session_affinity_yield_on_saturated_deployment=True, + ) + session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + await cache_all_full.async_set_cache(session_cache_key, {"model_id": "dep-b"}) + _saturate_deployment(cache_all_full, model_id="dep-a", model_group="group", max_parallel_requests=3) + _saturate_deployment(cache_all_full, model_id="dep-b", model_group="group", max_parallel_requests=3) + + filtered_all_full: Final = await callback_yield.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=[], + request_kwargs=session_kwargs, + ) + assert [d["model_info"]["id"] for d in filtered_all_full] == ["dep-b"], ( + "with every deployment full the turn must stay behind the pin, exactly like today" + ) + + cache_yield_off: Final = DualCache() + callback_off: Final = DeploymentAffinityCheck( + cache=cache_yield_off, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + await cache_yield_off.async_set_cache(session_cache_key, {"model_id": "dep-b"}) + _saturate_deployment(cache_yield_off, model_id="dep-b", model_group="group", max_parallel_requests=3) + + filtered_off: Final = await callback_off.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=[], + request_kwargs=session_kwargs, + ) + assert [d["model_info"]["id"] for d in filtered_off] == ["dep-b"], ( + "with the setting off a saturated pin must keep narrowing to it, exactly like today" + ) From d96f7fb6bc5f484f48319d061978da5e2b9ec84d Mon Sep 17 00:00:00 2001 From: akshatshahh Date: Sat, 10 Oct 2026 17:15:39 -0700 Subject: [PATCH 2/6] fix(router): satisfy type-discipline gate for session-affinity yield helper --- .../pre_call_checks/deployment_affinity_check.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) 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 dcf1d4d71e7..772bed84456 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -339,7 +339,9 @@ class DeploymentAffinityCheck(CustomLogger): return f"{model_id}_max_parallel_requests_client" @staticmethod - def _deployment_model_id(deployment: dict) -> str | None: + def _deployment_model_id( + deployment: dict, # mutable-ok: router deployment payloads are plain dicts by codebase convention; read-only access here + ) -> str | None: model_info: Final = deployment.get("model_info") if not isinstance(model_info, dict): return None @@ -371,9 +373,9 @@ class DeploymentAffinityCheck(CustomLogger): async def _yield_session_pin_if_saturated( self, - session_deployment: dict, - healthy_deployments: list[dict], - ) -> list[dict] | None: + session_deployment: dict, # mutable-ok: router deployment payloads are plain dicts by codebase convention; read-only access here + healthy_deployments: Sequence[dict], # mutable-ok: router deployment payloads are plain dicts by codebase convention; read-only access here + ) -> Sequence[dict] | None: # mutable-ok: router deployment payloads are plain dicts by codebase convention; read-only access here """Idle deployments when the pinned one is saturated, else None. Opt-in via session_affinity_yield_on_saturated_deployment. A turn yields From 9fba4ab40214247850f25a2fc954e07e6e19908b Mon Sep 17 00:00:00 2001 From: akshatshahh Date: Sat, 10 Oct 2026 18:24:37 -0700 Subject: [PATCH 3/6] fix(router): keep type-discipline suppression within line length --- .../pre_call_checks/deployment_affinity_check.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 772bed84456..e7f6c1fe33a 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -340,7 +340,7 @@ class DeploymentAffinityCheck(CustomLogger): @staticmethod def _deployment_model_id( - deployment: dict, # mutable-ok: router deployment payloads are plain dicts by codebase convention; read-only access here + deployment: dict, # mutable-ok: deployment payloads are dicts by router convention; read-only here ) -> str | None: model_info: Final = deployment.get("model_info") if not isinstance(model_info, dict): @@ -373,9 +373,9 @@ class DeploymentAffinityCheck(CustomLogger): async def _yield_session_pin_if_saturated( self, - session_deployment: dict, # mutable-ok: router deployment payloads are plain dicts by codebase convention; read-only access here - healthy_deployments: Sequence[dict], # mutable-ok: router deployment payloads are plain dicts by codebase convention; read-only access here - ) -> Sequence[dict] | None: # mutable-ok: router deployment payloads are plain dicts by codebase convention; read-only access here + session_deployment: dict, # mutable-ok: read-only router dicts + healthy_deployments: Sequence[dict], # mutable-ok: read-only router dicts + ) -> Sequence[dict] | None: # mutable-ok: read-only router dicts """Idle deployments when the pinned one is saturated, else None. Opt-in via session_affinity_yield_on_saturated_deployment. A turn yields From 638fb7ab674fe125e71e17efe2ad254892426dc1 Mon Sep 17 00:00:00 2001 From: akshatshahh Date: Sun, 11 Oct 2026 10:43:33 -0700 Subject: [PATCH 4/6] fix(router): fully type session-affinity yield helpers for basedpyright gate --- .../deployment_affinity_check.py | 41 ++++++++++--------- 1 file changed, 22 insertions(+), 19 deletions(-) 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 e7f6c1fe33a..617d103a057 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -30,6 +30,7 @@ from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes @@ -340,12 +341,12 @@ class DeploymentAffinityCheck(CustomLogger): @staticmethod def _deployment_model_id( - deployment: dict, # mutable-ok: deployment payloads are dicts by router convention; read-only here + deployment: dict[str, object], # mutable-ok: deployment payloads are dicts by router convention; read-only here ) -> str | None: model_info: Final = deployment.get("model_info") if not isinstance(model_info, dict): return None - model_id: Final = model_info.get("id") + model_id: Final[object | None] = model_info.get("id") return str(model_id) if model_id is not None else None async def _deployment_has_capacity(self, model_id: str) -> bool: @@ -356,26 +357,24 @@ class DeploymentAffinityCheck(CustomLogger): unreadable limit value counts as capacity, so this never blocks routing. """ try: - limit: Final = await self.cache.async_get_cache(key=self._max_parallel_requests_cache_key(model_id)) + limit: Final[object | None] = await self.cache.async_get_cache( + key=self._max_parallel_requests_cache_key(model_id) + ) except Exception as e: # noqa: BLE001 # affinity is best-effort; never break routing verbose_router_logger.debug( "DeploymentAffinityCheck: failed to read max_parallel_requests limit. error=%s", e, ) return True - if limit is None: + if not isinstance(limit, MaxParallelRequestsLimit): return True - in_flight: Final = getattr(limit, "in_flight", None) - max_parallel_requests: Final = getattr(limit, "max_parallel_requests", None) - if not isinstance(in_flight, int) or not isinstance(max_parallel_requests, int): - return True - return in_flight < max_parallel_requests + return limit.in_flight < limit.max_parallel_requests async def _yield_session_pin_if_saturated( self, - session_deployment: dict, # mutable-ok: read-only router dicts - healthy_deployments: Sequence[dict], # mutable-ok: read-only router dicts - ) -> Sequence[dict] | None: # mutable-ok: read-only router dicts + session_deployment: dict[str, object], # mutable-ok: read-only router dicts + healthy_deployments: Sequence[dict[str, object]], # mutable-ok: read-only router dicts + ) -> list[dict[str, object]] | None: # mutable-ok: read-only router dicts """Idle deployments when the pinned one is saturated, else None. Opt-in via session_affinity_yield_on_saturated_deployment. A turn yields @@ -390,13 +389,17 @@ class DeploymentAffinityCheck(CustomLogger): pinned_model_id: Final = self._deployment_model_id(session_deployment) if pinned_model_id is None or await self._deployment_has_capacity(pinned_model_id): return None - idle_deployments: Final = [] - for deployment in healthy_deployments: - deployment_model_id: Final = self._deployment_model_id(deployment) - if deployment_model_id is None or deployment_model_id == pinned_model_id: - continue - if await self._deployment_has_capacity(deployment_model_id): - idle_deployments.append(deployment) + candidates: Final = [ + (deployment, model_id) + for deployment in healthy_deployments + if (model_id := self._deployment_model_id(deployment)) is not None + and model_id != pinned_model_id + ] + idle_deployments: Final = [ + deployment + for deployment, model_id in candidates + if await self._deployment_has_capacity(model_id) + ] if not idle_deployments: return None verbose_router_logger.debug( From 0536659d95463d190ec9cc0da51df4eb4e2f8e2e Mon Sep 17 00:00:00 2001 From: akshatshahh Date: Sun, 11 Oct 2026 11:23:17 -0700 Subject: [PATCH 5/6] fix(router): apply ruff format to session-affinity yield helpers --- .../pre_call_checks/deployment_affinity_check.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) 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 617d103a057..3ce8c979249 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -392,13 +392,10 @@ class DeploymentAffinityCheck(CustomLogger): candidates: Final = [ (deployment, model_id) for deployment in healthy_deployments - if (model_id := self._deployment_model_id(deployment)) is not None - and model_id != pinned_model_id + if (model_id := self._deployment_model_id(deployment)) is not None and model_id != pinned_model_id ] idle_deployments: Final = [ - deployment - for deployment, model_id in candidates - if await self._deployment_has_capacity(model_id) + deployment for deployment, model_id in candidates if await self._deployment_has_capacity(model_id) ] if not idle_deployments: return None From 297326c8d3d5e256470e8d0b269cfc087ba9f2f2 Mon Sep 17 00:00:00 2001 From: akshatshahh Date: Sun, 11 Oct 2026 19:38:01 -0700 Subject: [PATCH 6/6] fix(router): eliminate new basedpyright unknown-type errors in session-affinity yield --- .../deployment_affinity_check.py | 24 +++++-- .../test_session_id_affinity.py | 64 +++++++++++++++++++ 2 files changed, 81 insertions(+), 7 deletions(-) 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 3ce8c979249..3305c133315 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -346,7 +346,10 @@ class DeploymentAffinityCheck(CustomLogger): model_info: Final = deployment.get("model_info") if not isinstance(model_info, dict): return None - model_id: Final[object | None] = model_info.get("id") + typed_model_info: Final = cast( + dict[str, object], model_info + ) # cast-ok: router deployment payloads are str-keyed dicts + model_id: Final = typed_model_info.get("id") return str(model_id) if model_id is not None else None async def _deployment_has_capacity(self, model_id: str) -> bool: @@ -357,8 +360,11 @@ class DeploymentAffinityCheck(CustomLogger): unreadable limit value counts as capacity, so this never blocks routing. """ try: - limit: Final[object | None] = await self.cache.async_get_cache( - key=self._max_parallel_requests_cache_key(model_id) + cached_limit: Final = cast( # cast-ok: DualCache returns untyped values; isinstance narrows below + object | None, + await self.cache.async_get_cache( # pyright: ignore[reportUnknownMemberType] # DualCache accessors are untyped legacy; this read is best-effort + key=self._max_parallel_requests_cache_key(model_id) + ), ) except Exception as e: # noqa: BLE001 # affinity is best-effort; never break routing verbose_router_logger.debug( @@ -366,9 +372,9 @@ class DeploymentAffinityCheck(CustomLogger): e, ) return True - if not isinstance(limit, MaxParallelRequestsLimit): + if not isinstance(cached_limit, MaxParallelRequestsLimit): return True - return limit.in_flight < limit.max_parallel_requests + return cached_limit.in_flight < cached_limit.max_parallel_requests async def _yield_session_pin_if_saturated( self, @@ -508,8 +514,12 @@ class DeploymentAffinityCheck(CustomLogger): ) if session_deployment is not None: yielded_deployments: Final = await self._yield_session_pin_if_saturated( - session_deployment=session_deployment, - healthy_deployments=typed_healthy_deployments, + session_deployment=cast( + dict[str, object], session_deployment + ), # cast-ok: router deployment payloads are str-keyed dicts + healthy_deployments=cast( + Sequence[dict[str, object]], typed_healthy_deployments + ), # cast-ok: router deployment payloads are str-keyed dicts ) if yielded_deployments is not None: return yielded_deployments diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 8d7d2582ba8..885c9625194 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -896,3 +896,67 @@ async def test_session_affinity_yield_keeps_todays_behavior_when_all_saturated_o assert [d["model_info"]["id"] for d in filtered_off] == ["dep-b"], ( "with the setting off a saturated pin must keep narrowing to it, exactly like today" ) + + +@pytest.mark.asyncio +async def test_session_affinity_yield_keeps_pin_when_pinned_deployment_has_capacity(): + """No yield when the pinned deployment still has max_parallel_requests room. + + Regression test for https://github.com/BerriAI/litellm/issues/45868. + """ + cache: Final = DualCache() + callback: Final = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + enable_session_affinity_yield_on_saturated_deployment=True, + ) + deployments: Final = [ + { + "model_name": "group", + "litellm_params": {"model": "openai/model-a"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/model-b"}, + "model_info": {"id": "dep-b"}, + }, + ] + session_kwargs: Final = {"metadata": {"session_id": "s1"}} + session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + await cache.async_set_cache(session_cache_key, {"model_id": "dep-b"}) + + filtered: Final = await callback.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=[], + request_kwargs=session_kwargs, + ) + assert [d["model_info"]["id"] for d in filtered] == ["dep-b"], ( + "with room on the pinned deployment the session must stay pinned, no yield" + ) + + +@pytest.mark.asyncio +async def test_session_affinity_yield_fails_open_when_capacity_probe_errors(): + """A cache failure during the capacity probe must never break routing. + + Regression test for https://github.com/BerriAI/litellm/issues/45868. + """ + + class _FailingCache(DualCache): + async def async_get_cache(self, key: str, **kwargs: object) -> object: + raise RuntimeError("cache is down") + + callback: Final = DeploymentAffinityCheck( + cache=_FailingCache(), + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + enable_session_affinity_yield_on_saturated_deployment=True, + ) + assert await callback._deployment_has_capacity("dep-a") is True