Merge 297326c8d3 into 102619d47a

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)

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,

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