diff --git a/litellm/router.py b/litellm/router.py index cdede46ef5c..dc58484de72 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..3305c133315 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 @@ -86,6 +87,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 +97,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 +335,82 @@ 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, 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 + 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: + """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: + 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( + "DeploymentAffinityCheck: failed to read max_parallel_requests limit. error=%s", + e, + ) + return True + if not isinstance(cached_limit, MaxParallelRequestsLimit): + return True + return cached_limit.in_flight < cached_limit.max_parallel_requests + + async def _yield_session_pin_if_saturated( + self, + 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 + 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 + 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( + "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 +513,16 @@ 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=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 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..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 @@ -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,191 @@ 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" + ) + + +@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