From c8ed5ccca8500ccc0c8a3b499790c07e7e57a763 Mon Sep 17 00:00:00 2001 From: Kazuki Hashimoto <113189196+setouchi-h@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:39:36 +0900 Subject: [PATCH] fix(router): unregister strategy selectors from global callbacks on discard (#45795) * fix(router): unregister strategy selectors from global callbacks on discard Router.discard only removed its own bound methods from litellm's callback lists, so the strategy selectors it registered (least-busy, usage-based, latency-based, cost-based and routing group selectors) stayed behind and kept receiving events after the router was gone. Reuse the identity based unregister that update_settings already uses so a discarded router takes only its own selectors with it * fix(router): purge discarded selectors from every event callback list function_setup copies litellm.callbacks into the success and failure lists on every request, so a router that served traffic still leaked its selectors there after discard. Route the cleanup through remove_callback_from_all_lists, which now also covers input_callback, and fold the selector gathering into discard so the router coverage check sees no untested helper * fix(router): keep live routers, in-flight hooks and user_config limits intact on discard discard() now rebuilds every global callback list without the router's selectors, optional pre-call checks and own usage callbacks instead of removing them in place, so a callback loop already running never skips a hook. It then re-registers what the remaining live routers need, so a same-strategy survivor keeps its deduped handler. The proxy's per-request user_config router is discarded only after its call returns. * fix(router): bind user_config handlers to the request and tighten discard When a user_config request carries a logging object, its router is discarded before the request runs and its handlers ride on that logging object, so v1 rpm and v2 tpm limits count every request and concurrent requests count once. Without one, the router is discarded after the request as before. The router's pre-call loops also read the request-scoped handlers. discard() now refills only the handler keys it vacated, guards the callback lists and the live-router set with a lock, covers routers built without a model_list, and drops its Slack alerting, adaptive router hook and budget limiter sync task * fix(proxy): reuse one user_config router per config and bind its handlers to each request Every user_config request now goes through a cached router built once per config. Its strategy handlers leave the global callback lists and are bound to the request's logging object instead, so limits count each request once and nothing is left behind after the request. * fix(router): keep the proxy budget limiter on user_config requests A user_config router now swaps in only its strategy selectors, so the proxy's RouterBudgetLimiting stays on the request and keeps counting its spend. Usage logged after a strategy is retired, such as a request still in flight when its user_config router is evicted, now schedules one more push to Redis. * fix(proxy): init the retire fields on the key model budget limiter The key model budget limiter builds its own state instead of calling the base init, so with the proxy's Redis attached every spend increment raised on the missing retire flag and the key's model budget never counted spend. * fix(router): bind the deployment callbacks to user_config requests 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. * fix(router): retry an evicted router's budget writes and keep the adaptive header on user_config requests A retired budget limiter now keeps retrying its queued Redis increments every sync interval until Redis accepts them, instead of giving up after one failed write. The x-litellm-adaptive-router-model header now comes from one stateless hook in litellm.callbacks that holds no router, so user_config adaptive routers, whose per-router hooks are request-scoped, still answer with the header. * fix(router): lock request callback setup against router cleanup and drain in-flight budget writes on retire function_setup now edits the global callback lists under the same lock that router cleanup uses, so a router dropping its callbacks on another thread can no longer shift the indexes function_setup pops or put a dropped handler back. A budget limiter retired while a Redis write is still in flight now keeps retrying that batch if the write fails, instead of leaving the spend queued with no task left to send it. * fix(utils): set up loggers after releasing the callback lists lock function_setup held the callback lists lock while set_callbacks ran its one-time logger setup, which can block on a Slack package install, so other threads and Router.discard waited on it. The lock now covers only the list reads and writes. * test(router): annotate the retire test's locals as Final * fix(logging): skip replaced global callbacks on the input path too replace_global_callback already kept a same-key global logger out of success and failure dispatch. pre_call and post_call still sent it the input events, so a logger with an input hook could count a request start it never saw finish. Both now share one filter with the success path * test(logging): type the input-path regression test --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 73 ++- .../logging_callback_manager.py | 149 +++-- .../proxy/hooks/model_max_budget_limiter.py | 2 + .../pass_through_endpoints.py | 4 +- litellm/proxy/route_llm_request.py | 59 +- litellm/router.py | 178 ++++-- .../adaptive_router/__init__.py | 7 +- .../router_strategy/adaptive_router/hooks.py | 14 +- .../router_strategy/base_routing_strategy.py | 20 +- litellm/router_strategy/budget_limiter.py | 30 +- litellm/utils.py | 155 ++--- .../router_code_coverage.py | 12 + tests/unit/conftest.py | 4 + .../test_litellm_logging.py | 26 + .../test_logging_callback_manager.py | 2 + tests/unit/proxy/test_route_llm_request.py | 595 +++++++++++++++++- .../adaptive_router/test_hooks.py | 13 +- .../test_base_routing_strategy.py | 28 +- .../test_budget_limiter_hotpath.py | 57 ++ .../test_router/test_router_helper_utils.py | 321 +++++++++- tests/unit/test_utils.py | 108 ++++ 21 files changed, 1656 insertions(+), 201 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 45349bb0e5a..eb2e7b8c240 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -596,6 +596,7 @@ class Logging(LiteLLMLoggingBaseClass): prometheusLogger, \ slack_app custom_pricing: bool = False + replaced_global_callback_keys: frozenset[str] = frozenset() stream_options = None litellm_request_debug: bool = False streamed_anthropic_message_id: str | None = None @@ -650,11 +651,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, @@ -728,15 +729,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 ) @@ -853,10 +856,32 @@ class Logging(LiteLLMLoggingBaseClass): self.dynamic_async_failure_callbacks, callback ) + def replace_global_callback(self, callback: CustomLogger) -> None: + self.add_dynamic_callback(callback) + self.replaced_global_callback_keys = self.replaced_global_callback_keys | { + 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,))] @@ -893,9 +918,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 @@ -907,7 +932,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): @@ -1568,7 +1593,10 @@ class Logging(LiteLLMLoggingBaseClass): "api_call_start_time" ] # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) + callbacks: Final = ( + *self._without_replaced_global_callbacks(litellm.input_callback), + *(self.dynamic_input_callbacks or ()), + ) for callback in callbacks: try: if callback == "supabase" and supabaseClient is not None: @@ -1741,7 +1769,10 @@ class Logging(LiteLLMLoggingBaseClass): ) # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) + callbacks: Final = ( + *self._without_replaced_global_callbacks(litellm.input_callback), + *(self.dynamic_input_callbacks or ()), + ) for callback in callbacks: try: if callback == "sentry" and add_breadcrumb: @@ -4518,14 +4549,26 @@ class Logging(LiteLLMLoggingBaseClass): _filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks) return len(_filtered_failure_callbacks) > 0 + def _without_replaced_global_callbacks( + self, global_callbacks: Iterable[str | Callable[..., object] | CustomLogger] + ) -> tuple[str | Callable[..., object] | CustomLogger, ...]: + if not self.replaced_global_callback_keys: + return tuple(global_callbacks) + return tuple( + litellm.logging_callback_manager.without_custom_logger_keys( + global_callbacks, self.replaced_global_callback_keys + ) + ) + def get_combined_callback_list( self, dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None, global_callbacks: list[str | Callable[..., object] | CustomLogger], ) -> list[str | Callable[..., object] | CustomLogger]: + kept_global_callbacks: Final = self._without_replaced_global_callbacks(global_callbacks) if dynamic_success_callbacks is None: - return list(global_callbacks) - return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks)) + return list(kept_global_callbacks) + return list(dict.fromkeys((*dynamic_success_callbacks, *kept_global_callbacks))) def _remove_internal_litellm_callbacks(self, callbacks: list) -> list: """ diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 9f521ffc6cc..307353d081b 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -1,4 +1,6 @@ -from collections.abc import Callable +import threading +from collections.abc import Callable, Iterable, Iterator +from types import MethodType from typing import TYPE_CHECKING, Final import litellm @@ -15,6 +17,7 @@ else: _custom_logger_compatible_callbacks_literal = str _generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {} +_callback_lists_lock: Final = threading.RLock() class LoggingCallbackManager: @@ -43,16 +46,18 @@ class LoggingCallbackManager: Add a input callback to litellm.input_callback. Auto-routes async callbacks to litellm._async_input_callback. """ - if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) + with _callback_lists_lock: + if not isinstance(callback, str) and self._is_async_callable(callback): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) def add_litellm_service_callback(self, callback: CustomLogger | str | Callable[..., object]): """ Add a service callback to litellm.service_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback) def add_litellm_callback(self, callback: CustomLogger | str | Callable[..., object]): """ @@ -60,10 +65,11 @@ class LoggingCallbackManager: Ensures no duplicates are added. """ - self._safe_add_callback_to_list( - callback=callback, - parent_list=litellm.callbacks, - ) + with _callback_lists_lock: + self._safe_add_callback_to_list( + callback=callback, + parent_list=litellm.callbacks, + ) def add_litellm_success_callback(self, callback: CustomLogger | str | Callable[..., object]): """ @@ -71,37 +77,41 @@ class LoggingCallbackManager: Auto-routes async callbacks to litellm._async_success_callback. Special-cases 'dynamodb' and 'openmeter' as async callbacks. """ - if ( - isinstance(callback, str) - and callback in ("dynamodb", "openmeter") - or not isinstance(callback, str) - and self._is_async_callable(callback) - ): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback) + with _callback_lists_lock: + if ( + isinstance(callback, str) + and callback in ("dynamodb", "openmeter") + or not isinstance(callback, str) + and self._is_async_callable(callback) + ): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback) def add_litellm_failure_callback(self, callback: CustomLogger | str | Callable[..., object]): """ Add a failure callback to `litellm.failure_callback`. Auto-routes async callbacks to litellm._async_failure_callback. """ - if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback) + with _callback_lists_lock: + if not isinstance(callback, str) and self._is_async_callable(callback): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback) def add_litellm_async_success_callback(self, callback: CustomLogger | Callable[..., object] | str): """ Add a success callback to litellm._async_success_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) def add_litellm_async_failure_callback(self, callback: CustomLogger | Callable[..., object] | str): """ Add a failure callback to litellm._async_failure_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) def remove_callback_from_list_by_object(self, callback_list, obj, require_self=True): """ @@ -110,13 +120,14 @@ class LoggingCallbackManager: if not isinstance(callback_list, list): # Not list -> do nothing return - if require_self: - remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj] - else: - remove_list = [c for c in callback_list if c == obj] + with _callback_lists_lock: + if require_self: + remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj] + else: + remove_list = [c for c in callback_list if c == obj] - for c in remove_list: - callback_list.remove(c) + for c in remove_list: + callback_list.remove(c) def remove_callbacks_by_type(self, callback_list, callback_type): """ @@ -134,10 +145,11 @@ class LoggingCallbackManager: if not isinstance(callback_list, list): return - remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)] + with _callback_lists_lock: + remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)] - for c in remove_list: - callback_list.remove(c) + for c in remove_list: + callback_list.remove(c) def _add_string_callback_to_list( self, callback: str, parent_list: list[CustomLogger | Callable[..., object] | str] @@ -336,6 +348,55 @@ class LoggingCallbackManager: return "-".join(key_parts) + get_custom_logger_key = _get_custom_logger_key + callback_lists_lock: Final = _callback_lists_lock + + def drop_callbacks(self, callbacks: Iterable[object], method_owners: Iterable[object] = ()) -> frozenset[str]: + dropped_ids: Final = frozenset(map(id, callbacks)) + owner_ids: Final = frozenset(map(id, method_owners)) + + def keep(callback: object) -> bool: + if id(callback) in dropped_ids: + return False + return not (isinstance(callback, MethodType) and id(callback.__self__) in owner_ids) + + with _callback_lists_lock: + vacated_keys: Final = frozenset( + self._get_custom_logger_key(callback) + for callback in litellm.callbacks + if isinstance(callback, CustomLogger) and id(callback) in dropped_ids + ) + 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) + ] + return vacated_keys + + def with_custom_logger_keys( + self, candidates: Iterable[CustomLogger], keys: frozenset[str] + ) -> tuple[CustomLogger, ...]: + if not keys: + return () + class_names: Final = frozenset(key.split("-", 1)[0] for key in keys) + same_class: Final = (candidate for candidate in candidates if type(candidate).__name__ in class_names) + return tuple(candidate for candidate in same_class if self._get_custom_logger_key(candidate) in keys) + + def without_custom_logger_keys( + self, callbacks: Iterable[str | Callable[..., object] | CustomLogger], keys: frozenset[str] + ) -> Iterator[str | Callable[..., object] | CustomLogger]: + return ( + callback + for callback in callbacks + if not (isinstance(callback, CustomLogger) and self._get_custom_logger_key(callback) in keys) + ) + def _reset_all_callbacks(self): """ Reset all callbacks to an empty list @@ -363,19 +424,21 @@ class LoggingCallbackManager: _get_all_callbacks = get_all_callbacks - def remove_callback_from_all_lists(self, obj, require_self=False) -> None: + def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None: """ Remove a callback object from every callback list it may have been promoted into, so a re-initialized callback leaves no stale instance behind. """ - for callback_list in ( - litellm.callbacks, - litellm.success_callback, - litellm.failure_callback, - litellm._async_success_callback, - litellm._async_failure_callback, - ): - self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self) + with _callback_lists_lock: + for callback_list in ( + litellm.callbacks, + litellm.input_callback, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ): + self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self) def get_active_additional_logging_utils_from_custom_logger( self, diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 37cef41034b..bc9cee1cb1d 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -374,6 +374,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() self._detached_increment_operations = None + self._retired = False + self._flush_after_retire = None self.deployment_budget_config = None @with_service_target("model_budgets") diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fe7ac30e434..e1e8864678e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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 diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 25f9f1dff15..cb226929167 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,8 +1,13 @@ import asyncio -from collections.abc import Mapping +import hashlib +import json +from collections import OrderedDict +from collections.abc import Callable, Mapping from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Any, Final, Literal +from uuid import uuid4 import httpx from fastapi import HTTPException, status @@ -10,6 +15,7 @@ from pydantic import BeforeValidator, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router_utils.common_utils import is_proxy_admin_request @@ -54,7 +60,34 @@ _ROUTER_ONLY_SETTINGS: Final = frozenset( ) -def _route_user_config_request(data: dict, route_type: str): +class UserConfigRouters: + def __init__(self, max_size: int) -> None: + self._max_size: Final = max_size + self._routers: Final[OrderedDict[str, LitellmRouter]] = OrderedDict() # mutable-ok: LRU evicts in place + + def get(self, config: Mapping[str, object], build: Callable[[], LitellmRouter]) -> LitellmRouter: + key: Final = hashlib.sha256(json.dumps(config, sort_keys=True, default=str).encode()).hexdigest() + cached: Final = self._routers.get(key) + if cached is not None: + self._routers.move_to_end(key) + return cached + router: Final = build() + router.serve_through_request_logging() + self._routers[key] = router + if len(self._routers) > self._max_size: + self._routers.popitem(last=False)[1].discard() + return router + + def clear(self) -> None: + for router in self._routers.values(): + router.discard() + self._routers.clear() + + +user_config_routers: Final = UserConfigRouters(max_size=32) + + +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") @@ -63,10 +96,24 @@ def _route_user_config_request(data: dict, route_type: str): valid_args: Final = litellm.Router.get_valid_args() 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 + logging_obj: Final = _request_logging_obj(JSON_OBJECT_ADAPTER.validate_python(data), route_type) + cached_router: Final = user_config_routers.get(filtered_config, lambda: litellm.Router(**filtered_config)) + cached_router.bind_to_request_logging(logging_obj) + return await getattr(cached_router, f"{route_type}")(**{**data, "litellm_logging_obj": logging_obj}) + + +def _request_logging_obj(data: Mapping[str, object], route_type: str) -> LiteLLMLoggingObj: + arrived_with: Final = data.get("litellm_logging_obj") + if isinstance(arrived_with, LiteLLMLoggingObj): + return arrived_with + logging_obj, _ = litellm.utils.function_setup( + original_function=route_type, + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), # noqa: DTZ005 # Logging subtracts it from naive end times + is_async_call=True, + **{"litellm_call_id": str(uuid4()), **data}, + ) + return logging_obj def _is_a2a_agent_model(model_name: object) -> bool: diff --git a/litellm/router.py b/litellm/router.py index 1ce1e0a5a70..0ac21960be1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -38,6 +38,7 @@ from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast @@ -553,6 +554,7 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _CallbackT = TypeVar("_CallbackT") _DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) + _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( {"model", "api_base", "api_key", "api_version", "responses_api_version"} ) @@ -1047,6 +1049,8 @@ class RoutingArgs(enum.Enum): # as the per-request one built from a caller-supplied user_config, drops out on its # own rather than leaving entries behind that nothing can withdraw. _live_routers: Final["weakref.WeakSet[Router]"] = weakref.WeakSet() +_routers_holding_global_callbacks: Final["weakref.WeakSet[Router]"] = weakref.WeakSet() +_routers_holding_global_callbacks_lock: Final = threading.Lock() def _replay_live_router_model_cost() -> None: @@ -1579,22 +1583,18 @@ class Router: self._init_routing_groups(self._routing_groups_input) self._override_selectors: dict[str, RouterStrategySelector | None] = {} self._override_selectors_lock = threading.Lock() + self._request_scoped_selectors: tuple[CustomLogger, ...] = () + self._request_scoped_policies: tuple[CustomLogger, ...] = () self.access_groups = None + with _routers_holding_global_callbacks_lock: + _routers_holding_global_callbacks.add(self) ## 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) @@ -1690,20 +1690,116 @@ 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.retire_routing_state() + self._release_global_callbacks() - # 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 - ) + def _release_global_callbacks(self) -> None: + with _routers_holding_global_callbacks_lock: + _routers_holding_global_callbacks.discard(self) + slack_alerting_loggers: Final = self._slack_alerting_loggers() + adaptive_router_hooks: Final = self._adaptive_router_hooks() + vacated_keys: Final = litellm.logging_callback_manager.drop_callbacks( + ( + *self._callback_selectors(), + *self._override_selectors.values(), + *(self.optional_callbacks or ()), + *slack_alerting_loggers, + *adaptive_router_hooks, + ), + method_owners=(self, *slack_alerting_loggers), + ) + with _routers_holding_global_callbacks_lock: + survivors: Final = tuple(_routers_holding_global_callbacks) + candidates: Final = chain.from_iterable(router._global_callback_candidates() for router in survivors) + for candidate in litellm.logging_callback_manager.with_custom_logger_keys(candidates, vacated_keys): + self._register_router_selector(candidate) + adaptive_keys: Final = frozenset( + map(litellm.logging_callback_manager.get_custom_logger_key, adaptive_router_hooks) + ) + adaptive_survivor: Final = next((router for router in survivors if router.adaptive_routers), None) + if adaptive_survivor is not None and vacated_keys & adaptive_keys: + adaptive_survivor._sync_adaptive_router_hooks() + + def serve_through_request_logging(self) -> None: + self._request_scoped_selectors = (*self._callback_selectors(), *self._request_adaptive_router_hooks()) + self._request_scoped_policies = (*self._optional_loggers(), *self._slack_alerting_loggers()) + 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(): + logging_obj.add_dynamic_callback(policy) + + def _request_policies_not_set_globally(self) -> tuple[CustomLogger, ...]: + global_keys: Final = frozenset( + map( + litellm.logging_callback_manager.get_custom_logger_key, + (callback for callback in litellm.callbacks if isinstance(callback, CustomLogger)), + ) + ) + return tuple( + policy + for policy in self._request_scoped_policies + if litellm.logging_callback_manager.get_custom_logger_key(policy) not in global_keys + ) + + def _request_adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + return tuple( + AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) + for tagged in chain.from_iterable(self.adaptive_routers.values()) + ) + + def _routing_callbacks(self) -> Iterable[object]: + if not (self._request_scoped_selectors or self._request_scoped_policies): + return litellm.callbacks + replaced_keys: Final = frozenset( + map(litellm.logging_callback_manager.get_custom_logger_key, self._request_scoped_selectors) + ) + return chain( + litellm.logging_callback_manager.without_custom_logger_keys(litellm.callbacks, replaced_keys), + self._request_scoped_selectors, + self._request_policies_not_set_globally(), + ) + + def retire_routing_state(self) -> None: + for strategy in (*self._global_callback_candidates(), *self._override_selectors.values()): + if isinstance(strategy, BaseRoutingStrategy | RouterBudgetLimiting): + strategy.retire() + + def _slack_alerting_loggers(self) -> tuple[CustomLogger, ...]: + slack_alerting_logger: Final = getattr(self, "slack_alerting_logger", None) + return (slack_alerting_logger,) if isinstance(slack_alerting_logger, CustomLogger) else () + + def _adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + own_strategy_ids: Final = frozenset( + id(tagged.strategy) for tagged in chain.from_iterable(self.adaptive_routers.values()) + ) + hooks: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + return tuple( + hook + for hook in hooks + if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own_strategy_ids + ) + + def _optional_loggers(self) -> tuple[CustomLogger, ...]: + return tuple(callback for callback in self.optional_callbacks or () if isinstance(callback, CustomLogger)) + + def _global_callback_candidates(self) -> tuple[CustomLogger, ...]: + return (*self._callback_selectors(), *self._optional_loggers(), *self._slack_alerting_loggers()) @staticmethod def _create_redis_cache( @@ -1809,7 +1905,7 @@ class Router: return selector @staticmethod - def _register_router_selector(selector: RouterStrategySelector) -> None: + def _register_router_selector(selector: CustomLogger) -> None: if isinstance(selector, LeastBusyLoggingHandler): if isinstance(litellm.input_callback, list): litellm.logging_callback_manager.add_litellm_input_callback(selector) @@ -1818,23 +1914,28 @@ 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 _unregister_router_selectors(self, selectors: Sequence[object]) -> None: """ Drop router-owned strategy selectors from litellm's global callback - lists by identity. Used before re-init (`routing_strategy_init` / - `_init_routing_groups`) so repeated `update_settings` calls don't - accumulate dead selectors that keep receiving callback events. + lists. Used before re-init (`routing_strategy_init` / + `_init_routing_groups`) and by `discard` so dead selectors stop + receiving callback events. `function_setup` copies `litellm.callbacks` + into the per-event lists on every request, so those are purged too. """ for selector in selectors: if isinstance(selector, BaseRoutingStrategy): selector.retire() - selector_ids: Final = {id(s) for s in selectors if s is not None} - if not selector_ids: - return - if isinstance(litellm.callbacks, list): - litellm.callbacks = [c for c in litellm.callbacks if id(c) not in selector_ids] - if isinstance(litellm.input_callback, list): - litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids] + litellm.logging_callback_manager.drop_callbacks(selectors) def _apply_updated_routing_strategy_args(self) -> None: """ @@ -9804,7 +9905,7 @@ class Router: Raises: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): _callback.pre_call_check(deployment) @@ -9825,7 +9926,7 @@ class Router: Raises: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): try: await _callback.async_pre_call_check(deployment, parent_otel_span) @@ -9962,7 +10063,7 @@ class Router: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ returned_healthy_deployments = healthy_deployments - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): try: returned_healthy_deployments = await _callback.async_filter_deployments( @@ -10617,6 +10718,7 @@ class Router: changes, otherwise a released router keeps recording turns through its hook.""" from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, ) for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): @@ -10626,6 +10728,8 @@ class Router: litellm.logging_callback_manager.add_litellm_callback( AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) ) + if self.adaptive_routers: + litellm.logging_callback_manager.add_litellm_callback(AdaptiveRouterResponseHeaderHook()) def init_adaptive_router_deployment(self, deployment: Deployment) -> None: """ diff --git a/litellm/router_strategy/adaptive_router/__init__.py b/litellm/router_strategy/adaptive_router/__init__.py index d7f55ebced9..1f7e6b93169 100644 --- a/litellm/router_strategy/adaptive_router/__init__.py +++ b/litellm/router_strategy/adaptive_router/__init__.py @@ -1,6 +1,9 @@ """Adaptive router strategy. See README.md for design overview.""" from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter -from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, +) -__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"] +__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook", "AdaptiveRouterResponseHeaderHook"] diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 92207e07c52..3427b81be68 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -170,12 +170,7 @@ def _assistant_content_and_tool_calls(response_obj: Any) -> tuple[object, Sequen return content, tool_calls -class AdaptiveRouterPostCallHook(CustomLogger): - """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" - - def __init__(self, adaptive_router: AdaptiveRouter) -> None: - self.adaptive_router = adaptive_router - +class AdaptiveRouterResponseHeaderHook(CustomLogger): async def async_post_call_response_headers_hook( self, data: Mapping[str, object], @@ -200,6 +195,13 @@ class AdaptiveRouterPostCallHook(CustomLogger): return None return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen} + +class AdaptiveRouterPostCallHook(CustomLogger): + """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" + + def __init__(self, adaptive_router: AdaptiveRouter) -> None: + self.adaptive_router = adaptive_router + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): await self._record(kwargs, response_obj, response_status=200) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 51400e6da24..2cea7389808 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -24,6 +24,8 @@ class BaseRoutingStrategy(ABC): self.dual_cache = dual_cache self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] self._sync_task: asyncio.Task[None] | None = None + self._retired = False + self._flush_after_retire: asyncio.Task[None] | None = None if should_batch_redis_writes: self.setup_sync_task(default_sync_interval) @@ -47,13 +49,25 @@ class BaseRoutingStrategy(ABC): def retire(self) -> None: self.cancel_sync_task() - if not self.redis_increment_operation_queue: + self._retired = True + self._schedule_flush_after_retire() + + def _schedule_flush_after_retire(self) -> None: + if self._flush_after_retire is not None and not self._flush_after_retire.done(): + return + if not self.redis_increment_operation_queue or self.dual_cache.redis_cache is None: return try: loop: Final = asyncio.get_running_loop() except RuntimeError: return - loop.create_task(self._push_in_memory_increments_to_redis()) + self._flush_after_retire = loop.create_task(self._drain_increments_after_retire()) + + async def _drain_increments_after_retire(self) -> None: + if not self.redis_increment_operation_queue: + return + await self._push_in_memory_increments_to_redis() + await self._drain_increments_after_retire() async def cleanup(self): """Cleanup method to be called when shutting down""" @@ -98,6 +112,8 @@ class BaseRoutingStrategy(ABC): self.redis_increment_operation_queue.append(increment_op) self.add_to_in_memory_keys_to_update(key=key) + if self._retired: + self._schedule_flush_after_retire() return result @with_service_target("router_usage") diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 147ce60a99b..eab232eae18 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -128,7 +128,9 @@ class RouterBudgetLimiting(CustomLogger): self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() self._detached_increment_operations: tuple[RedisPipelineIncrementOperation, ...] | None = None - asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) + self._sync_task: asyncio.Task[None] = asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) + self._retired = False + self._flush_after_retire: asyncio.Task[None] | None = None self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config self.deployment_budget_config: GenericBudgetConfigType | None = None self.tag_budget_config: GenericBudgetConfigType | None = None @@ -140,6 +142,30 @@ class RouterBudgetLimiting(CustomLogger): if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(self) + def retire(self) -> None: + self._sync_task.cancel() + self._retired = True + self._schedule_flush_after_retire() + + def _schedule_flush_after_retire(self) -> None: + if self._flush_after_retire is not None and not self._flush_after_retire.done(): + return + if not self._has_unwritten_increments(): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + self._flush_after_retire = loop.create_task(self._drain_increments_after_retire()) + + def _has_unwritten_increments(self) -> bool: + return bool(self.redis_increment_operation_queue) or self._detached_increment_operations is not None + + async def _drain_increments_after_retire(self) -> None: + while self._has_unwritten_increments(): + if not await self._push_in_memory_increments_to_redis(): + await asyncio.sleep(DEFAULT_REDIS_SYNC_INTERVAL) + @with_service_target("router_budgets") async def async_filter_deployments( self, @@ -427,6 +453,8 @@ class RouterBudgetLimiting(CustomLogger): if self.dual_cache.redis_cache is None: return self.redis_increment_operation_queue.append(increment_op) + if self._retired: + self._schedule_flush_after_retire() def _get_redis_increment_queue_lock(self) -> asyncio.Lock: return self._redis_increment_queue_lock diff --git a/litellm/utils.py b/litellm/utils.py index 1aad9baf06a..1a114ffe9f0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1126,83 +1126,90 @@ def function_setup( dynamic_callbacks: Final[list[str | Callable | CustomLogger] | None] = kwargs.pop("callbacks", None) all_callbacks: Final = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) - if len(all_callbacks) > 0: - for callback in all_callbacks: - # check if callback is a string - e.g. "lago", "openmeter" - if isinstance(callback, str): - callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( - callback, - internal_usage_cache=None, - llm_router=None, - ) - if callback is None or any( - type(cb) is type(callback) for cb in litellm._async_success_callback - ): # don't double add a callback - continue - if callback not in litellm.input_callback: - litellm.input_callback.append(callback) - if callback not in litellm.success_callback: - litellm.logging_callback_manager.add_litellm_success_callback(callback) - if callback not in litellm.failure_callback: - litellm.logging_callback_manager.add_litellm_failure_callback(callback) - if callback not in litellm._async_success_callback: - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - if callback not in litellm._async_failure_callback: - litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) - print_verbose(f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}") + with litellm.logging_callback_manager.callback_lists_lock: + if len(all_callbacks) > 0: + for callback in all_callbacks: + # check if callback is a string - e.g. "lago", "openmeter" + if isinstance(callback, str): + callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( + callback, + internal_usage_cache=None, + llm_router=None, + ) + if callback is None or any( + type(cb) is type(callback) for cb in litellm._async_success_callback + ): # don't double add a callback + continue + if callback not in litellm.input_callback: + litellm.input_callback.append(callback) + if callback not in litellm.success_callback: + litellm.logging_callback_manager.add_litellm_success_callback(callback) + if callback not in litellm.failure_callback: + litellm.logging_callback_manager.add_litellm_failure_callback(callback) + if callback not in litellm._async_success_callback: + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + if callback not in litellm._async_failure_callback: + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + print_verbose( + f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}" + ) - if ( - len(litellm.input_callback) > 0 or len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0 - ) and len(callback_list) == 0: - callback_list = list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) + callbacks_to_set_up: Final = ( + list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) + if len(callback_list) == 0 + else [] + ) + if callbacks_to_set_up: + callback_list = callbacks_to_set_up + ## ASYNC CALLBACKS - safety net for callbacks added via direct append + if len(litellm.input_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.input_callback): + if coroutine_checker.is_async_callable(callback): + litellm._async_input_callback.append(callback) + removed_async_items.append(index) + + # Pop the async items from input_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.input_callback.pop(index) + if len(litellm.success_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.success_callback): + if coroutine_checker.is_async_callable(callback): + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + removed_async_items.append(index) + elif callback == "dynamodb" or callback == "openmeter": + # dynamo is an async callback, it's used for the proxy and needs to be async + # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + removed_async_items.append(index) + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): + add_custom_logger_callback_to_specific_event(callback, "success") + + # Pop the async items from success_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.success_callback.pop(index) + + if len(litellm.failure_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.failure_callback): + if coroutine_checker.is_async_callable(callback): + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + removed_async_items.append(index) + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): + add_custom_logger_callback_to_specific_event(callback, "failure") + + # Pop the async items from failure_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.failure_callback.pop(index) + if callbacks_to_set_up: get_set_callbacks: Final[Callable[[], Callable[..., None]]] = litellm_utils.get_set_callbacks - get_set_callbacks()(callback_list=callback_list, function_id=function_id) - ## ASYNC CALLBACKS - safety net for callbacks added via direct append - if len(litellm.input_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.input_callback): - if coroutine_checker.is_async_callable(callback): - litellm._async_input_callback.append(callback) - removed_async_items.append(index) - - # Pop the async items from input_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.input_callback.pop(index) - if len(litellm.success_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.success_callback): - if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - removed_async_items.append(index) - elif callback == "dynamodb" or callback == "openmeter": - # dynamo is an async callback, it's used for the proxy and needs to be async - # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - removed_async_items.append(index) - elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): - add_custom_logger_callback_to_specific_event(callback, "success") - - # Pop the async items from success_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.success_callback.pop(index) - - if len(litellm.failure_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.failure_callback): - if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) - removed_async_items.append(index) - elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): - add_custom_logger_callback_to_specific_event(callback, "failure") - - # Pop the async items from failure_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.failure_callback.pop(index) + get_set_callbacks()(callback_list=callbacks_to_set_up, function_id=function_id) ### 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"]): diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 55435c51648..af2f386a362 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -146,6 +146,18 @@ 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 + "_callback_selectors", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "bind_to_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py + "serve_through_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py + "_request_adaptive_router_hooks", # Run by serve_through_request_logging, tested through the user_config adaptive test + "_release_global_callbacks", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "_routing_callbacks", # Read by the pre-call checks, tested through the user_config limit tests in test_route_llm_request.py + "_request_policies_not_set_globally", # Read by bind_to_request_logging, tested through the user_config budget test in test_route_llm_request.py + "_optional_loggers", # Read by discard and serve_through_request_logging, tested through the budget limiter discard test + "retire_routing_state", # Run by discard, tested through the budget limiter discard test + "_slack_alerting_loggers", # Run by discard, tested through the Slack alerting discard test in test_router_helper_utils.py + "_adaptive_router_hooks", # Run by discard, tested through the adaptive router discard tests in test_router_helper_utils.py + "_global_callback_candidates", # Run by discard, tested through the survivor discard tests in test_router_helper_utils.py "without_deployments_that_cannot_batch", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py "batch_create_routing", # Tested through the mixed-group acreate_batch tests in test_router.py "is_batch_create_routing", # Tested through the chat-after-batch-create test in test_router.py diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index c96e3e2c031..e64f5cca5b8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -252,6 +252,7 @@ def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> It @pytest.fixture(autouse=True) def isolate_litellm_globals() -> Iterator[None]: original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS} + original_callback_holders: Final = frozenset(litellm_router_module._routers_holding_global_callbacks) original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS} original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)} original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__} @@ -263,6 +264,7 @@ def isolate_litellm_globals() -> Iterator[None]: _flush_client_caches() for name in CALLBACK_LISTS: setattr(litellm, name, []) + litellm_router_module._routers_holding_global_callbacks.clear() for name in RESET_TO_NONE_GLOBALS: setattr(litellm, name, None) for name in MODULE_LEVEL_CLIENTS: @@ -277,6 +279,8 @@ def isolate_litellm_globals() -> Iterator[None]: litellm.__dict__.update(original_clients) for name, value in (original_callbacks | original_reset | original_restored).items(): setattr(litellm, name, value) + litellm_router_module._routers_holding_global_callbacks.clear() + litellm_router_module._routers_holding_global_callbacks.update(original_callback_holders) for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items(): logger.setLevel(level) logger.disabled = disabled diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 79645ffe8f6..4f8a13c1182 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -402,6 +402,32 @@ def test_get_combined_callback_list_preserves_insertion_order(logging_obj): ) == ["prometheus", "langfuse", "datadog", "otel", "s3", "gcs_bucket", "arize", "logfire"] +class _InputEventRecorder(CustomLogger): + def __init__(self, label: str, events: list[str]) -> None: + self._label = label + self.events = events + + def log_pre_api_call(self, model: object, messages: object, kwargs: object) -> None: + self.events.append(f"{self._label} pre") + + def log_post_api_call(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.events.append(f"{self._label} post") + + +def test_input_events_reach_the_replacement_and_not_the_global_callback_it_replaced( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + events: Final[list[str]] = [] + monkeypatch.setattr(litellm, "input_callback", [_InputEventRecorder("global", events)]) + logging_obj.replace_global_callback(_InputEventRecorder("replacement", events)) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + + logging_obj.pre_call(input="", api_key="") + logging_obj.post_call(original_response="ok", input="", api_key="") + + assert events == ["replacement pre", "replacement post"] + + def test_get_masked_api_base(logging_obj): api_base = "https://api.openai.com/v1" masked_api_base = logging_obj._get_masked_api_base(api_base) diff --git a/tests/unit/litellm_core_utils/test_logging_callback_manager.py b/tests/unit/litellm_core_utils/test_logging_callback_manager.py index f8260702068..25a52137d65 100644 --- a/tests/unit/litellm_core_utils/test_logging_callback_manager.py +++ b/tests/unit/litellm_core_utils/test_logging_callback_manager.py @@ -145,6 +145,7 @@ def test_remove_callback_from_all_lists(): obj = TestLogger() manager.add_litellm_callback(obj) + manager.add_litellm_input_callback(obj) manager.add_litellm_success_callback(obj) manager.add_litellm_failure_callback(obj) manager.add_litellm_async_success_callback(obj) @@ -153,6 +154,7 @@ def test_remove_callback_from_all_lists(): manager.remove_callback_from_all_lists(obj) assert obj not in litellm.callbacks + assert obj not in litellm.input_callback assert obj not in litellm.success_callback assert obj not in litellm.failure_callback assert obj not in litellm._async_success_callback diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 8928b79bcd4..6044475652c 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,19 +1,42 @@ +import asyncio import contextlib +import gc +import threading +import weakref import httpx import json import litellm import pytest +import redis import respx +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping +from datetime import datetime from types import SimpleNamespace +from uuid import uuid4 from typing import Final from unittest.mock import MagicMock +from fakeredis import TcpFakeServer from fastapi import HTTPException -from litellm.proxy.route_llm_request import JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, route_request -from litellm.types.router import RouterGeneralSettings +from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.route_llm_request import ( + JSON_OBJECT_ADAPTER, + ProxyModelNotFoundError, + UserConfigRouters, + route_request, + user_config_routers, +) +from litellm.proxy.utils import ProxyLogging +from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ADAPTIVE_ROUTER_RESPONSE_HEADER, +) +from litellm.types.router import RouterErrors, RouterGeneralSettings, RouterRateLimitError @pytest.mark.parametrize( @@ -2434,3 +2457,571 @@ 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(autouse=True) +def discard_cached_user_config_routers() -> Iterator[None]: + yield + user_config_routers.clear() + + +@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 _wait_out_a_rate_limit_minute_rollover() -> None: + seconds_left: Final = 60 - datetime.now().second + if seconds_left < 10: + await asyncio.sleep(seconds_left) + + +def _proxy_logging_obj(model: str = "rate-limited") -> LiteLLMLoggingObj: + logging_obj, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model=model, + messages=[{"role": "user", "content": "hi"}], + litellm_call_id=str(uuid4()), + ) + assert isinstance(logging_obj, LiteLLMLoggingObj) + return logging_obj + + +def _mock_deployment(model_name: str, **deployment_limits: int) -> dict[str, object]: + return { + "model_name": model_name, + **deployment_limits, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "mock_response": "ok"}, + } + + +async def _user_config_request_outcome( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + logging_obj: LiteLLMLoggingObj | None, +) -> str: + data: Final = { + "model": "rate-limited", + "messages": [{"role": "user", "content": "hi"}], + "num_retries": 0, + **({"litellm_logging_obj": logging_obj} if logging_obj is not None else {}), + "user_config": { + **routing_settings, + **({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}), + "model_list": [_mock_deployment("rate-limited", **deployment_limits)], + }, + } + try: + await (await route_request(data, None, None, "acompletion")) + except (litellm.RateLimitError, RouterRateLimitError): + return "429" + finally: + await asyncio.sleep(0.3) + return "200" + + +async def _tpm_for_three_requests() -> int: + messages: Final = [{"role": "user", "content": "hi"}] + probe: Final = await litellm.acompletion(model="openai/gpt-5.5", messages=messages, mock_response="ok") + assert isinstance(probe, litellm.ModelResponse) + probe_usage: Final = getattr(probe, "usage", None) + assert isinstance(probe_usage, litellm.Usage) + return 2 * probe_usage.total_tokens + litellm.token_counter(messages=messages) + + +async def _request_cost() -> float: + probe: Final = await litellm.acompletion( + model="openai/gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok" + ) + return litellm.completion_cost(completion_response=probe) + + +async def _proxy_router_outcomes( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + requests: int, +) -> list[str]: + router: Final = litellm.Router( + **routing_settings, + num_retries=0, + model_list=[_mock_deployment("proxy-rate-limited", **deployment_limits)], + **({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}), + ) + + async def outcome() -> str: + try: + await router.acompletion(model="proxy-rate-limited", messages=[{"role": "user", "content": "hi"}]) + except (litellm.RateLimitError, RouterRateLimitError): + return "429" + finally: + await asyncio.sleep(0.3) + return "200" + + try: + return [await outcome() for _ in range(requests)] + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_a_user_config_adaptive_request_answers_with_the_model_it_picked() -> None: + data: Final = { + "model": "smart-cheap-router", + "messages": [{"role": "user", "content": "Write a one line python function that adds two numbers"}], + "metadata": {}, + "user_config": { + "model_list": [ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["fast", "smart"]}, + }, + }, + { + **_mock_deployment("fast"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 2, "strengths": []}}, + }, + { + **_mock_deployment("smart"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]}}, + }, + ] + }, + } + + await (await route_request(data, None, None, "acompletion")) + headers: Final = await ProxyLogging(user_api_key_cache=DualCache()).post_call_response_headers_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), response=None + ) + + assert headers == {ADAPTIVE_ROUTER_RESPONSE_HEADER: data["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY]} + + +def _global_router_callbacks() -> tuple[list[object], list[object]]: + return list(litellm.callbacks), list(litellm.input_callback) + + +@pytest.mark.parametrize( + "routing_settings, deployment_limits", + [ + ({"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}), + ({"routing_strategy": "usage-based-routing-v2"}, {"tpm": 30}), + ({"routing_strategy": "usage-based-routing"}, {"rpm": 2}), + ({"optional_pre_call_checks": ["enforce_model_rate_limits"]}, {"rpm": 2}), + ], +) +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.parametrize("with_redis", [True, False]) +@pytest.mark.asyncio +async def test_user_config_router_enforces_its_limit_across_proxy_requests( + fake_redis_port: int, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + arrives_with_logging_obj: bool, + with_redis: bool, +) -> None: + await _wait_out_a_rate_limit_minute_rollover() + callbacks_before: Final = _global_router_callbacks() + outcomes: Final = [ + await _user_config_request_outcome( + fake_redis_port if with_redis else None, + routing_settings, + deployment_limits, + _proxy_logging_obj() if arrives_with_logging_obj else None, + ) + for _ in range(6) + ] + + assert "429" in outcomes, outcomes + assert _global_router_callbacks() == callbacks_before + assert outcomes == await _proxy_router_outcomes( + fake_redis_port if with_redis else None, routing_settings, deployment_limits, len(outcomes) + ) + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_limit( + fake_redis_port: int, arrives_with_logging_obj: bool +) -> None: + tpm_for_three_requests: Final = await _tpm_for_three_requests() + + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, + {"routing_strategy": "usage-based-routing-v2"}, + {"tpm": tpm_for_three_requests}, + _proxy_logging_obj() if arrives_with_logging_obj else None, + ) + + await _wait_out_a_rate_limit_minute_rollover() + concurrent_outcomes: Final = await asyncio.gather(request(), request()) + later_outcomes: Final = [await request() for _ in range(2)] + + assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"] + + +@pytest.mark.asyncio +async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None: + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}, _proxy_logging_obj() + ) + + await _wait_out_a_rate_limit_minute_rollover() + outcomes: Final = await asyncio.gather(*(request() for _ in range(4))) + + assert outcomes.count("200") <= 2, outcomes + assert "429" in outcomes, outcomes + + +@pytest.mark.parametrize( + "routing_strategy, limited_by", + [("usage-based-routing-v2", "rpm"), ("usage-based-routing-v2", "tpm"), ("usage-based-routing", "rpm")], +) +@pytest.mark.parametrize("proxy_router_shares_redis", [True, False]) +@pytest.mark.asyncio +async def test_a_proxy_router_on_the_same_strategy_leaves_the_user_config_limit_unchanged( + fake_redis_port: int, routing_strategy: str, limited_by: str, proxy_router_shares_redis: bool +) -> None: + deployment_limits: Final = {"rpm": 3} if limited_by == "rpm" else {"tpm": await _tpm_for_three_requests()} + + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, {"routing_strategy": routing_strategy}, deployment_limits, _proxy_logging_obj() + ) + + async def outcomes() -> list[str]: + first_instance: Final = [await request() for _ in range(2)] + user_config_routers.clear() + second_instance: Final = [await request() for _ in range(3)] + return [*first_instance, *second_instance] + + await _wait_out_a_rate_limit_minute_rollover() + alone: Final = await outcomes() + user_config_routers.clear() + await asyncio.sleep(0.3) + redis.Redis(host="127.0.0.1", port=fake_redis_port).flushall() + proxy_router: Final = litellm.Router( + routing_strategy=routing_strategy, + model_list=[_mock_deployment("proxy-model")], + redis_host="127.0.0.1" if proxy_router_shares_redis else None, + redis_port=fake_redis_port if proxy_router_shares_redis else None, + ) + try: + beside_proxy_router: Final = await outcomes() + finally: + proxy_router.discard() + + assert "429" in alone, alone + assert beside_proxy_router == alone + + +@pytest.mark.parametrize("user_config_settings", [{}, {"optional_pre_call_checks": ["router_budget_limiting"]}]) +@pytest.mark.asyncio +async def test_user_config_requests_stay_under_the_proxy_provider_budget( + user_config_settings: Mapping[str, object], +) -> None: + proxy_router: Final = litellm.Router( + model_list=[_mock_deployment("proxy-model")], + provider_budget_config={"openai": {"budget_limit": await _request_cost() / 2, "time_period": "1d"}}, + ) + try: + await proxy_router.acompletion(model="proxy-model", messages=[{"role": "user", "content": "hi"}]) + await asyncio.sleep(0.3) + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await _user_config_request_outcome(None, user_config_settings, {}, _proxy_logging_obj()) + finally: + proxy_router.discard() + + +@pytest.mark.asyncio +async def test_a_user_config_provider_budget_limits_its_own_requests() -> None: + budget_limit: Final = await _request_cost() / 2 + + def request() -> Coroutine[None, None, str]: + budget: Final = {"openai": {"budget_limit": budget_limit, "time_period": "1d"}} + return _user_config_request_outcome(None, {"provider_budget_config": budget}, {}, _proxy_logging_obj()) + + first: Final = await request() + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await request() + assert first == "200" + + +@pytest.mark.asyncio +async def test_a_request_in_flight_when_its_user_config_router_is_dropped_still_counts_against_its_budget( + fake_redis_port: int, +) -> None: + budget_limit: Final = 1.5 * await _request_cost() + + async def request() -> None: + await ( + await route_request( + { + "model": "budgeted", + "messages": [{"role": "user", "content": "hi"}], + "litellm_logging_obj": _proxy_logging_obj("budgeted"), + "user_config": { + "redis_host": "127.0.0.1", + "redis_port": fake_redis_port, + "provider_budget_config": {"openai": {"budget_limit": budget_limit, "time_period": "1d"}}, + "model_list": [ + { + "model_name": "budgeted", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "mock_response": "ok", + "mock_delay": 0.5, + }, + } + ], + }, + }, + None, + None, + "acompletion", + ) + ) + + await request() + in_flight: Final = asyncio.create_task(request()) + await asyncio.sleep(0.2) + user_config_routers.clear() + await in_flight + await asyncio.sleep(0.3) + + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + 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 = { + "model_list": [ + _mock_deployment("smart"), + _mock_deployment("fast"), + { + "model_name": "adaptive", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["smart", "fast"]}, + }, + }, + ] + } + await ( + await route_request( + { + "model": "adaptive", + "messages": [{"role": "user", "content": "write a sorting function"}], + "metadata": {"session_id": "session-1"}, + "litellm_logging_obj": _proxy_logging_obj("adaptive"), + "user_config": user_config, + }, + None, + None, + "acompletion", + ) + ) + await asyncio.sleep(0.3) + + router: Final = user_config_routers.get(user_config, lambda: pytest.fail("the user_config router was not reused")) + snapshot: Final = await router.adaptive_routers["adaptive"][0].strategy.get_state_snapshot() + assert snapshot["feedback_contexts_live"] == 1, snapshot + + +@pytest.mark.asyncio +async def test_user_config_routers_discard_the_least_recently_used_router_past_capacity() -> None: + routers: Final = UserConfigRouters(max_size=2) + + def config(model_name: str) -> dict[str, object]: + return {"routing_strategy": "usage-based-routing-v2", "model_list": [_mock_deployment(model_name)]} + + def build(model_name: str) -> litellm.Router: + return litellm.Router(**config(model_name)) + + def must_be_cached() -> litellm.Router: + pytest.fail("a cached router was rebuilt") + + first_handler: Final = weakref.ref(routers.get(config("first"), lambda: build("first")).lowesttpm_logger_v2) + second_handler: Final = weakref.ref(routers.get(config("second"), lambda: build("second")).lowesttpm_logger_v2) + routers.get(config("first"), must_be_cached) + routers.get(config("third"), lambda: build("third")) + await asyncio.sleep(0.2) + gc.collect() + + assert routers.get(config("first"), must_be_cached) is not None + assert first_handler() is not None + assert second_handler() is None + routers.clear() diff --git a/tests/unit/router_strategy/adaptive_router/test_hooks.py b/tests/unit/router_strategy/adaptive_router/test_hooks.py index ad61f43c5a0..0ed22658144 100644 --- a/tests/unit/router_strategy/adaptive_router/test_hooks.py +++ b/tests/unit/router_strategy/adaptive_router/test_hooks.py @@ -10,6 +10,7 @@ from litellm.router_strategy.adaptive_router.config import ( ) from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, _recent_tool_results, _resolve_session_key, ) @@ -337,8 +338,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): """The header hook returns the `x-litellm-adaptive-router-model` header so proxy header construction picks it up (works for both streaming and non-streaming; `async_post_call_success_hook` is too late for streaming).""" - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"adaptive_router_chosen_model": "smart"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -348,8 +348,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"litellm_session_id": "sess-A"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -359,8 +358,7 @@ async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_no_metadata(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -370,8 +368,7 @@ async def test_post_call_response_headers_hook_noop_when_no_metadata(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_not_dict(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": "not-a-dict"}, user_api_key_dict=MagicMock(), response=MagicMock(), diff --git a/tests/unit/router_strategy/test_base_routing_strategy.py b/tests/unit/router_strategy/test_base_routing_strategy.py index dc75c3d2916..19ff11dd3f6 100644 --- a/tests/unit/router_strategy/test_base_routing_strategy.py +++ b/tests/unit/router_strategy/test_base_routing_strategy.py @@ -1,15 +1,20 @@ import json import logging -from typing import Any, Dict, List, Optional, Set, Union +from typing import Any, Dict, Final, List, Optional, Set, Union import pytest import asyncio +import threading from unittest.mock import MagicMock, patch +from fakeredis import TcpFakeServer + from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, RedisPipelineIncrementOperation from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy @@ -162,3 +167,24 @@ async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_erro assert caplog.records == [] assert base_strategy.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_an_increment_counted_after_retire_still_reaches_redis(): + server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis") + threading.Thread(target=server.serve_forever, daemon=True).start() + redis_cache: Final = RedisCache(host="127.0.0.1", port=server.server_address[1]) + strategy: Final = BaseRoutingStrategy( + dual_cache=DualCache(redis_cache=redis_cache, in_memory_cache=InMemoryCache()), + should_batch_redis_writes=True, + default_sync_interval=60, + ) + try: + strategy.retire() + await strategy._increment_value_in_current_window("counted-after-retire", 3, 60) + await asyncio.sleep(0.2) + + assert float(await redis_cache.async_get_cache("counted-after-retire")) == 3 + finally: + server.shutdown() + server.server_close() diff --git a/tests/unit/router_strategy/test_budget_limiter_hotpath.py b/tests/unit/router_strategy/test_budget_limiter_hotpath.py index 4667ac67efa..10704a29347 100644 --- a/tests/unit/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/unit/router_strategy/test_budget_limiter_hotpath.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import gc import logging from types import SimpleNamespace @@ -492,6 +493,8 @@ def _new_router_budget_limiter( budget_limiter._redis_increment_queue_lock = queue_lock if queue_lock is not None else asyncio.Lock() budget_limiter._redis_increment_flush_lock = asyncio.Lock() budget_limiter._detached_increment_operations = None + budget_limiter._retired = False + budget_limiter._flush_after_retire = None return budget_limiter @@ -784,3 +787,57 @@ async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations: await limiter._push_in_memory_increments_to_redis() assert redis_cache.values[_SPEND_KEY] == 10.0 assert limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_spend_queued_when_a_limiter_is_retired_reaches_redis_after_a_failed_write() -> None: + pipeline_started: Final = asyncio.Event() + pipeline_completed: Final = asyncio.Event() + redis_cache: Final = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + pipeline_completed=pipeline_completed, + should_fail_pipeline=True, + ) + budget_limiter: Final = _new_router_budget_limiter( + redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)] + ) + budget_limiter._sync_task = asyncio.create_task(asyncio.sleep(3600)) + + budget_limiter.retire() + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + assert redis_cache.values == {} + redis_cache.should_fail_pipeline = False + await asyncio.wait_for(pipeline_completed.wait(), timeout=5) + + assert redis_cache.values == {_SPEND_KEY: 10.0} + assert budget_limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_spend_in_flight_when_a_limiter_is_retired_reaches_redis_after_that_write_fails() -> None: + pipeline_started: Final = asyncio.Event() + allow_pipeline_to_complete: Final = asyncio.Event() + pipeline_completed: Final = asyncio.Event() + redis_cache: Final = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + pipeline_completed=pipeline_completed, + ) + budget_limiter: Final = _new_router_budget_limiter( + redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)] + ) + budget_limiter._sync_task = asyncio.create_task(budget_limiter.periodic_sync_in_memory_spend_with_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + + budget_limiter.retire() + allow_pipeline_to_complete.set() + with contextlib.suppress(asyncio.CancelledError): + await budget_limiter._sync_task + redis_cache.should_fail_pipeline = False + await asyncio.wait_for(pipeline_completed.wait(), timeout=5) + + assert redis_cache.values == {_SPEND_KEY: 10.0} + assert budget_limiter.redis_increment_operation_queue == [] diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 7088818aa01..af8e7f53462 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1,8 +1,12 @@ +import asyncio import functools +import gc 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 @@ -12,8 +16,12 @@ from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo +from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.types.litellm_params import RoutingStrategyName +from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( CallTypes, ModelResponse, @@ -1852,6 +1860,315 @@ def test_discard(model_list): assert len(litellm.service_callback) == 0 +_CALLBACK_LIST_NAMES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", +) + + +def _callback_ids(list_name: str) -> list[int]: + return [id(callback) for callback in getattr(litellm, list_name)] + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("least-busy", "leastbusy_logger"), + ("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_releases_strategy_selector_across_router_lifecycles( + routing_strategy: RoutingStrategyName, selector_attr: str +): + def retire() -> weakref.ReferenceType[object]: + router: Final = Router(model_list=[], routing_strategy=routing_strategy) + selector: Final = getattr(router, selector_attr) + assert any(callback is selector for callback in litellm.callbacks) + router.discard() + return weakref.ref(selector) + + selector_refs: Final = tuple(retire() for _ in range(3)) + await asyncio.sleep(0) + gc.collect() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + assert all(selector_ref() is None for selector_ref in selector_refs) + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("least-busy", "leastbusy_logger"), + ("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_purges_selector_from_event_lists_after_serving_a_request( + routing_strategy: RoutingStrategyName, selector_attr: str +): + router: Final = Router( + model_list=[{"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}], + routing_strategy=routing_strategy, + ) + selector: Final = getattr(router, selector_attr) + await router.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + assert any(callback is selector for callback in litellm.success_callback) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def test_discard_keeps_callbacks_owned_by_other_live_routers(): + survivor: Final = Router(model_list=[], routing_strategy="latency-based-routing") + unrelated: Final = CustomLogger() + litellm.logging_callback_manager.add_litellm_callback(unrelated) + before_ids: Final = {list_name: _callback_ids(list_name) for list_name in _CALLBACK_LIST_NAMES} + retired: Final = Router(model_list=[], routing_strategy="least-busy") + assert any(callback is retired.leastbusy_logger for callback in litellm.callbacks) + + retired.discard() + + assert {list_name: _callback_ids(list_name) for list_name in before_ids} == before_ids + assert any(callback is survivor.lowestlatency_logger for callback in litellm.callbacks) + 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=[ + {"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} + ], + routing_groups=[ + {"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "latency-based-routing"} + ], + ) + _, group_selector = router._get_routing_context("grouped-model") + assert any(callback is group_selector for callback in litellm.callbacks) + + router.discard() + + 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_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") + assert not _is_registered(survivor.lowestlatency_logger) + + retired.discard() + + assert _is_registered(survivor.lowestlatency_logger) + survivor.discard() + + +def test_discard_drops_its_slack_alerting_callbacks(): + router: Final = Router(model_list=_GROUPED_MODEL_LIST, alerting_config=AlertingConfig(webhook_url="test")) + assert any(isinstance(callback, SlackAlerting) for callback in litellm.callbacks) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def _adaptive_router_model_list(name: str) -> list[DeploymentTypedDict]: + return [ + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}, + { + "model_name": name, + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["gpt-5.5"]}, + }, + }, + ] + + +def _registered_adaptive_router_strategies() -> tuple[object, ...]: + return tuple( + callback.adaptive_router for callback in litellm.callbacks if isinstance(callback, AdaptiveRouterPostCallHook) + ) + + +def test_discard_drops_its_adaptive_router_hook(): + router: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + assert _registered_adaptive_router_strategies() + + router.discard() + + assert _registered_adaptive_router_strategies() == () + + +def test_discard_registers_the_adaptive_router_hook_of_a_survivor(): + survivor: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + retired: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + + retired.discard() + + survivor_strategies: Final = tuple(tagged.strategy for tagged in survivor.adaptive_routers["adaptive"]) + assert _registered_adaptive_router_strategies() == survivor_strategies + survivor.discard() + + +@pytest.mark.asyncio +async def test_discard_releases_its_router_budget_limiter(): + def build_budget_limiter_then_discard() -> weakref.ReferenceType[object]: + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + budget_limiter: Final = router.get_router_deployment_budget_limiter() + assert budget_limiter is not None + router.discard() + return weakref.ref(budget_limiter) + + budget_limiter_ref: Final = build_budget_limiter_then_discard() + await asyncio.sleep(0) + gc.collect() + + 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) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 9b6700e9d39..cdff03af3b4 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -3891,6 +3891,114 @@ class TestCallbackAsyncSyncSeparation: assert my_sync_cb not in litellm._async_input_callback +class _DropsACallbackFromAnotherThreadWhenComparedToDynamodb: + def __init__(self, dropped: Callable[..., object]) -> None: + self.dropped = dropped + self.drop_thread: threading.Thread | None = None + + def __eq__(self, other: object) -> bool: + if isinstance(other, str) and other == "dynamodb" and self.drop_thread is None: + self.drop_thread = threading.Thread( + target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),) + ) + self.drop_thread.start() + self.drop_thread.join(timeout=0.2) + return self is other + + __hash__ = object.__hash__ + + +def test_function_setup_moves_async_callbacks_while_another_thread_drops_a_callback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def dropped_callback(*args: object, **kwargs: object) -> None: + return None + + async def async_callback(*args: object, **kwargs: object) -> None: + return None + + drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToDynamodb(dropped_callback) + monkeypatch.setattr(litellm.utils, "callback_list", ["initialized"]) + for list_name in ( + "callbacks", + "input_callback", + "failure_callback", + "service_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, list_name, []) + monkeypatch.setattr(litellm, "success_callback", [dropped_callback, drops_on_compare, async_callback]) + + litellm.utils.function_setup( + "completion", + litellm.utils.Rules(), + datetime.now(), + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="callback-drop-race", + ) + assert drops_on_compare.drop_thread is not None + drops_on_compare.drop_thread.join() + + assert litellm.success_callback == [drops_on_compare] + assert litellm._async_success_callback == [async_callback] + + +class _DropsACallbackFromAnotherThreadWhenComparedToSentry: + def __init__(self, dropped: Callable[..., object]) -> None: + self.dropped = dropped + self.drop_thread: threading.Thread | None = None + self.drop_finished_during_logger_setup = False + + def __eq__(self, other: object) -> bool: + if isinstance(other, str) and other == "sentry" and self.drop_thread is None: + self.drop_thread = threading.Thread( + target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),) + ) + self.drop_thread.start() + self.drop_thread.join(timeout=5) + self.drop_finished_during_logger_setup = not self.drop_thread.is_alive() + return self is other + + __hash__ = object.__hash__ + + +def test_function_setup_lets_other_threads_drop_callbacks_while_it_sets_up_loggers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def dropped_callback(*args: object, **kwargs: object) -> None: + return None + + drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToSentry(dropped_callback) + monkeypatch.setattr(litellm.utils, "callback_list", []) + for list_name in ( + "callbacks", + "input_callback", + "failure_callback", + "service_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, list_name, []) + monkeypatch.setattr(litellm, "success_callback", [drops_on_compare]) + + litellm.utils.function_setup( + "completion", + litellm.utils.Rules(), + datetime.now(), + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="logger-setup-outside-lock", + ) + assert drops_on_compare.drop_thread is not None + drops_on_compare.drop_thread.join() + + assert drops_on_compare.drop_finished_during_logger_setup + + class TestMetadataNoneHandling: """ Test that metadata=None in kwargs doesn't cause TypeError.