scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(router): bind the deployment callbacks to user_config requests

Browse source 
A cached user_config router serves requests through each request's logging
object but bound only its selectors and policies there, so a failing
deployment was never cooled down and its usage counters missed requests.
This commit is contained in:
mateo-berri 2026-10-20 07:20:42 -07:00
parent 9bb64d2a5d
commit cd7b358d28
5 changed files with 200 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
litellm/litellm_core_utils/litellm_logging.py
View file

@ -647,11 +647,11 @@ class Logging(LiteLLMLoggingBaseClass):
litellm_call_id: str,
function_id: str,
litellm_trace_id: str | None = None,
dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = None,
dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None,
dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = None,
dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
applied_guardrails: list[str] | None = None,
kwargs: dict | None = None,
log_raw_request_response: bool = False,
@ -724,15 +724,17 @@ class Logging(LiteLLMLoggingBaseClass):
self.raw_request_only = raw_request_only
# Initialize dynamic callbacks
self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks
self.dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = dynamic_input_callbacks
self.dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_success_callbacks
)
self.dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_async_success_callbacks
)
self.dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = dynamic_failure_callbacks
self.dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = (
self.dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_failure_callbacks
)
self.dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_async_failure_callbacks
)
@ -854,10 +856,26 @@ class Logging(LiteLLMLoggingBaseClass):
litellm.logging_callback_manager.get_custom_logger_key(callback)
}
def add_dynamic_function_callbacks(
self,
*,
async_success: Callable[..., object],
failure: Callable[..., object],
async_failure: Callable[..., object],
) -> None:
self.dynamic_async_success_callbacks = self._with_dynamic_callback(
self.dynamic_async_success_callbacks, async_success
)
self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, failure)
self.dynamic_async_failure_callbacks = self._with_dynamic_callback(
self.dynamic_async_failure_callbacks, async_failure
)
@staticmethod
def _with_dynamic_callback(
callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger
) -> list[str | Callable | CustomLogger]:
callbacks: Sequence[str | Callable[..., object] | CustomLogger] | None,
callback: Callable[..., object] | CustomLogger,
) -> list[str | Callable[..., object] | CustomLogger]:
existing: Final = tuple(callbacks or ())
return [*existing, *(() if callback in existing else (callback,))]
@ -894,9 +912,9 @@ class Logging(LiteLLMLoggingBaseClass):
def _process_dynamic_callback_list(
self,
callback_list: list[str | Callable | CustomLogger] | None,
callback_list: list[str | Callable[..., object] | CustomLogger] | None,
dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"],
) -> list[str | Callable | CustomLogger] | None:
) -> list[str | Callable[..., object] | CustomLogger] | None:
"""
Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks
@ -908,7 +926,7 @@ class Logging(LiteLLMLoggingBaseClass):
if callback_list is None:
return None
processed_list: Final[list[str | Callable | CustomLogger]] = []
processed_list: Final[list[str | Callable[..., object] | CustomLogger]] = []
for callback in callback_list:
if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks:
for callback_instance in self._resolve_dynamic_callback_string(callback):

4
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -880,8 +880,8 @@ def _build_passthrough_failure_request_payload(
@dataclass(frozen=True, slots=True)
class _TeamCallbackWiring:
success_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg
failure_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg
success_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg
failure_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg
logging_kwargs: dict[str, str | dict[str, str]] | None = None # mutable-ok: Logging.__init__ arg

5
litellm/router.py
View file

@ -1709,6 +1709,11 @@ class Router:
self._release_global_callbacks()
def bind_to_request_logging(self, logging_obj: LiteLLMLogging) -> None:
logging_obj.add_dynamic_function_callbacks(
async_success=self.deployment_callback_on_success,
failure=self.deployment_callback_on_failure,
async_failure=self.async_deployment_callback_on_failure,
)
for selector in self._request_scoped_selectors:
logging_obj.replace_global_callback(selector)
for policy in self._request_policies_not_set_globally():

8
litellm/utils.py
View file

@ -1197,10 +1197,10 @@ def function_setup(
for index in reversed(removed_async_items):
litellm.failure_callback.pop(index)
### DYNAMIC CALLBACKS ###
dynamic_success_callbacks: list[str | Callable | CustomLogger] | None = None
dynamic_async_success_callbacks: list[str | Callable | CustomLogger] | None = None
dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None
dynamic_async_failure_callbacks: Final[list[str | Callable | CustomLogger] | None] = None
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | CustomLogger] | None] = None
if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list):
removed_async_items = []
for index, callback in enumerate(kwargs["success_callback"]):

161
tests/unit/proxy/test_route_llm_request.py
View file

@ -11,7 +11,7 @@ import redis
import respx
from collections.abc import Coroutine, Iterator, Mapping
from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping
from datetime import datetime
from types import SimpleNamespace
from uuid import uuid4
@ -2764,6 +2764,165 @@ async def test_a_request_in_flight_when_its_user_config_router_is_dropped_still_
await request()
async def _user_config_request(model: str, user_config: Mapping[str, object], arrives_with_logging_obj: bool) -> object:
return await (
await route_request(
{
"model": model,
"messages": list(_HELLO_MESSAGES),
"num_retries": 0,
**({"litellm_logging_obj": _proxy_logging_obj(model)} if arrives_with_logging_obj else {}),
"user_config": user_config,
},
None,
None,
"acompletion",
)
)
def _ordered_pair(model_name: str, first_base: str, second_base: str) -> list[dict[str, object]]:
return [
{
"model_name": model_name,
"litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "api_base": api_base, "order": order},
}
for order, api_base in ((1, first_base), (2, second_base))
]
@pytest.mark.parametrize("arrives_with_logging_obj", [True, False])
@pytest.mark.asyncio
async def test_a_rate_limited_user_config_deployment_cools_down_for_the_next_requests(
monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
rate_limited_base: Final = "https://rate-limited.example/v1"
healthy_base: Final = "https://healthy.example/v1"
user_config: Final = {"model_list": _ordered_pair("flaky", rate_limited_base, healthy_base)}
with respx.mock(assert_all_called=False) as respx_mock:
rate_limited: Final = respx_mock.post(f"{rate_limited_base}/chat/completions").mock(
return_value=httpx.Response(429, json={"error": {"message": "slow down"}})
)
respx_mock.post(f"{healthy_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5"))
for _ in range(3):
await _user_config_request("flaky", user_config, arrives_with_logging_obj)
await asyncio.sleep(0.3)
assert rate_limited.call_count == 1
async def _remaining_usage_headers_after_429_then_200s(
send: Callable[[], Awaitable[object]], api_base: str
) -> list[tuple[object, object] | str]:
async def remaining_usage() -> tuple[object, object] | str:
try:
response: Final = await send()
except litellm.RateLimitError:
return "429"
finally:
await asyncio.sleep(0.3)
assert isinstance(response, litellm.ModelResponse)
headers: Final = response._hidden_params["additional_headers"]
return headers.get("x-ratelimit-remaining-requests"), headers.get("x-ratelimit-remaining-tokens")
with respx.mock(assert_all_called=False) as respx_mock:
respx_mock.post(f"{api_base}/chat/completions").mock(
side_effect=[
httpx.Response(429, json={"error": {"message": "slow down"}}),
_chat_completion_response("gpt-5.5"),
_chat_completion_response("gpt-5.5"),
]
)
return [await remaining_usage() for _ in range(3)]
@pytest.mark.parametrize("arrives_with_logging_obj", [True, False])
@pytest.mark.asyncio
async def test_user_config_remaining_usage_headers_match_a_proxy_router_with_the_same_deployment(
monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
await _wait_out_a_rate_limit_minute_rollover()
proxy_base: Final = "https://proxy-limited.example/v1"
user_config_base: Final = "https://user-config-limited.example/v1"
def deployment(model_name: str, api_base: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "openai/gpt-5.5",
"api_key": "sk-test",
"api_base": api_base,
"rpm": 10,
"tpm": 1000,
},
}
proxy_router: Final = litellm.Router(model_list=[deployment("proxy-limited", proxy_base)])
try:
proxy_router_headers: Final = await _remaining_usage_headers_after_429_then_200s(
lambda: proxy_router.acompletion(model="proxy-limited", messages=list(_HELLO_MESSAGES), num_retries=0),
proxy_base,
)
finally:
proxy_router.discard()
user_config: Final = {"model_list": [deployment("user-config-limited", user_config_base)]}
user_config_headers: Final = await _remaining_usage_headers_after_429_then_200s(
lambda: _user_config_request("user-config-limited", user_config, arrives_with_logging_obj), user_config_base
)
assert user_config_headers == proxy_router_headers
async def _first_deployment_calls_over_successes_then_failures(
send: Callable[[], Awaitable[object]], first_base: str, second_base: str
) -> int:
successes: Final = [_chat_completion_response("gpt-5.5") for _ in range(6)]
failures: Final = [httpx.Response(500, json={"error": {"message": "boom"}}) for _ in range(6)]
with respx.mock(assert_all_called=False) as respx_mock:
first: Final = respx_mock.post(f"{first_base}/chat/completions").mock(side_effect=[*successes, *failures])
respx_mock.post(f"{second_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5"))
for _ in range(len(successes) + len(failures)):
await send()
await asyncio.sleep(0.3)
return first.call_count
@pytest.mark.parametrize("arrives_with_logging_obj", [True, False])
@pytest.mark.asyncio
async def test_user_config_successes_keep_a_deployment_out_of_cooldown_like_a_proxy_router(
monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
await _wait_out_a_rate_limit_minute_rollover()
proxy_router: Final = litellm.Router(
model_list=_ordered_pair("proxy-pair", "https://proxy-first.example/v1", "https://proxy-second.example/v1")
)
try:
proxy_router_calls: Final = await _first_deployment_calls_over_successes_then_failures(
lambda: proxy_router.acompletion(model="proxy-pair", messages=list(_HELLO_MESSAGES), num_retries=0),
"https://proxy-first.example/v1",
"https://proxy-second.example/v1",
)
finally:
proxy_router.discard()
user_config: Final = {
"model_list": _ordered_pair("user-pair", "https://user-first.example/v1", "https://user-second.example/v1")
}
user_config_calls: Final = await _first_deployment_calls_over_successes_then_failures(
lambda: _user_config_request("user-pair", user_config, arrives_with_logging_obj),
"https://user-first.example/v1",
"https://user-second.example/v1",
)
assert user_config_calls == proxy_router_calls
@pytest.mark.asyncio
async def test_user_config_adaptive_router_learns_from_proxy_requests() -> None:
user_config: Final = {