diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 3614bd397e9..d5ac694fa24 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -159,12 +159,20 @@ class PrometheusLogger(CustomLogger): @staticmethod def get_instance() -> PrometheusLogger | None: - """Find the PrometheusLogger instance from litellm.callbacks, if registered.""" + """The registered PrometheusLogger, however it was registered. + + Searching ``litellm.callbacks`` alone misses the equally supported + ``litellm_settings.success_callback: ["prometheus"]``, which lands the + logger on the success-callback lists instead. Callers of this method + publish metrics from outside the request hooks, so a miss shows up as a + metric that registers and then never leaves zero. + """ import litellm - for cb in litellm.callbacks: - if isinstance(cb, PrometheusLogger): - return cb + instances: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=PrometheusLogger) + for instance in instances: + if isinstance(instance, PrometheusLogger): + return instance return None def __init__( diff --git a/litellm/proxy/db/log_db_metrics.py b/litellm/proxy/db/log_db_metrics.py index a43ae9cbf49..1536af4cb0d 100644 --- a/litellm/proxy/db/log_db_metrics.py +++ b/litellm/proxy/db/log_db_metrics.py @@ -20,6 +20,10 @@ if TYPE_CHECKING: from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.db.db_pool_metrics import SupportsPoolSample +# Marks a P2024 that has already been counted, so nested decorated helpers do +# not each count the same timeout. +_POOL_TIMEOUT_COUNTED: Final = "_litellm_db_pool_timeout_counted" + # One sampler per process. The pool it reads is per-process too, so there is # nothing to key this by. _pool_metrics_sampler: Final = DBPoolMetricsSampler() @@ -66,15 +70,25 @@ async def _sample_db_pool_metrics() -> None: def _record_db_pool_timeout_if_exhausted(e: Exception) -> None: - """Count a pool exhaustion, without ever displacing the error that caused it. + """Count a pool exhaustion once, without ever displacing the error itself. - The caller re-raises the original exception after this returns. Anything - that escaped here would replace a P2024 with a metrics error, during the - incident this counter exists to record. + Decorated database helpers nest: ``get_object_permission`` is called from + inside ``get_key_object``, and both carry this decorator. One P2024 therefore + passes through several ``except`` blocks on its way up, so the exception is + marked the first time it is counted and skipped afterwards. + + The caller re-raises after this returns. Anything escaping here would replace + a P2024 with a metrics error, during the incident this counter exists to + record. """ try: if not PrismaDBExceptionHandler.is_connection_pool_timeout_error(e): return + if getattr(e, _POOL_TIMEOUT_COUNTED, False): + return + # rebind-ok: marking the in-flight exception is how an outer decorator + # learns this timeout was already counted by an inner one + setattr(e, _POOL_TIMEOUT_COUNTED, True) logger: Final = _prometheus_logger() if logger is not None: logger.record_db_pool_timeout() diff --git a/tests/test_litellm/integrations/test_prometheus_db_pool_metrics.py b/tests/test_litellm/integrations/test_prometheus_db_pool_metrics.py index e2bb433dcd0..b94aaf09484 100644 --- a/tests/test_litellm/integrations/test_prometheus_db_pool_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_db_pool_metrics.py @@ -225,3 +225,37 @@ def test_the_gauge_publishes_the_corrected_pending_not_the_raw_engine_reading(lo ) assert _value("litellm_db_pool_pending_acquirers") == 0.0 + + +def test_get_instance_finds_a_logger_registered_via_success_callback(logger): + """litellm_settings.success_callback: ["prometheus"] is an equally supported + registration and lands the logger on the success-callback lists rather than + litellm.callbacks. Searching only the latter left every pool metric + registered and permanently at zero, verified on a live proxy.""" + import litellm + from litellm.integrations.prometheus import PrometheusLogger + + saved_callbacks = litellm.callbacks + saved_success = litellm.success_callback + try: + litellm.callbacks = [] + litellm.success_callback = [logger] + assert PrometheusLogger.get_instance() is logger + finally: + litellm.callbacks = saved_callbacks + litellm.success_callback = saved_success + + +def test_get_instance_returns_none_when_prometheus_is_not_registered(logger): + import litellm + from litellm.integrations.prometheus import PrometheusLogger + + saved_callbacks = litellm.callbacks + saved_success = litellm.success_callback + try: + litellm.callbacks = [] + litellm.success_callback = [] + assert PrometheusLogger.get_instance() is None + finally: + litellm.callbacks = saved_callbacks + litellm.success_callback = saved_success diff --git a/tests/test_litellm/proxy/db/test_log_db_metrics.py b/tests/test_litellm/proxy/db/test_log_db_metrics.py index d9f51fec68e..84cdcf1ced6 100644 --- a/tests/test_litellm/proxy/db/test_log_db_metrics.py +++ b/tests/test_litellm/proxy/db/test_log_db_metrics.py @@ -1,7 +1,7 @@ import os import sys from datetime import datetime -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from prisma.errors import DataError, UniqueViolationError @@ -103,3 +103,60 @@ async def test_a_stalled_pool_sample_does_not_delay_the_database_call(): for task in asyncio.all_tasks(): if task is not asyncio.current_task() and task.get_coro().__name__ == "_hang": task.cancel() + + +@pytest.mark.asyncio +async def test_one_pool_timeout_is_counted_once_across_nested_decorated_calls(): + """get_object_permission carries @log_db_metrics and is called from inside + get_key_object, which also carries it, so a single P2024 passes through + several except blocks on its way up. Counting per block would inflate the + exhaustion metric by the nesting depth.""" + from litellm.proxy.db.log_db_metrics import log_db_metrics + + logger = MagicMock() + error = _pool_timeout_error() + + @log_db_metrics + async def inner(**kwargs): + raise error + + @log_db_metrics + async def outer(**kwargs): + return await inner(**kwargs) + + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + + with ( + patch("litellm.proxy.db.log_db_metrics._prometheus_logger", return_value=logger), + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj), + ): + with pytest.raises(DataError): + await outer(table_name="x") + + assert logger.record_db_pool_timeout.call_count == 1 + + +@pytest.mark.asyncio +async def test_two_separate_pool_timeouts_are_counted_separately(): + """The dedup must key on the exception instance, not suppress the metric.""" + from litellm.proxy.db.log_db_metrics import log_db_metrics + + logger = MagicMock() + + @log_db_metrics + async def failing(**kwargs): + raise _pool_timeout_error() + + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + + with ( + patch("litellm.proxy.db.log_db_metrics._prometheus_logger", return_value=logger), + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj), + ): + for _ in range(2): + with pytest.raises(DataError): + await failing(table_name="x") + + assert logger.record_db_pool_timeout.call_count == 2