mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 285ceba838 into 85fdea0fd1
This commit is contained in:
commit
67435cda98
4 changed files with 208 additions and 6 deletions
|
|
@ -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__
|
||||
|
|
|
|||
|
|
@ -2301,13 +2301,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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -2259,3 +2259,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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue