scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-24 05:45:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

Merge de8856df73 into c58fc0c996

Browse source
This commit is contained in:
zhongerdeajie 2026-10-17 19:52:07 -04:00 • committed by GitHub
commit f967e7e604
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 309 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

113
litellm/litellm_core_utils/logging_callback_manager.py
View file

@ -1,5 +1,6 @@
from collections.abc import Callable
from collections.abc import Callable, Sequence, Set
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:
"""
@ -38,15 +47,30 @@ 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.
"""
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]):
"""
@ -54,15 +78,27 @@ class LoggingCallbackManager:
"""
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.
"""
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]):
@ -252,11 +288,20 @@ 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:
callback: the callback to add.
parent_list: the list to add the callback to.
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.
"""
# Check max callbacks limit first
if not self._check_callback_list_size(parent_list):
@ -273,6 +318,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):
@ -298,10 +344,38 @@ 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:
custom_logger: the logger to add.
parent_list: the list to add the logger to.
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__
@ -377,6 +451,35 @@ class LoggingCallbackManager:
):
self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self)
def remove_callbacks_by_id(self, selector_ids: Set[int]) -> None:
"""
Remove every callback whose object identity is in ``selector_ids`` from
each callback list it may have been promoted into.
Router strategy selectors are registered by identity (see
``add_litellm_callback(..., dedupe_by_identity=True)``), so discarding a
Router must drop them by ``id()`` rather than by class key.
"""
for callback_list in (
litellm.callbacks,
litellm.input_callback,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback, # pyright: ignore[reportPrivateUsage] # internal callback list promoted into by `function_setup`
litellm._async_failure_callback, # pyright: ignore[reportPrivateUsage] # internal callback list promoted into by `function_setup`
):
self._remove_callbacks_by_id_in_place(callback_list, selector_ids)
def _remove_callbacks_by_id_in_place(
self,
callback_list: Sequence[CustomLogger | Callable[..., object] | str],
selector_ids: Set[int],
) -> None:
if not isinstance(callback_list, list): # immutable sequence -> nothing to filter
return
for callback in [c for c in callback_list if id(c) in selector_ids]:
callback_list.remove(callback)
def get_active_additional_logging_utils_from_custom_logger(
self,
) -> set[AdditionalLoggingUtils]:

31
litellm/router.py
View file

@ -1623,6 +1623,16 @@ 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).
selectors: Final = [getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()]
selectors.extend(getattr(self, "_override_selectors", {}).values())
for group in getattr(self, "_group_selectors", {}).values():
selectors.extend(group.values())
self._unregister_router_selectors(selectors)
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)
@ -1743,20 +1753,28 @@ class Router:
@staticmethod
def _register_router_selector(selector: RouterStrategySelector) -> 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 would treat a second Router's selector as a duplicate
# of the first one and silently drop it, leaving the new Router's
# latency/usage cache permanently empty (issue #44575).
if isinstance(selector, LeastBusyLoggingHandler):
if isinstance(litellm.input_callback, list):
litellm.logging_callback_manager.add_litellm_input_callback(selector)
litellm.logging_callback_manager.add_litellm_input_callback(selector, dedupe_by_identity=True)
else:
litellm.input_callback = [selector]
if isinstance(litellm.callbacks, list):
litellm.logging_callback_manager.add_litellm_callback(selector)
litellm.logging_callback_manager.add_litellm_callback(selector, dedupe_by_identity=True)
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.
`_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):
@ -1764,10 +1782,7 @@ class Router:
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_callbacks_by_id(selector_ids)
def _apply_updated_routing_strategy_args(self) -> None:
"""

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

@ -342,3 +342,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

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

@ -2254,3 +2254,140 @@ 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
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