scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(router): dedupe strategy selectors by identity so every Router's selector gets registered

Browse source 
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
This commit is contained in:
zhongerdeajie 2026-11-01 12:21:08 +08:00
parent 60fc121ce8
commit 285ceba838
4 changed files with 208 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

81
litellm/litellm_core_utils/logging_callback_manager.py
View file

@ -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__

11
litellm/router_python_core.py
View file

@ -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(

41
tests/unit/litellm_core_utils/test_logging_callback_manager.py
View file

@ -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

81
tests/unit/router_strategy/test_router_routing_groups.py
View file

@ -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