From 285ceba838f3db5b9aaadbea8154083216ddd641 Mon Sep 17 00:00:00 2001 From: zhongerdeajie <1283667545@qq.com> Date: Sun, 1 Nov 2026 12:21:08 +0800 Subject: [PATCH] fix(router): dedupe strategy selectors by identity so every Router's selector gets registered A second Router with latency-based-routing in one process (notebook re-run, tests, multi-router SDK setups) silently lost its default selector: the LoggingCallbackManager dedupes CustomLoggers by a class-level key, and the latency handler's key has no distinguishing scalar attributes, so the second instance collided with the first one already in litellm.callbacks and was dropped. Calls made through the new Router then fed latency data into the old Router's handler/cache, leaving the new Router's `{model_group}_map` cache permanently empty and its routing blind (#44575). Add a `dedupe_by_identity` opt-in to the callback manager (identity-based membership for CustomLogger instances; duplicate adds of the same instance are still no-ops and MAX_CALLBACKS still applies) and use it in `_register_router_selector`, which is the only place per-instance strategy selectors are registered. Once an instance is registered by identity, later lists it is promoted into (`function_setup` copies `litellm.callbacks` into the success/failure event lists) keep identity mode too, or the second instance is dropped again there. Fixes #44575 --- .../logging_callback_manager.py | 81 ++++++++++++++++++- litellm/router_python_core.py | 11 ++- .../test_logging_callback_manager.py | 41 ++++++++++ .../test_router_routing_groups.py | 81 +++++++++++++++++++ 4 files changed, 208 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 56c6436eae6..bc5dfdc9a10 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -2,6 +2,7 @@ import threading from collections.abc import Callable, Iterable, Iterator, Mapping from types import MethodType from typing import TYPE_CHECKING, Final +from weakref import WeakSet import litellm from litellm._logging import verbose_logger @@ -19,6 +20,14 @@ else: _generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {} _callback_lists_lock: Final = threading.RLock() +# 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() + def _instance_attributes(target: object) -> Mapping[str, object]: return vars(target) @@ -45,16 +54,35 @@ class LoggingCallbackManager: except Exception: return False - def add_litellm_input_callback(self, callback: CustomLogger | str | Callable[..., object]): + def add_litellm_input_callback( + self, + callback: CustomLogger | str | Callable[..., object], + dedupe_by_identity: bool = False, + ): """ Add a input callback to litellm.input_callback. Auto-routes async callbacks to litellm._async_input_callback. + + Args: + callback: the callback to add. + dedupe_by_identity: if True, CustomLogger instances are deduped by + object identity instead of the class-level key. Required for + per-instance stateful loggers (router strategy selectors) where + two live instances of the same class must coexist. """ 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) + self._safe_add_callback_to_list( + callback=callback, + parent_list=litellm._async_input_callback, + dedupe_by_identity=dedupe_by_identity, + ) else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) + self._safe_add_callback_to_list( + callback=callback, + parent_list=litellm.input_callback, + dedupe_by_identity=dedupe_by_identity, + ) def add_litellm_service_callback(self, callback: CustomLogger | str | Callable[..., object]): """ @@ -63,16 +91,28 @@ class LoggingCallbackManager: 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]): + def add_litellm_callback( + self, + callback: CustomLogger | str | Callable[..., object], + dedupe_by_identity: bool = False, + ): """ Add a callback to litellm.callbacks Ensures no duplicates are added. + + Args: + callback: the callback to add. + dedupe_by_identity: if True, CustomLogger instances are deduped by + object identity instead of the class-level key. Required for + per-instance stateful loggers (router strategy selectors) where + two live instances of the same class must coexist. """ with _callback_lists_lock: self._safe_add_callback_to_list( callback=callback, parent_list=litellm.callbacks, + dedupe_by_identity=dedupe_by_identity, ) def add_litellm_success_callback(self, callback: CustomLogger | str | Callable[..., object]): @@ -268,11 +308,17 @@ class LoggingCallbackManager: self, callback: CustomLogger | Callable[..., object] | str, parent_list: list[CustomLogger | Callable[..., object] | str], + dedupe_by_identity: bool = False, ): """ Safe add a callback to a list, if the callback is already in the list, do not add it again. Ensures no duplicates are added for `str`, `Callable`, and `CustomLogger` callbacks. + + Args: + dedupe_by_identity: forwarded to `_add_custom_logger_to_list` — dedupe + ``CustomLogger`` instances by object identity instead of the + class-level key. """ # Check max callbacks limit first if not self._check_callback_list_size(parent_list): @@ -289,6 +335,7 @@ class LoggingCallbackManager: self._add_custom_logger_to_list( custom_logger=callback, parent_list=parent_list, + dedupe_by_identity=dedupe_by_identity, ) elif callable(callback): @@ -314,10 +361,36 @@ class LoggingCallbackManager: self, custom_logger: CustomLogger, parent_list: list[CustomLogger | Callable[..., object] | str], + dedupe_by_identity: bool = False, ): """ Add a custom logger to a list, if another instance of the same custom logger exists in the list, do not add it again. + + Args: + dedupe_by_identity: if True, dedupe by object identity instead of + the class-level key. Required for per-instance stateful loggers + (router strategy selectors) where two live instances of the same + class must coexist. """ + # 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). 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: + verbose_logger.debug( + "Custom logger instance %s already exists in %s, not adding again..", + custom_logger, + parent_list, + ) + return + # Check if an instance of the same class already exists in the list custom_logger_key: Final = self._get_custom_logger_key(custom_logger) custom_logger_type_name: Final = type(custom_logger).__name__ diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 30b05456987..68bcb586074 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -2275,13 +2275,20 @@ class PythonCore(DeprecatedRouterMembers): return selector def _register_router_selector(self, selector: CustomLogger) -> None: + # Strategy selectors are per-instance stateful callbacks (each holds + # its own Router cache), so they must be deduped by identity: the + # class-level key treats a second Router's selector as a duplicate of + # the first one and silently drops it, leaving the new Router's + # latency/usage cache permanently empty (issue #44575). Least-busy + # never hit this by luck — its instances differ in a scalar attribute, + # so their keys never collided. if isinstance(selector, LeastBusyLoggingHandler): if isinstance(self._host.callback_list("input_callback"), list): - self._host.callbacks.add_litellm_input_callback(selector) + self._host.callbacks.add_litellm_input_callback(selector, dedupe_by_identity=True) else: self._host.set_callback_list("input_callback", [selector]) if isinstance(self._host.callback_list("callbacks"), list): - self._host.callbacks.add_litellm_callback(selector) + self._host.callbacks.add_litellm_callback(selector, dedupe_by_identity=True) def _callback_selectors(self) -> tuple[RouterStrategySelector, ...]: return tuple( 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 25a52137d65..5d26072f1ab 100644 --- a/tests/unit/litellm_core_utils/test_logging_callback_manager.py +++ b/tests/unit/litellm_core_utils/test_logging_callback_manager.py @@ -344,3 +344,44 @@ def test_duplicate_multiple_loggers_test(): assert langfuse_count == 1, "Should have exactly one LangfusePromptManagement instance" assert otel_count == 1, "Should have exactly one OpenTelemetry instance" + + +def test_add_custom_logger_default_mode_dedupes_same_class_instances(): + """ + Default mode keeps one instance per CustomLogger class: the class-level + key treats a second instance of the same class as a duplicate. + """ + manager = LoggingCallbackManager() + + class TestLogger(CustomLogger): + pass + + first, second = TestLogger(), TestLogger() + + manager.add_litellm_callback(first) + manager.add_litellm_callback(second) + + assert sum(1 for c in litellm.callbacks if c is first) == 1 + assert all(c is not second for c in litellm.callbacks) + + +def test_add_custom_logger_dedupe_by_identity_allows_two_instances(): + """ + `dedupe_by_identity=True` is required for per-instance stateful loggers + (router strategy selectors): two live instances of the same class must + coexist, while re-adding the same instance stays a no-op. + """ + manager = LoggingCallbackManager() + + class TestLogger(CustomLogger): + pass + + first, second = TestLogger(), TestLogger() + + manager.add_litellm_callback(first, dedupe_by_identity=True) + manager.add_litellm_callback(second, dedupe_by_identity=True) + assert sum(1 for c in litellm.callbacks if c is first) == 1 + assert sum(1 for c in litellm.callbacks if c is second) == 1 + + manager.add_litellm_callback(first, dedupe_by_identity=True) + assert sum(1 for c in litellm.callbacks if c is first) == 1 diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 5eeb46c3518..deb604f3f1f 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -2255,3 +2255,84 @@ async def test_update_settings_changed_routing_strategy_args_flushes_replaced_se await t except asyncio.CancelledError: 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 + default selector registered: the latency handler's class-level dedup key is + identical for every instance (no scalar instance attributes), so key-based + dedup silently dropped the second Router's selector and its latency cache + never got populated (issue #44575). + """ + _reset_global_callback_lists(monkeypatch) + + first = _build_router(routing_strategy="latency-based-routing") + second = _build_router(routing_strategy="latency-based-routing") + + first_selector = first.lowestlatency_logger + second_selector = second.lowestlatency_logger + assert first_selector is not None and second_selector is not None + assert first_selector is not second_selector + assert any(c is first_selector for c in litellm.callbacks) + assert any(c is second_selector for c in litellm.callbacks), ( + "second Router's latency selector was dropped by the callback dedup" + ) + + +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