From 297326c8d3d5e256470e8d0b269cfc087ba9f2f2 Mon Sep 17 00:00:00 2001 From: akshatshahh Date: Sun, 11 Oct 2026 19:38:01 -0700 Subject: [PATCH] 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