mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
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:
parent
e59e34bed3
commit
a5794a1ff0
2 changed files with 38 additions and 8 deletions
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue