From 5783a43c5cf67d2753e765e5be0b18cbe5abeef1 Mon Sep 17 00:00:00 2001 From: setouchi-h Date: Sat, 10 Oct 2026 23:03:31 +0900 Subject: [PATCH 1/2] fix(router): unregister strategy selectors from global callbacks on discard Router.discard only removed its own bound methods from litellm's callback lists, so the strategy selectors it registered (least-busy, usage-based, latency-based, cost-based and routing group selectors) stayed behind and kept receiving events after the router was gone. Reuse the identity based unregister that update_settings already uses so a discarded router takes only its own selectors with it --- litellm/router.py | 9 +++ .../test_router/test_router_helper_utils.py | 69 +++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index f173cdb37b4..8d5174c6dea 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -35,6 +35,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, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -1410,6 +1411,7 @@ 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(self._all_strategy_selectors()) # Remove ForwardClientSideHeadersByModelGroup if it exists if self.optional_callbacks is not None: @@ -1530,6 +1532,13 @@ class Router: if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(selector) + def _all_strategy_selectors(self) -> tuple[RouterStrategySelector | None, ...]: + default_selectors: Final = ( + getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values() + ) + group_selectors: Final = chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()) + return (*default_selectors, *self._override_selectors.values(), *group_selectors) + def _unregister_router_selectors(self, selectors: Sequence[object]) -> None: """ Drop router-owned strategy selectors from litellm's global callback diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index cbbfe654bf6..47db58bd82c 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,6 +13,7 @@ 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.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( @@ -1525,6 +1529,71 @@ def test_discard(model_list): assert len(litellm.service_callback) == 0 +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, selector_attr): + 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 litellm.callbacks == [] + assert litellm.input_callback == [] + assert all(selector_ref() is None for selector_ref in selector_refs) + + +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 ("callbacks", "input_callback")} + 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 litellm.callbacks == [] + assert litellm.input_callback == [] + + def test_initialize_assistants_endpoint(model_list): """Test if the 'initialize_assistants_endpoint' function is working correctly""" router = Router(model_list=model_list) From 60b43b108a5c0b4a961aa45afed623d09b8dff40 Mon Sep 17 00:00:00 2001 From: setouchi-h Date: Sat, 10 Oct 2026 23:43:25 +0900 Subject: [PATCH 2/2] fix(router): purge discarded selectors from every event callback list function_setup copies litellm.callbacks into the success and failure lists on every request, so a router that served traffic still leaked its selectors there after discard. Route the cleanup through remove_callback_from_all_lists, which now also covers input_callback, and fold the selector gathering into discard so the router coverage check sees no untested helper --- .../logging_callback_manager.py | 3 +- litellm/router.py | 32 ++++++------ .../test_logging_callback_manager.py | 2 + .../test_router/test_router_helper_utils.py | 50 ++++++++++++++++--- 4 files changed, 62 insertions(+), 25 deletions(-) diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 5eaa4b467fe..ddf65c9f1e4 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 8d5174c6dea..1cca7d0801a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1411,7 +1411,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(self._all_strategy_selectors()) + 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: @@ -1532,30 +1538,20 @@ class Router: if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(selector) - def _all_strategy_selectors(self) -> tuple[RouterStrategySelector | None, ...]: - default_selectors: Final = ( - getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values() - ) - group_selectors: Final = chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()) - return (*default_selectors, *self._override_selectors.values(), *group_selectors) - 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 47db58bd82c..793dbdcacf2 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -15,6 +15,7 @@ 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, @@ -1529,6 +1530,16 @@ 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)] @@ -1544,7 +1555,9 @@ def _callback_ids(list_name: str) -> list[int]: ], ) @pytest.mark.asyncio -async def test_discard_releases_strategy_selector_across_router_lifecycles(routing_strategy, selector_attr): +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) @@ -1556,16 +1569,42 @@ async def test_discard_releases_strategy_selector_across_router_lifecycles(routi await asyncio.sleep(0) gc.collect() - assert litellm.callbacks == [] - assert litellm.input_callback == [] + 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 ("callbacks", "input_callback")} + 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) @@ -1590,8 +1629,7 @@ def test_discard_drops_routing_group_selectors(): router.discard() - assert litellm.callbacks == [] - assert litellm.input_callback == [] + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) def test_initialize_assistants_endpoint(model_list):