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"""