From 491aceaaf534cb6362139c8170bdcd1a01d00ab2 Mon Sep 17 00:00:00 2001 From: zhongerdeajie <1283667545@qq.com> Date: Sat, 10 Oct 2026 13:01:17 +0800 Subject: [PATCH] fix(router): carry identity-based selector dedup into the event callback lists The previous commit registered each Router's strategy selector by identity in `litellm.callbacks`, but normal requests still dropped the second one: `function_setup` copies the general callbacks into the success/failure (and async) event lists through manager methods that deduped by class key, so a second latency Router's selector never reached the dispatch and its cache stayed empty in real traffic. Record loggers registered with `dedupe_by_identity` in a WeakSet on the callback manager and keep identity semantics whenever they are promoted into another list, so both routers' selectors receive events. Make `Router.discard()` unregister the selectors it registered (default, group and override), and extend `_unregister_router_selectors` to the event lists, so a throwaway per-request `user_config` Router cannot leak a selector per request until MAX_CALLBACKS silently blocks later callback registration. Drop the real clock from the new regression test and drive the second-router case through `function_setup` plus the success dispatch instead of calling the selectors directly. Co-Authored-By: Claude Code --- .../logging_callback_manager.py | 17 ++- litellm/router.py | 29 +++- .../test_router_routing_groups.py | 124 ++++++++++++++++-- 3 files changed, 152 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 614810feae0..34d7a7edaf5 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -1,5 +1,6 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Final +from weakref import WeakSet import litellm from litellm._logging import verbose_logger @@ -16,6 +17,14 @@ else: _generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {} +# CustomLoggers registered with `dedupe_by_identity=True` (router strategy +# selectors) keep identity-based membership when another code path re-adds them +# to a different list. `function_setup` copies `litellm.callbacks` into the +# success/failure event lists, and the class-key dedup there would otherwise +# drop a second live instance again, leaving the newest Router's cache empty +# (issue #44575). Weakly held so a discarded logger does not leak. +_identity_deduped_loggers: Final[WeakSet[CustomLogger]] = WeakSet() + class LoggingCallbackManager: """ @@ -351,8 +360,12 @@ class LoggingCallbackManager: # Per-instance loggers are deduped by identity: their class-level key # does not distinguish two live instances, and silently dropping the # second one leaves the newest Router with an unregistered callback - # (see issue #44575). - if dedupe_by_identity: + # (see issue #44575). Once a logger is registered by identity, every + # later list it is promoted into (e.g. `function_setup` copying + # `litellm.callbacks` into the success/failure event lists) must keep + # that mode, or the second instance is dropped again there. + if dedupe_by_identity or any(existing is custom_logger for existing in _identity_deduped_loggers): + _identity_deduped_loggers.add(custom_logger) if not any(existing is custom_logger for existing in parent_list): parent_list.append(custom_logger) else: diff --git a/litellm/router.py b/litellm/router.py index 22a56fb2c1b..263afe4e4eb 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1403,6 +1403,20 @@ class Router: # Stop contributing to cost-map rebuilds straight away rather than waiting # for this router to be collected. _live_routers.discard(self) + # Drop the per-instance strategy selectors this router registered. They + # are registered by identity, so class-key dedup no longer implicitly + # caps them at one per class: without this a discarded router (e.g. the + # throwaway `user_config` Router built per request) leaks its selectors + # into every callback list for the life of the process (issue #44575). + self._unregister_router_selectors( + [getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()] + + list(getattr(self, "_override_selectors", {}).values()) + + [ + selector + for selectors in getattr(self, "_group_selectors", {}).values() + for selector in selectors.values() + ] + ) 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) @@ -1539,8 +1553,11 @@ class Router: """ 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. + `_init_routing_groups`) and on `discard`, so repeated `update_settings` + calls (or discarded routers) don't accumulate dead selectors that keep + receiving callback events. Covers the success/failure event lists too: + `function_setup` promotes general callbacks into them, and the + per-instance selectors are no longer capped at one per class. """ for selector in selectors: if isinstance(selector, BaseRoutingStrategy): @@ -1552,6 +1569,14 @@ class Router: 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] + if isinstance(litellm.success_callback, list): + litellm.success_callback = [c for c in litellm.success_callback if id(c) not in selector_ids] + if isinstance(litellm.failure_callback, list): + litellm.failure_callback = [c for c in litellm.failure_callback if id(c) not in selector_ids] + if isinstance(litellm._async_success_callback, list): + litellm._async_success_callback = [c for c in litellm._async_success_callback if id(c) not in selector_ids] + if isinstance(litellm._async_failure_callback, list): + litellm._async_failure_callback = [c for c in litellm._async_failure_callback if id(c) not in selector_ids] def _apply_updated_routing_strategy_args(self) -> None: """ diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index dcb806fc641..e6544d15343 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -2256,6 +2256,20 @@ async def test_update_settings_changed_routing_strategy_args_flushes_replaced_se pass +def _reset_global_callback_lists(monkeypatch) -> None: + """Every global callback list, so a Router built here does not see selectors + promoted by an earlier test (monkeypatch would not restore those).""" + for name in ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, name, []) + + def test_two_routers_same_latency_strategy_both_register_selectors(monkeypatch): """ Two Routers with `latency-based-routing` in one process must both see their @@ -2264,8 +2278,7 @@ def test_two_routers_same_latency_strategy_both_register_selectors(monkeypatch): dedup silently dropped the second Router's selector and its latency cache never got populated (issue #44575). """ - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) + _reset_global_callback_lists(monkeypatch) first = _build_router(routing_strategy="latency-based-routing") second = _build_router(routing_strategy="latency-based-routing") @@ -2279,19 +2292,102 @@ def test_two_routers_same_latency_strategy_both_register_selectors(monkeypatch): "second Router's latency selector was dropped by the callback dedup" ) - # Each selector feeds its own Router's cache: an event delivered through - # the global callback list must land in the emitting router's cache. - kwargs = { - "litellm_params": { - "metadata": {"model_group": "filtered-model"}, - "model_info": {"id": "deploy-1"}, - } - } - start_time = datetime.datetime.now() - for callback in litellm.callbacks: - if isinstance(callback, CustomLogger): - callback.log_success_event(kwargs, None, start_time, start_time + datetime.timedelta(seconds=0.5)) +def test_second_latency_router_receives_events_through_the_normal_logging_path(monkeypatch): + """ + Registering both selectors in `litellm.callbacks` is not enough: the normal + request path (`function_setup`) copies those callbacks into the success and + failure event lists, and that copy used to dedupe by class key and drop the + second Router's selector again. Drive `function_setup` and the success + dispatch here instead of calling the selectors directly (issue #44575). + """ + _reset_global_callback_lists(monkeypatch) + + first = _build_router(routing_strategy="latency-based-routing") + second = _build_router(routing_strategy="latency-based-routing") + second_selector = second.lowestlatency_logger + + start_time = datetime.datetime(2026, 1, 1) + logging_obj, _ = function_setup( + "completion", + Rules(), + start_time, + litellm_call_id=str(uuid.uuid4()), + messages=[{"role": "user", "content": "hi"}], + model="filtered-model", + metadata={"model_group": "filtered-model"}, + model_info={"id": "deploy-1"}, + ) + + # The promotion into the event lists is where the second instance used to be + # dropped, which the `litellm.callbacks` assertion above cannot see. + assert any(cb is second_selector for cb in litellm.success_callback), ( + "function_setup dropped the second Router's selector from the success event list" + ) + assert any(cb is second_selector for cb in litellm.failure_callback) + assert any(cb is second_selector for cb in litellm._async_success_callback) + + logging_obj.success_handler( + result=None, + start_time=start_time, + end_time=start_time + datetime.timedelta(seconds=0.5), + ) assert first.cache.get_cache("filtered-model_map") is not None assert second.cache.get_cache("filtered-model_map") is not None + + +def test_discard_unregisters_only_the_discarded_routers_selectors(monkeypatch): + """ + A discarded Router must take its identity-registered selectors (default, + group and override) out of every global callback list, or a per-request + `user_config` Router leaks a selector per request until MAX_CALLBACKS blocks + every later guardrail registration. A second, still-live Router keeps working. + """ + _reset_global_callback_lists(monkeypatch) + + first = _build_router( + routing_strategy="latency-based-routing", + routing_groups=[ + { + "group_name": "fast", + "models": ["filtered-model"], + "routing_strategy": "latency-based-routing", + } + ], + ) + second = _build_router(routing_strategy="latency-based-routing") + discarded_selectors = [ + first.lowestlatency_logger, + first._group_selectors["fast"]["latency-based-routing"], + ] + live_selector = second.lowestlatency_logger + + start_time = datetime.datetime(2026, 1, 1) + logging_obj, _ = function_setup( + "completion", + Rules(), + start_time, + litellm_call_id=str(uuid.uuid4()), + messages=[{"role": "user", "content": "hi"}], + model="filtered-model", + metadata={"model_group": "filtered-model"}, + model_info={"id": "deploy-1"}, + ) + assert all(any(cb is selector for cb in litellm.success_callback) for selector in discarded_selectors) + + first.discard() + + for selector in discarded_selectors: + assert _selector_is_not_global(selector), "discard left a selector in a global callback list" + assert any(cb is live_selector for cb in litellm.callbacks) + assert any(cb is live_selector for cb in litellm.success_callback) + + # The live Router still receives success events through the normal path. + logging_obj.success_handler( + result=None, + start_time=start_time, + end_time=start_time + datetime.timedelta(seconds=0.5), + ) + assert second.cache.get_cache("filtered-model_map") is not None +