fix(proxy): share single _ProxyDBLogger instance across callback lists

cost_tracking() was creating two separate _ProxyDBLogger instances — one
for sync callbacks and one for async. This doubles DB connection pressure
and causes state divergence between the two loggers. Reuse a single
instance for both callback registrations.
This commit is contained in:
yryzhan 2026-05-20 15:28:48 +02:00
parent e59e34bed3
commit a5794a1ff0
2 changed files with 38 additions and 8 deletions

View file

@ -1892,10 +1892,9 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False):
def cost_tracking():
global prisma_client
if prisma_client is not None:
litellm.logging_callback_manager.add_litellm_callback(_ProxyDBLogger())
litellm.logging_callback_manager.add_litellm_async_success_callback(
_ProxyDBLogger()
)
db_logger = _ProxyDBLogger()
litellm.logging_callback_manager.add_litellm_callback(db_logger)
litellm.logging_callback_manager.add_litellm_async_success_callback(db_logger)
async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
@ -2710,11 +2709,9 @@ def run_ollama_serve():
with open(os.devnull, "w") as devnull:
subprocess.Popen(command, stdout=devnull, stderr=devnull)
except Exception as e:
verbose_proxy_logger.debug(
f"""
verbose_proxy_logger.debug(f"""
LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve`
"""
)
""")
def _get_process_rss_mb() -> Optional[float]:

View file

@ -7487,3 +7487,36 @@ class TestSortModelsByDisplayName:
all_models=models, sort_by="model_name", sort_order="asc"
)
assert [m["model_name"] for m in sorted_models] == ["alpha", "beta"]
class TestCostTrackingSingleDBLogger:
def test_cost_tracking_creates_single_db_logger_instance(self):
from unittest.mock import MagicMock, patch
from litellm.proxy.proxy_server import cost_tracking
mock_callback_manager = MagicMock()
with patch("litellm.proxy.proxy_server.prisma_client", new="fake-client"):
with patch("litellm.logging_callback_manager", mock_callback_manager):
cost_tracking()
add_callback_call = mock_callback_manager.add_litellm_callback.call_args[0][0]
add_async_call = (
mock_callback_manager.add_litellm_async_success_callback.call_args[0][0]
)
assert add_callback_call is add_async_call
def test_cost_tracking_noop_when_prisma_client_is_none(self):
from unittest.mock import MagicMock, patch
from litellm.proxy.proxy_server import cost_tracking
mock_callback_manager = MagicMock()
with patch("litellm.proxy.proxy_server.prisma_client", new=None):
with patch("litellm.logging_callback_manager", mock_callback_manager):
cost_tracking()
mock_callback_manager.add_litellm_callback.assert_not_called()
mock_callback_manager.add_litellm_async_success_callback.assert_not_called()