Merge 60b43b108a into 82ca10cc2d

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,

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:
         """

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

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)