diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 8bab7215b17..0827199027d 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -38,7 +38,7 @@ else: LitellmRouter = Any -def _route_user_config_request(data: dict, route_type: str): +async def _route_user_config_request(data: dict, route_type: str): """Route a request using the user-provided router config.""" router_config: Final = data.pop("user_config") @@ -48,9 +48,10 @@ def _route_user_config_request(data: dict, route_type: str): filtered_config: Final = {k: v for k, v in router_config.items() if k in valid_args} user_router: Final = litellm.Router(**filtered_config) - ret_val: Final = getattr(user_router, f"{route_type}")(**data) - user_router.discard() - return ret_val + try: + return await getattr(user_router, f"{route_type}")(**data) + finally: + user_router.discard() def _is_a2a_agent_model(model_name: object) -> bool: diff --git a/litellm/router.py b/litellm/router.py index 725c829b6cb..ece559cb90e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -38,7 +38,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial from itertools import chain -from types import MappingProxyType +from types import MappingProxyType, MethodType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast import anyio @@ -516,6 +516,24 @@ def model_info_is_active_for_environment(model_info: Mapping[str, object] | None _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _CallbackT = TypeVar("_CallbackT") + +def _drop_from_global_callback_lists(callbacks: Iterable[object], method_owner: object | None = None) -> None: + dropped_ids: Final = frozenset(id(callback) for callback in callbacks if callback is not None) + + def keep(callback: object) -> bool: + if id(callback) in dropped_ids: + return False + return method_owner is None or not (isinstance(callback, MethodType) and callback.__self__ is method_owner) + + litellm.callbacks = [callback for callback in litellm.callbacks if keep(callback)] + litellm.input_callback = [callback for callback in litellm.input_callback if keep(callback)] + litellm.success_callback = [callback for callback in litellm.success_callback if keep(callback)] + litellm.failure_callback = [callback for callback in litellm.failure_callback if keep(callback)] + litellm.service_callback = [callback for callback in litellm.service_callback if keep(callback)] + litellm._async_success_callback = [callback for callback in litellm._async_success_callback if keep(callback)] + litellm._async_failure_callback = [callback for callback in litellm._async_failure_callback if keep(callback)] + + _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( {"model", "api_base", "api_key", "api_version", "responses_api_version"} ) @@ -1527,20 +1545,12 @@ class Router: self._override_selectors_lock = threading.Lock() self.access_groups = None ## USAGE TRACKING ## - if isinstance(litellm._async_success_callback, list): - litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) - else: - litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) + litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) if isinstance(litellm.success_callback, list): litellm.logging_callback_manager.add_litellm_success_callback(self.sync_deployment_callback_on_success) else: litellm.success_callback = [self.sync_deployment_callback_on_success] - if isinstance(litellm._async_failure_callback, list): - litellm.logging_callback_manager.add_litellm_async_failure_callback( - self.async_deployment_callback_on_failure - ) - else: - litellm._async_failure_callback = [self.async_deployment_callback_on_failure] + litellm.logging_callback_manager.add_litellm_async_failure_callback(self.async_deployment_callback_on_failure) ## COOLDOWNS ## if isinstance(litellm.failure_callback, list): litellm.logging_callback_manager.add_litellm_failure_callback(self.deployment_callback_on_failure) @@ -1636,27 +1646,10 @@ class Router: # Stop contributing to cost-map rebuilds straight away rather than waiting # for this router to be collected. _live_routers.discard(self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_success_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.success_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_failure_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.failure_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.input_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.service_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, self) - self._unregister_router_selectors( - ( - *(getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()), - *self._override_selectors.values(), - *chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()), - ) - ) - - # Remove ForwardClientSideHeadersByModelGroup if it exists - if self.optional_callbacks is not None: - for callback in self.optional_callbacks: - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.callbacks, callback, require_self=False - ) + self._unregister_router_selectors((*self._callback_selectors(), *self._override_selectors.values())) + _drop_from_global_callback_lists(self.optional_callbacks or (), method_owner=self) + for router in tuple(_live_routers): + router._register_global_callbacks() @staticmethod def _create_redis_cache( @@ -1771,6 +1764,22 @@ class Router: if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(selector) + def _callback_selectors(self) -> tuple[RouterStrategySelector, ...]: + return tuple( + selector + for selector in ( + *(getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()), + *chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()), + ) + if selector is not None + ) + + def _register_global_callbacks(self) -> None: + for selector in self._callback_selectors(): + self._register_router_selector(selector) + for callback in self.optional_callbacks or (): + litellm.logging_callback_manager.add_litellm_callback(callback) + def _unregister_router_selectors(self, selectors: Sequence[object]) -> None: """ Drop router-owned strategy selectors from litellm's global callback @@ -1780,11 +1789,9 @@ class Router: into the per-event lists on every request, so those are purged too. """ for selector in selectors: - if selector is None: - continue if isinstance(selector, BaseRoutingStrategy): selector.retire() - litellm.logging_callback_manager.remove_callback_from_all_lists(selector) + _drop_from_global_callback_lists(selectors) def _apply_updated_routing_strategy_args(self) -> None: """ diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 4ea242258f6..b65968f4bd4 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -139,6 +139,9 @@ ignored_function_names = [ "async_admit_deployment_turn", # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py "_deployment_as_configured_now", # Read by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py "_admit_turn_through_routing_strategy", # Run by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py + "_drop_from_global_callback_lists", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "_callback_selectors", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "_register_global_callbacks", # Run by discard, tested through the survivor discard tests in test_router_helper_utils.py ] diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 4e4e2f23950..418fa669f16 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,16 +1,20 @@ +import asyncio +import threading import httpx import litellm import pytest import respx +from collections.abc import Iterator from types import SimpleNamespace from typing import Final from unittest.mock import MagicMock +from fakeredis import TcpFakeServer from fastapi import HTTPException -from litellm.proxy.route_llm_request import ProxyModelNotFoundError, route_request +from litellm.proxy.route_llm_request import ProxyModelNotFoundError, _route_user_config_request, route_request @pytest.mark.parametrize( @@ -1982,3 +1986,52 @@ async def test_team_key_search_of_registered_bedrock_store_reaches_bedrock_with_ posted_url ) assert "Credential=AKIAREGISTRYKEY00000/" in signed_auth, signed_auth + + +@pytest.fixture +def fake_redis_port() -> Iterator[int]: + server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis") + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield server.server_address[1] + server.shutdown() + server.server_close() + + +async def _user_config_request_outcome(redis_port: int) -> str: + data: Final = { + "model": "rate-limited", + "messages": [{"role": "user", "content": "hi"}], + "num_retries": 0, + "user_config": { + "routing_strategy": "usage-based-routing-v2", + "redis_host": "127.0.0.1", + "redis_port": redis_port, + "model_list": [ + { + "model_name": "rate-limited", + "rpm": 2, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "mock_response": "ok"}, + } + ], + }, + } + try: + await _route_user_config_request(data, "acompletion") + except litellm.RateLimitError: + return "429" + finally: + await asyncio.sleep(0.3) + return "200" + + +@pytest.mark.asyncio +async def test_user_config_router_enforces_its_rpm_limit_across_requests(fake_redis_port: int) -> None: + outcomes: Final = [await _user_config_request_outcome(fake_redis_port) for _ in range(5)] + + assert outcomes[0] == "200" + assert "429" in outcomes, outcomes + assert all(getattr(litellm, list_name) == [] for list_name in ("callbacks", "input_callback")), ( + litellm.callbacks, + litellm.input_callback, + ) diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 91f719599cb..03ea18dc1b5 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -4,7 +4,8 @@ import json import os import weakref from datetime import datetime, timezone -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -1839,6 +1840,32 @@ def test_discard_keeps_callbacks_owned_by_other_live_routers(): survivor.discard() +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_leaves_same_strategy_survivor_receiving_events( + routing_strategy: RoutingStrategyName, selector_attr: str +): + model_list: Final[list[DeploymentTypedDict]] = [ + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} + ] + retired: Final = Router(model_list=model_list, routing_strategy=routing_strategy) + survivor: Final = Router(model_list=model_list, routing_strategy=routing_strategy) + + retired.discard() + await survivor.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + + assert any(callback is getattr(survivor, selector_attr) for callback in litellm.success_callback) + survivor.discard() + + def test_discard_drops_routing_group_selectors(): router: Final = Router( model_list=[ @@ -1856,6 +1883,109 @@ def test_discard_drops_routing_group_selectors(): assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) +_GROUPED_MODEL_LIST: Final[list[DeploymentTypedDict]] = [ + {"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} +] + + +class _LabelledHook(CustomLogger): + def __init__(self, label: str) -> None: + super().__init__() + self.label = label + + +def _visit_hooks_reloading_routing_groups_mid_loop( + router: Router, hooks: tuple[_LabelledHook, ...] +) -> Iterator[CustomLogger]: + for callback in litellm.callbacks: + if callback is hooks[0]: + router.update_settings( + routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}] + ) + if callback in hooks: + yield callback + + +def test_routing_group_reload_keeps_every_hook_of_an_in_flight_callback_loop(): + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}], + ) + hooks: Final = tuple(_LabelledHook(label) for label in ("guardrail_a", "guardrail_b", "guardrail_c")) + for hook in hooks: + litellm.logging_callback_manager.add_litellm_callback(hook) + + visited: Final = tuple(_visit_hooks_reloading_routing_groups_mid_loop(router, hooks)) + + assert visited == hooks + router.discard() + + +def _routers_seen_in_success_loop_discarding_first(first: Router, second: Router) -> Iterator[Router]: + for callback in litellm.success_callback: + if not isinstance(callback, MethodType): + continue + if callback.__self__ is first: + first.discard() + yield first + if callback.__self__ is second: + yield second + + +def test_discard_mid_success_loop_keeps_the_next_router_callback(): + first: Final = Router(model_list=_GROUPED_MODEL_LIST) + second: Final = Router(model_list=_GROUPED_MODEL_LIST) + + seen: Final = tuple(_routers_seen_in_success_loop_discarding_first(first, second)) + + assert seen == (first, second) + second.discard() + + +@pytest.mark.asyncio +async def test_discard_purges_optional_pre_call_checks_from_event_lists_after_serving_a_request(): + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + optional_pre_call_checks=["enforce_model_rate_limits", "prompt_caching", "deployment_affinity"], + ) + await router.acompletion(model="grouped-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def test_discard_registers_same_config_survivor_optional_pre_call_checks(): + retired: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"]) + survivor: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"]) + + retired.discard() + + assert survivor.optional_callbacks + assert all(_is_registered(check) for check in survivor.optional_callbacks) + survivor.discard() + + +def _is_registered(callback: object) -> bool: + return any(registered is callback for registered in litellm.callbacks) + + +@pytest.mark.asyncio +async def test_discard_releases_per_request_override_selector(): + def build_override_then_discard() -> weakref.ReferenceType[object]: + router: Final = Router(model_list=_GROUPED_MODEL_LIST) + selector: Final = router._get_override_strategy_selector("usage-based-routing-v2") + assert selector is not None + router.discard() + return weakref.ref(selector) + + selector_ref: Final = build_override_then_discard() + await asyncio.sleep(0) + gc.collect() + + assert selector_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)