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.
This commit is contained in:
mateo-berri 2026-10-29 11:18:51 -07:00
parent ebf84f8649
commit 5e12c67aae
7 changed files with 520 additions and 96 deletions

View file

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

View file

@ -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,

View file

@ -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)

View file

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

View file

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

View file

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

View file

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