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.