From 5e12c67aaed33acddd7187cdacf695770d1010e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:18:51 -0700 Subject: [PATCH] fix(router): take only the failing discovered Kubernetes pod out of rotation With kubernetes_pod_discovery, a 429 or a high 5xx rate from one pod cooled down the whole deployment, so the healthy pods behind the same Service went idle for the cooldown window. Each discovered pod is now judged by the deployment's cooldown rules with its sibling pods as the alternatives: the failing pod leaves rotation for the cooldown time and the deployment cools down only once no live pod is left. --- litellm/router.py | 69 ++--- litellm/router_utils/cooldown_handlers.py | 89 +++++- .../router_utils/fallback_event_handlers.py | 3 + .../router_utils/kubernetes_pod_discovery.py | 72 ++++- .../router_surface_allowlist.json | 1 + .../test_kubernetes_pod_discovery.py | 282 ++++++++++++++++++ .../test_router/test_router_helper_utils.py | 100 ++++--- 7 files changed, 520 insertions(+), 96 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index ff31f48c49e..0a1b67ff063 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -327,6 +327,7 @@ from litellm.router_utils.handle_error import ( ) from litellm.router_utils.health_state_cache import DeploymentHealthCache from litellm.router_utils.kubernetes_pod_discovery import ( + KubernetesPod, KubernetesPodDiscovery, async_resolve_pods_after, async_resolve_pods_after_bound, @@ -9516,27 +9517,28 @@ class Router: return model_group: Final = metadata.get("model_group", None) model_info: Final = litellm_params.get("model_info", {}) or {} - deployment_id: Final = model_info.get("id", None) - if model_group is None or deployment_id is None or self.get_deployment(model_id=str(deployment_id)) is None: + raw_deployment_id: Final = model_info.get("id", None) + if model_group is None or raw_deployment_id is None: return + deployment_id: Final = str(raw_deployment_id) + if self.get_deployment(model_id=deployment_id) is None: + return + typed_litellm_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params) if self.session_affinity_repin_on_failover: await self._repin_session_after_success( - _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params), + typed_litellm_params, _OPAQUE_VALUE_ADAPTER.validate_python(completion_response), ) # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits - increment_deployment_successes_for_current_minute( - litellm_router_instance=self, - deployment_id=str(deployment_id), - ) + self._count_deployment_success(deployment_id, typed_litellm_params) total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0) - counted_tokens: Final = get_counted_usage_tokens(litellm_params) + counted_tokens: Final = get_counted_usage_tokens(typed_litellm_params) deployment_name: Final = metadata.get("deployment", None) return await self._increment_deployment_usage( - deployment_id=str(deployment_id), + deployment_id=deployment_id, deployment_name=deployment_name if isinstance(deployment_name, str) else None, model_group=model_group, total_tokens=total_tokens if counted_tokens is None else max(0, total_tokens - counted_tokens), @@ -9694,26 +9696,28 @@ class Router: """ if is_batch_retrieve_call_type(kwargs.get("call_type")): return None - id = None - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - model_info: Final = kwargs["litellm_params"].get("model_info", {}) or {} - id = model_info.get("id", None) - if model_group is None or id is None: - return None - elif isinstance(id, int): - id = str(id) + litellm_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(kwargs["litellm_params"]) + metadata: Final = litellm_params.get("metadata") + if metadata is None: + return None + model_group: Final = _REQUEST_KWARGS_ADAPTER.validate_python(metadata).get("model_group") + deployment_id: Final = _MODEL_INFO_ADAPTER.validate_python(litellm_params.get("model_info") or {}).get("id") + if model_group is None or deployment_id is None: + return None + return self._count_deployment_success(str(deployment_id), litellm_params) - if id is not None: - key: Final = increment_deployment_successes_for_current_minute( - litellm_router_instance=self, - deployment_id=id, + def _kubernetes_pod_for_call(self, deployment_id: str, call_litellm_params: object) -> KubernetesPod | None: + return self.kubernetes_pod_discovery.pod_for_call(self.get_model_info(id=deployment_id), call_litellm_params) + + def _count_deployment_success(self, deployment_id: str, call_litellm_params: object) -> str: + pod: Final = self._kubernetes_pod_for_call(deployment_id, call_litellm_params) + if pod is not None: + increment_deployment_successes_for_current_minute( + litellm_router_instance=self, deployment_id=pod.counter_id ) - return key - - return None + return increment_deployment_successes_for_current_minute( + litellm_router_instance=self, deployment_id=deployment_id + ) def deployment_callback_on_failure( self, @@ -9822,19 +9826,16 @@ class Router: litellm_router_instance=self, deployment_id=deployment_id, ) + request_kwargs: Final = cast("dict[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict result: Final = set_cooldown_deployments( litellm_router_instance=self, exception_status=exception_status, original_exception=exception, deployment=deployment_id, time_to_cooldown=_time_to_cooldown, - requested_model_group=( - get_litellm_metadata_from_kwargs( - cast("dict[str, object]", kwargs) # cast-ok: untyped kwargs dict - ) - or {} - ).get("model_group"), - request_kwargs=cast("dict[str, object]", kwargs), # cast-ok: kwargs is the untyped request dict + requested_model_group=(get_litellm_metadata_from_kwargs(request_kwargs) or {}).get("model_group"), + request_kwargs=request_kwargs, + failed_pod=self._kubernetes_pod_for_call(deployment_id, request_kwargs.get("litellm_params")), ) # setting deployment_id in cooldown deployments return result diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index a1b1e327802..93c94ab961b 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCacheValue from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback +from litellm.router_utils.kubernetes_pod_discovery import KubernetesPod from litellm.types.router import Deployment from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -43,6 +44,7 @@ from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES from .router_callbacks.track_deployment_metrics import ( get_deployment_failures_for_current_minute, get_deployment_successes_for_current_minute, + increment_deployment_failures_for_current_minute, ) if TYPE_CHECKING: @@ -181,6 +183,7 @@ def _should_cooldown_based_on_deployment_policy( dep_policy: Mapping[str, int] | None, dep_allowed_fails: int | None, is_single_deployment_model_group: bool, + counter_id: str | None = None, ) -> bool: """Resolve deployment-level allowed-fails and delegate to the shared counting logic. @@ -224,7 +227,7 @@ def _should_cooldown_based_on_deployment_policy( return should_cooldown_based_on_allowed_fails_policy( litellm_router_instance=litellm_router_instance, - deployment=deployment, + deployment=deployment if counter_id is None else counter_id, original_exception=original_exception, allowed_fails_override=allowed_fails_override, cooldown_time_override=cooldown_time_override, @@ -461,12 +464,29 @@ def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_ return bool(fallback_group) +def _is_single_deployment_model_group( + litellm_router_instance: LitellmRouter, + deployment: str, + requested_model_group: str | None, +) -> bool: + model_group: Final = litellm_router_instance.get_model_group(id=deployment) + return ( + model_group is not None + and len(model_group) == 1 + and not ( + litellm_router_instance.routing_group_has_alternatives(requested_model_group) + or litellm_router_instance.team_model_has_alternatives(deployment) + ) + ) + + def _should_cooldown_deployment( litellm_router_instance: LitellmRouter, deployment: str, exception_status: str | int, original_exception: Exception, requested_model_group: str | None = None, + failed_pod: KubernetesPod | None = None, ) -> bool: """ Helper that decides if a deployment should be put in cooldown @@ -486,13 +506,10 @@ def _should_cooldown_deployment( - v1 logic (Legacy): if allowed fails or allowed fail policy set, coolsdown if num fails in this minute > allowed fails """ - model_group: Final = litellm_router_instance.get_model_group(id=deployment) - is_single_deployment_model_group = False - if model_group is not None and len(model_group) == 1: - is_single_deployment_model_group = not ( - litellm_router_instance.routing_group_has_alternatives(requested_model_group) - or litellm_router_instance.team_model_has_alternatives(deployment) - ) + is_single_deployment_model_group: Final = failed_pod is None and _is_single_deployment_model_group( + litellm_router_instance, deployment, requested_model_group + ) + counter_id: Final = deployment if failed_pod is None else failed_pod.counter_id is_lone_deployment: Final = is_single_deployment_model_group and not _has_fallback_group( litellm_router_instance, requested_model_group ) @@ -516,6 +533,7 @@ def _should_cooldown_deployment( dep_policy, dep_allowed_fails, is_lone_deployment, + counter_id, ) ## BASE CASE - single deployment @@ -524,10 +542,10 @@ def _should_cooldown_deployment( and _is_allowed_fails_set_on_router(litellm_router_instance=litellm_router_instance) is False ): num_successes_this_minute: Final = get_deployment_successes_for_current_minute( - litellm_router_instance=litellm_router_instance, deployment_id=deployment + litellm_router_instance=litellm_router_instance, deployment_id=counter_id ) num_fails_this_minute: Final = get_deployment_failures_for_current_minute( - litellm_router_instance=litellm_router_instance, deployment_id=deployment + litellm_router_instance=litellm_router_instance, deployment_id=counter_id ) total_requests_this_minute: Final = num_successes_this_minute + num_fails_this_minute @@ -564,7 +582,7 @@ def _should_cooldown_deployment( return False return should_cooldown_based_on_allowed_fails_policy( litellm_router_instance=litellm_router_instance, - deployment=deployment, + deployment=counter_id, original_exception=original_exception, ) @@ -579,6 +597,7 @@ def set_cooldown_deployments( time_to_cooldown: float | None = None, requested_model_group: str | None = None, request_kwargs: Mapping[str, object] | None = None, + failed_pod: KubernetesPod | None = None, ) -> bool: """ Add a model to the list of models being cooled down for that minute, if it exceeds the allowed fails / minute @@ -591,6 +610,11 @@ def set_cooldown_deployments( - True if the deployment should be put in cooldown - False if the deployment should not be put in cooldown """ + if failed_pod is not None: + increment_deployment_failures_for_current_minute( + litellm_router_instance=litellm_router_instance, deployment_id=failed_pod.counter_id + ) + verbose_router_logger.debug("checks 'should_run_cooldown_logic'") if ( @@ -608,6 +632,18 @@ def set_cooldown_deployments( verbose_router_logger.debug("should_run_cooldown_logic returned False") return False + if failed_pod is not None and litellm_router_instance.kubernetes_pod_discovery.has_other_live_pod(failed_pod): + _cool_down_kubernetes_pod_if_needed( + litellm_router_instance=litellm_router_instance, + failed_pod=failed_pod, + deployment=deployment, + exception_status=exception_status, + original_exception=original_exception, + time_to_cooldown=time_to_cooldown, + requested_model_group=requested_model_group, + ) + return False + exception_status_int: Final = cast_exception_status_to_int(exception_status) verbose_router_logger.debug("Attempting to add %s to cooldown list", deployment) @@ -641,6 +677,37 @@ def set_cooldown_deployments( _set_cooldown_deployments = set_cooldown_deployments +def _cool_down_kubernetes_pod_if_needed( + litellm_router_instance: LitellmRouter, + failed_pod: KubernetesPod, + deployment: str, + exception_status: str | int, + original_exception: Exception, + time_to_cooldown: float | None, + requested_model_group: str | None, +) -> None: + if not _should_cooldown_deployment( + litellm_router_instance=litellm_router_instance, + deployment=deployment, + exception_status=exception_status, + original_exception=original_exception, + requested_model_group=requested_model_group, + failed_pod=failed_pod, + ): + return + cooldown_seconds: Final = ( + time_to_cooldown if time_to_cooldown is not None else litellm_router_instance.cooldown_time + ) + litellm_router_instance.kubernetes_pod_discovery.cool_down_pod(failed_pod, cooldown_seconds) + verbose_router_logger.info( + "Kubernetes pod %s behind %s is out of rotation for %ss, its deployment %s stays up", + failed_pod.ip, + failed_pod.service_host, + cooldown_seconds, + deployment, + ) + + async def async_get_cooldown_deployments( litellm_router_instance: LitellmRouter, parent_otel_span: Span | None, diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 6d1de1f9f6c..68fc390c549 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -176,6 +176,9 @@ def _trigger_cooldown_for_failed_deployment( deployment=deployment_id, time_to_cooldown=time_to_cooldown, request_kwargs=kwargs, + failed_pod=litellm_router.kubernetes_pod_discovery.pod_for_call( + deployment_dict, model_call_details.get("litellm_params") + ), ) verbose_router_logger.debug("Triggered cooldown for fallback deployment %s", deployment_id) diff --git a/litellm/router_utils/kubernetes_pod_discovery.py b/litellm/router_utils/kubernetes_pod_discovery.py index 34d99db1f6e..2f1bfa9f26d 100644 --- a/litellm/router_utils/kubernetes_pod_discovery.py +++ b/litellm/router_utils/kubernetes_pod_discovery.py @@ -24,6 +24,8 @@ from typing import ( ) import httpx +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_router_logger from litellm.constants import ( @@ -52,6 +54,24 @@ class KubernetesPodLookup: error: str | None +class _CallParams(TypedDict, total=False): + api_base: ReadOnly[object] + + +_CALL_PARAMS_ADAPTER: Final = TypeAdapter(_CallParams) + + +@dataclass(frozen=True, slots=True) +class KubernetesPod: + service_host: str + port: int | None + ip: str + + @property + def counter_id(self) -> str: + return f"kubernetes_pod:{self.service_host}:{self.port}:{self.ip}" + + def _retry_count_from_metadata(metadata: object) -> int | None: if not isinstance(metadata, Mapping): return None @@ -102,6 +122,7 @@ class KubernetesPodDiscovery: self.proxy_environment = proxy_environment self.clock = clock self._cache: Mapping[_CacheKey, _PodSet] = MappingProxyType({}) + self._cooled_pods: Mapping[KubernetesPod, float] = MappingProxyType({}) self._refreshing: frozenset[_CacheKey] = frozenset() self._timed_out: frozenset[_CacheKey] = frozenset() self._https_warning_hosts: frozenset[str] = frozenset() @@ -208,6 +229,41 @@ class KubernetesPodDiscovery: return KubernetesPodLookup(service_host=host, port=port, pod_ips=(), error=str(result)) return KubernetesPodLookup(service_host=host, port=port, pod_ips=self._pod_ips(result), error=None) + def pod_for_call(self, deployment: object, call_litellm_params: object) -> KubernetesPod | None: + eligible: Final = self._eligible(deployment) + if eligible is None or not isinstance(call_litellm_params, Mapping): + return None + call_api_base: Final = _CALL_PARAMS_ADAPTER.validate_python(call_litellm_params).get("api_base") + if not isinstance(call_api_base, str): + return None + _, (host, port), _ = eligible + try: + pod_ip: Final = httpx.URL(call_api_base).host + except httpx.InvalidURL: + return None + with self._lock: + cached: Final = self._cache.get((host, port)) + if cached is None or pod_ip not in cached.ips: + return None + return KubernetesPod(service_host=host, port=port, ip=pod_ip) + + def has_other_live_pod(self, pod: KubernetesPod) -> bool: + now: Final = self.clock() + with self._lock: + cached: Final = self._cache.get((pod.service_host, pod.port)) + sibling_ips: Final = () if cached is None else tuple(ip for ip in cached.ips if ip != pod.ip) + return any(not self._is_cooled(replace(pod, ip=ip), now) for ip in sibling_ips) + + def cool_down_pod(self, pod: KubernetesPod, seconds: float) -> None: + now: Final = self.clock() + with self._lock: + self._cooled_pods = MappingProxyType( + {**{cooled: until for cooled, until in self._cooled_pods.items() if until > now}, pod: now + seconds} + ) + + def _is_cooled(self, pod: KubernetesPod, now: float) -> bool: + return self._cooled_pods.get(pod, now) > now + @staticmethod def _session_id(request_kwargs: Mapping[str, object] | None) -> str | None: if request_kwargs is None: @@ -451,18 +507,28 @@ class KubernetesPodDiscovery: if cached is None or not cached.ips: return None pod_count: Final = len(cached.ips) + live_ips: Final = tuple( + ip + for ip in cached.ips + if not self._is_cooled(KubernetesPod(service_host=key[0], port=key[1], ip=ip), now) + ) + candidate_ips: Final = live_ips or cached.ips ranked_ips: Final = ( tuple( sorted( - cached.ips, + candidate_ips, key=lambda address: (-session_rank(address), address), ) ) if session_id is not None else () ) - ip: Final = cached.ips[cached.cursor] if session_id is None else ranked_ips[retry_count % pod_count] - next_cursor: Final = (cached.cursor + 1) % len(cached.ips) if session_id is None else cached.cursor + ip: Final = ( + candidate_ips[cached.cursor % len(candidate_ips)] + if session_id is None + else ranked_ips[retry_count % len(candidate_ips)] + ) + next_cursor: Final = (cached.cursor + 1) % len(candidate_ips) if session_id is None else cached.cursor selection: Final = ( "round_robin" if session_id is None diff --git a/tests/code_coverage_tests/router_surface_allowlist.json b/tests/code_coverage_tests/router_surface_allowlist.json index 52e5bb43625..0d83ee4326e 100644 --- a/tests/code_coverage_tests/router_surface_allowlist.json +++ b/tests/code_coverage_tests/router_surface_allowlist.json @@ -180,6 +180,7 @@ "get_model_info", "get_model_list", "get_router_model_info", + "kubernetes_pod_discovery", "log_retry", "model_list", "provider_default_deployment_ids", diff --git a/tests/unit/router_utils/test_kubernetes_pod_discovery.py b/tests/unit/router_utils/test_kubernetes_pod_discovery.py index 115e2a49bc5..0e000e2c3fc 100644 --- a/tests/unit/router_utils/test_kubernetes_pod_discovery.py +++ b/tests/unit/router_utils/test_kubernetes_pod_discovery.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import copy import hashlib import json @@ -12,6 +13,7 @@ from itertools import chain, count, repeat from typing import Final, TypeAlias, cast import httpx +import openai import pytest import respx from openai import AsyncOpenAI, OpenAI @@ -19,11 +21,14 @@ from openai import AsyncOpenAI, OpenAI import litellm from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, + DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS, KUBERNETES_POD_ROUTING_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) from litellm.router import CustomRoutingStrategyBase, Router +from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments from litellm.router_utils.kubernetes_pod_discovery import ( + KubernetesPod, KubernetesPodDiscovery, async_resolve_pods_after, async_resolve_pods_after_bound, @@ -2262,3 +2267,280 @@ def test_unresolved_deployment_does_not_record_pod_selection( assert resolved is deployment assert KUBERNETES_POD_ROUTING_KEY not in resolved + + +_PLAIN_URL: Final = "http://plain-vllm.ns.svc.cluster.local:8000/v1" +_SERVICE_HOST: Final = "vllm-headless.ns.svc.cluster.local" + + +class _ManualClock: + def __init__(self) -> None: + self.seconds = 0.0 + + def __call__(self) -> float: + return self.seconds + + def advance(self, seconds: float) -> None: + self.seconds += seconds + + +def _stub_async_three_pod_dns(monkeypatch: pytest.MonkeyPatch) -> None: + async def getaddrinfo( + host: str | None, + port: str | int | None, + family: int = 0, + type: int = 0, + proto: int = 0, + flags: int = 0, + ) -> list[_AddrInfo]: + return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0)) + + monkeypatch.setattr(asyncio.get_running_loop(), "getaddrinfo", getaddrinfo) + + +def _failing_pods(status_code: int, *pod_ips: str) -> Callable[[httpx.Request], httpx.Response]: + def response_for(request: httpx.Request) -> httpx.Response: + if request.url.host not in pod_ips: + return httpx.Response(200, json=_CHAT_RESPONSE, request=request) + return httpx.Response(status_code, json={"error": {"message": "pod failed", "type": "error"}}, request=request) + + return response_for + + +def _ordered_deployment(deployment: dict[str, object], deployment_id: str, order: int) -> dict[str, object]: + params: Final = deployment["litellm_params"] + assert isinstance(params, Mapping) + return {**deployment, "litellm_params": {**params, "order": order}, "model_info": {"id": deployment_id}} + + +def _pods_and_plain_router() -> Router: + return Router( + model_list=[ + _ordered_deployment(_deployment(), "pods", 1), + _ordered_deployment(_deployment(_PLAIN_URL, kubernetes_pod_discovery=None), "plain", 2), + ], + num_retries=0, + ) + + +async def _send_chats(router: Router, model: str, requests: int) -> None: + for _ in range(requests): + with contextlib.suppress(openai.APIStatusError): + await router.acompletion(model=model, messages=[{"role": "user", "content": "hello"}]) + + +def _hosts(route: respx.Route) -> tuple[str, ...]: + return tuple(call.request.url.host for call in route.calls) + + +@pytest.mark.asyncio +async def test_pod_429_takes_only_that_pod_out_and_keeps_its_deployment_serving( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _stub_async_three_pod_dns(monkeypatch) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=False) as respx_mock: + pod_route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( + side_effect=_failing_pods(429, "10.0.0.3") + ) + respx_mock.post(f"{_PLAIN_URL}/chat/completions").mock(return_value=httpx.Response(200, json=_CHAT_RESPONSE)) + _patch_router_discovery( + monkeypatch, KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) + ) + router: Final = _pods_and_plain_router() + + await _send_chats(router, "gpu-model", 9) + cooled_deployments: Final = await async_get_cooldown_deployments(router, None) + + assert _hosts(pod_route) == ( + "10.0.0.1", + "10.0.0.2", + "10.0.0.3", + "10.0.0.1", + "10.0.0.2", + "10.0.0.1", + "10.0.0.2", + "10.0.0.1", + "10.0.0.2", + ) + assert cooled_deployments == [] + + +@pytest.mark.asyncio +async def test_deployment_cools_down_once_its_last_live_pod_fails(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_async_three_pod_dns(monkeypatch) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=False) as respx_mock: + pod_route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( + side_effect=_failing_pods(429, "10.0.0.1", "10.0.0.2", "10.0.0.3") + ) + respx_mock.post(f"{_PLAIN_URL}/chat/completions").mock(return_value=httpx.Response(200, json=_CHAT_RESPONSE)) + _patch_router_discovery( + monkeypatch, KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) + ) + router: Final = _pods_and_plain_router() + + await _send_chats(router, "gpu-model", 5) + cooled_deployments: Final = await async_get_cooldown_deployments(router, None) + + assert tuple(sorted(_hosts(pod_route))) == ("10.0.0.1", "10.0.0.2", "10.0.0.3") + assert cooled_deployments == ["pods"] + + +@pytest.mark.asyncio +async def test_pod_429_in_a_single_deployment_group_takes_that_pod_out(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_async_three_pod_dns(monkeypatch) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=True) as respx_mock: + pod_route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( + side_effect=_failing_pods(429, "10.0.0.3") + ) + _patch_router_discovery( + monkeypatch, KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) + ) + router: Final = Router(model_list=[_deployment()], num_retries=0) + + await _send_chats(router, "gpu-model", 7) + + assert _hosts(pod_route) == ("10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.0.1", "10.0.0.2", "10.0.0.1", "10.0.0.2") + + +@pytest.mark.asyncio +async def test_pod_failing_most_of_its_calls_with_500s_is_taken_out(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_async_three_pod_dns(monkeypatch) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + requests_until_threshold: Final = DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS * 3 + with respx.mock(assert_all_called=True) as respx_mock: + pod_route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( + side_effect=_failing_pods(500, "10.0.0.3") + ) + _patch_router_discovery( + monkeypatch, KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) + ) + router: Final = Router(model_list=[_deployment()], num_retries=0) + + await _send_chats(router, "gpu-model", requests_until_threshold + 4) + hosts: Final = _hosts(pod_route) + + assert hosts.count("10.0.0.3") == DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS + assert hosts[requests_until_threshold:] == ("10.0.0.1", "10.0.0.2", "10.0.0.1", "10.0.0.2") + + +@pytest.mark.asyncio +async def test_pod_failing_half_its_calls_stays_in_rotation(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_async_three_pod_dns(monkeypatch) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + third_pod_calls: Final = count() + third_pod_requests: Final = DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS * 3 + + def response_for(request: httpx.Request) -> httpx.Response: + third_pod_failure: Final = request.url.host == "10.0.0.3" and next(third_pod_calls) % 2 == 1 + if third_pod_failure: + return httpx.Response(500, json={"error": {"message": "pod failed", "type": "error"}}, request=request) + return httpx.Response(200, json=_CHAT_RESPONSE, request=request) + + with respx.mock(assert_all_called=True) as respx_mock: + pod_route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(side_effect=response_for) + _patch_router_discovery( + monkeypatch, KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) + ) + router: Final = Router(model_list=[_deployment()], num_retries=0) + + await _send_chats(router, "gpu-model", third_pod_requests * 3) + + assert _hosts(pod_route).count("10.0.0.3") == third_pod_requests + + +@pytest.mark.asyncio +async def test_fallback_pod_429_takes_only_that_pod_out(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_async_three_pod_dns(monkeypatch) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=True) as respx_mock: + pod_route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( + side_effect=_failing_pods(429, "10.0.0.3") + ) + respx_mock.post(f"{_PLAIN_URL}/chat/completions").mock( + return_value=httpx.Response(500, json={"error": {"message": "primary failed", "type": "error"}}) + ) + _patch_router_discovery( + monkeypatch, KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) + ) + primary: Final = { + "model_name": "primary-model", + "litellm_params": {"model": "openai/primary", "api_base": _PLAIN_URL, "api_key": "fake"}, + "model_info": {"id": "primary"}, + } + router: Final = Router( + model_list=[primary, _deployment()], + fallbacks=[{"primary-model": ["gpu-model"]}], + num_retries=0, + ) + + await _send_chats(router, "primary-model", 6) + + assert _hosts(pod_route) == ("10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.0.1", "10.0.0.2", "10.0.0.1") + + +def test_cooled_pod_moves_only_its_own_sessions_until_its_cooldown_ends(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_sync_dns(monkeypatch, "10.0.0.1", "10.0.0.2", "10.0.0.3") + clock: Final = _ManualClock() + discovery: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment) + sessions: Final = tuple(f"session-{index}" for index in range(30)) + + def session_hosts() -> tuple[str, ...]: + return tuple( + httpx.URL( + _api_base(discovery.resolve_deployment(_deployment(), {"metadata": {"session_id": session}})) + ).host + for session in sessions + ) + + before: Final = session_hosts() + cooled_pod: Final = discovery.pod_for_call(_deployment(), {"api_base": "http://10.0.0.3:8000/v1"}) + assert cooled_pod is not None + discovery.cool_down_pod(cooled_pod, 5.0) + during: Final = session_hosts() + clock.advance(5.0) + after: Final = session_hosts() + + assert "10.0.0.3" in before + assert "10.0.0.3" not in during + assert all( + during_host == before_host for before_host, during_host in zip(before, during) if before_host != "10.0.0.3" + ) + assert after == before + + +def test_every_pod_cooled_keeps_rotating_over_all_pods(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_sync_dns(monkeypatch, "10.0.0.1", "10.0.0.2", "10.0.0.3") + discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) + first_host: Final = httpx.URL(_api_base(discovery.resolve_deployment(_deployment()))).host + pods: Final = tuple( + KubernetesPod(service_host=_SERVICE_HOST, port=8000, ip=ip) for ip in ("10.0.0.1", "10.0.0.2", "10.0.0.3") + ) + for pod in pods: + discovery.cool_down_pod(pod, 5.0) + + later_hosts: Final = tuple(httpx.URL(_api_base(discovery.resolve_deployment(_deployment()))).host for _ in pods) + + assert first_host == "10.0.0.1" + assert later_hosts == ("10.0.0.2", "10.0.0.3", "10.0.0.1") + assert not discovery.has_other_live_pod(pods[0]) + + +def test_pod_for_call_names_only_a_discovered_pod_of_a_discovery_deployment(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_sync_dns(monkeypatch, "10.0.0.1", "10.0.0.2") + discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) + discovery.resolve_deployment(_deployment()) + + assert discovery.pod_for_call(_deployment(), {"api_base": "http://10.0.0.2:8000/v1"}) == KubernetesPod( + service_host=_SERVICE_HOST, port=8000, ip="10.0.0.2" + ) + assert discovery.pod_for_call(_deployment(), {"api_base": _SERVICE_URL}) is None + assert discovery.pod_for_call(_deployment(), {"api_base": "http://10.0.0.9:8000/v1"}) is None + assert ( + discovery.pod_for_call(_deployment(kubernetes_pod_discovery=None), {"api_base": "http://10.0.0.2:8000/v1"}) + is None + ) + assert discovery.pod_for_call(_deployment(), {}) is None + assert discovery.pod_for_call(_deployment(), {"api_base": "http://10.0.0.2:port/v1"}) is None diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 513136ae150..29f530b7807 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.litellm_params import RoutingStrategyName from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( @@ -456,9 +457,7 @@ async def test_deployment_callback_on_success(sync_mode): ] router = Router(model_list=model_list) # Get the actual deployment ID that was generated - gpt_deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + gpt_deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") deployment_id = gpt_deployment["model_info"]["id"] standard_logging_payload = create_standard_logging_payload() @@ -494,6 +493,31 @@ async def test_deployment_callback_on_success(sync_mode): assert tpm_key is not None +def test_sync_deployment_callback_on_success_skips_a_call_the_router_did_not_route(): + router: Final = Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"}, + "model_info": {"id": "unrouted-deployment"}, + } + ] + ) + + key: Final = router.sync_deployment_callback_on_success( + kwargs={"litellm_params": {"model_info": {"id": "unrouted-deployment"}}}, + completion_response=litellm.ModelResponse(model="gpt-5-mini"), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert key is None + assert ( + get_deployment_successes_for_current_minute(litellm_router_instance=router, deployment_id="unrouted-deployment") + == 0 + ) + + @pytest.mark.asyncio async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(): import time @@ -563,9 +587,16 @@ async def _finish_35_token_websocket_turn(routed: Mapping[str, object], call_id: litellm_call_id=call_id, function_id=call_id, ) - turn.update_environment_variables(litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={}) + turn.update_environment_variables( + litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={} + ) await turn.dispatch_success_handlers( - [{"type": "response.completed", "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}}], + [ + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}, + } + ], prefer_async_handlers=True, ) @@ -968,9 +999,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") litellm_logging_obj = Logging( model="gpt-5-mini", @@ -985,9 +1014,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): router.routing_strategy_pre_call_checks(deployment) else: ## NO EXCEPTION - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) ## WITH EXCEPTION - rate limit error with patch.object( @@ -1002,27 +1029,19 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): ), ): with pytest.raises(litellm.RateLimitError): - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) ## WITH EXCEPTION - generic error - with patch.object( - callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error")) - ): + with patch.object(callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error"))): with pytest.raises(Exception, match="Error"): - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) @pytest.mark.parametrize( "set_supported_environments, supported_environments, is_supported", [(True, ["staging"], True), (False, None, True), (True, ["development"], False)], ) -def test_create_deployment( - model_list, set_supported_environments, supported_environments, is_supported -): +def test_create_deployment(model_list, set_supported_environments, supported_environments, is_supported): """Test if the '_create_deployment' function is working correctly""" router = Router(model_list=model_list) @@ -1056,20 +1075,14 @@ def test_deployment_is_active_for_environment( ): """Test if the '_deployment_is_active_for_environment' function is working correctly""" router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") if set_supported_environments: os.environ["LITELLM_ENVIRONMENT"] = "staging" deployment["model_info"]["supported_environments"] = supported_environments if is_supported: - assert ( - router.deployment_is_active_for_environment(deployment=deployment) is True - ) + assert router.deployment_is_active_for_environment(deployment=deployment) is True else: - assert ( - router.deployment_is_active_for_environment(deployment=deployment) is False - ) + assert router.deployment_is_active_for_environment(deployment=deployment) is False def test_set_model_list(model_list): @@ -1474,27 +1487,19 @@ async def test_get_model_group_io_token_usage_sums_across_deployments(): minute = get_utc_datetime().strftime("%H-%M") keys_and_values = [ ( - RouterCacheEnum.ITPM.value.format( - id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute - ), + RouterCacheEnum.ITPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), 30, ), ( - RouterCacheEnum.OTPM.value.format( - id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute - ), + RouterCacheEnum.OTPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), 10, ), ( - RouterCacheEnum.ITPM.value.format( - id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute - ), + RouterCacheEnum.ITPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), 70, ), ( - RouterCacheEnum.OTPM.value.format( - id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute - ), + RouterCacheEnum.OTPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), 20, ), ] @@ -1816,9 +1821,7 @@ def test_initialize_alerting(model_list): from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.types.router import AlertingConfig - router = Router( - model_list=model_list, alerting_config=AlertingConfig(webhook_url="test") - ) + router = Router(model_list=model_list, alerting_config=AlertingConfig(webhook_url="test")) router._initialize_alerting() callback_added = False @@ -2093,7 +2096,6 @@ async def test_discard_releases_per_request_override_selector(): assert selector_ref() is None - def test_discard_registers_the_selector_of_a_same_strategy_survivor_built_without_a_model_list(): retired: Final = Router(model_list=[], routing_strategy="latency-based-routing") survivor: Final = Router(routing_strategy="latency-based-routing") @@ -2171,6 +2173,7 @@ async def test_discard_releases_its_router_budget_limiter(): assert budget_limiter_ref() is None + def test_initialize_assistants_endpoint(model_list): """Test if the 'initialize_assistants_endpoint' function is working correctly""" router = Router(model_list=model_list) @@ -3188,6 +3191,7 @@ def test_resolve_model_name_from_model_id_keeps_a_wildcard_name_only_for_the_ids router = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-fake"}}]) assert router.resolve_model_name_from_model_id(model_id, custom_llm_provider) == expected + def test_get_valid_args(): """Test get_valid_args static method returns valid Router.__init__ arguments"""