diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 9f521ffc6cc..b3adede7e59 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -363,13 +363,14 @@ 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.input_callback, litellm.success_callback, litellm.failure_callback, litellm._async_success_callback, diff --git a/litellm/router.py b/litellm/router.py index 99adc6c5fa6..cac156f6773 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -37,6 +37,7 @@ from collections.abc import ( 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 @@ -1622,6 +1623,13 @@ class Router: litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.input_callback, self) litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.service_callback, self) litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, self) + self._unregister_router_selectors( + ( + *(getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()), + *self._override_selectors.values(), + *chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()), + ) + ) # Remove ForwardClientSideHeadersByModelGroup if it exists if self.optional_callbacks is not None: @@ -1746,20 +1754,17 @@ class Router: 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 selector is None: + continue 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.remove_callback_from_all_lists(selector) def _apply_updated_routing_strategy_args(self) -> None: """ 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/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 4ac917d706d..df07528a7cb 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1,5 +1,8 @@ +import asyncio +import gc import json import os +import weakref from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -10,7 +13,9 @@ 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.custom_logger import CustomLogger from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.types.litellm_params import RoutingStrategyName from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( ModelResponse, @@ -1526,6 +1531,108 @@ 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() + + +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) + + def test_initialize_assistants_endpoint(model_list): """Test if the 'initialize_assistants_endpoint' function is working correctly""" router = Router(model_list=model_list)