mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(proxy): reconcile budget reservations when no database logger
Without a prisma client, cost_tracking registered nothing while auth still reserved spend and bound the reservation to success callbacks. Nothing reconciled on success, so the in-memory reservation stuck until TTL and later requests failed with BudgetExceededError under custom_auth_run_common_checks. Register a success-only BudgetReservationReconcileLogger when prisma is absent so reservations settle to response_cost (or release on reconcile failure). Fixes #45895.
This commit is contained in:
parent
fb2a1450dd
commit
e6d9fa40cc
4 changed files with 226 additions and 3 deletions
61
litellm/proxy/hooks/budget_reservation_reconcile_logger.py
Normal file
61
litellm/proxy/hooks/budget_reservation_reconcile_logger.py
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
"""Reconcile in-memory budget reservations when no database logger is registered.
|
||||
|
||||
See https://github.com/BerriAI/litellm/issues/45895.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
budget_reservation_from_metadata,
|
||||
get_litellm_metadata_from_kwargs,
|
||||
)
|
||||
|
||||
class BudgetReservationReconcileLogger(CustomLogger):
|
||||
"""Settle auth's spend-counter reservation on success when ProxyDBLogger is absent.
|
||||
|
||||
With ``custom_auth_run_common_checks`` and no ``database_url``, auth still reserves
|
||||
spend in the in-memory counter and ``bind_budget_reservation_to_callbacks`` marks
|
||||
``callback_bound=True``. ``ProxyDBLogger`` (the only success path that reconciled)
|
||||
is registered only when ``prisma_client is not None``, so the reservation stuck
|
||||
until TTL and every later request on that key failed with ``BudgetExceededError``.
|
||||
"""
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
del response_obj, start_time, end_time
|
||||
metadata: Final = get_litellm_metadata_from_kwargs(kwargs=dict(kwargs))
|
||||
if not isinstance(metadata, Mapping):
|
||||
return
|
||||
budget_reservation: Final = budget_reservation_from_metadata(metadata)
|
||||
if not isinstance(budget_reservation, dict) or budget_reservation.get("finalized") is True:
|
||||
return
|
||||
|
||||
sl_object = kwargs.get("standard_logging_object")
|
||||
response_cost: object
|
||||
if isinstance(sl_object, Mapping):
|
||||
response_cost = sl_object.get("response_cost")
|
||||
else:
|
||||
response_cost = kwargs.get("response_cost")
|
||||
actual_cost: Final = float(response_cost or 0.0)
|
||||
|
||||
from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation
|
||||
|
||||
try:
|
||||
await reconcile_budget_reservation(
|
||||
budget_reservation=budget_reservation,
|
||||
actual_cost=actual_cost,
|
||||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
"Failed to reconcile budget reservation without a database logger; releasing instead"
|
||||
)
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
release_or_invalidate_budget_reservation,
|
||||
)
|
||||
|
||||
await release_or_invalidate_budget_reservation(budget_reservation=budget_reservation)
|
||||
|
|
@ -3326,6 +3326,17 @@ def cost_tracking():
|
|||
litellm.logging_callback_manager.add_litellm_callback(db_logger)
|
||||
litellm.logging_callback_manager.add_litellm_async_success_callback(db_logger)
|
||||
litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger())
|
||||
else:
|
||||
# No DB logger means nothing reconciled auth's in-memory reservation after
|
||||
# success (callback_bound=True skipped the request-end release). Register a
|
||||
# success-only reconciler so no-database custom_auth budgets do not stick. #45895
|
||||
from litellm.proxy.hooks.budget_reservation_reconcile_logger import (
|
||||
BudgetReservationReconcileLogger,
|
||||
)
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_async_success_callback(
|
||||
BudgetReservationReconcileLogger()
|
||||
)
|
||||
|
||||
|
||||
async def _drain_spend_event_producer_on_shutdown() -> None:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,146 @@
|
|||
"""Unit tests for BudgetReservationReconcileLogger (#45895)."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.hooks.budget_reservation_reconcile_logger import (
|
||||
BudgetReservationReconcileLogger,
|
||||
)
|
||||
|
||||
|
||||
def _reservation(callback_bound: bool = True) -> dict:
|
||||
return {
|
||||
"reserved_cost": 0.00255,
|
||||
"entries": [{"counter_key": "spend:key:hashed", "reserved_cost": 0.00255}],
|
||||
"finalized": False,
|
||||
"callback_bound": callback_bound,
|
||||
"input_cost": 0.0,
|
||||
"input_tokens": None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_reconciles_reservation_to_response_cost(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
reservation = _reservation()
|
||||
reconcile = AsyncMock(return_value=())
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
|
||||
reconcile,
|
||||
)
|
||||
logger = BudgetReservationReconcileLogger()
|
||||
|
||||
await logger.async_log_success_event(
|
||||
{
|
||||
"litellm_params": {
|
||||
"metadata": {"user_api_key_budget_reservation": reservation},
|
||||
},
|
||||
"response_cost": 0.0,
|
||||
},
|
||||
response_obj=object(),
|
||||
start_time=datetime.now(UTC),
|
||||
end_time=datetime.now(UTC),
|
||||
)
|
||||
|
||||
reconcile.assert_awaited_once()
|
||||
assert reconcile.await_args.kwargs["budget_reservation"] is reservation
|
||||
assert reconcile.await_args.kwargs["actual_cost"] == 0.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_uses_standard_logging_response_cost(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
reservation = _reservation()
|
||||
reconcile = AsyncMock(return_value=())
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
|
||||
reconcile,
|
||||
)
|
||||
logger = BudgetReservationReconcileLogger()
|
||||
|
||||
await logger.async_log_success_event(
|
||||
{
|
||||
"litellm_params": {
|
||||
"litellm_metadata": {"user_api_key_budget_reservation": reservation},
|
||||
},
|
||||
"standard_logging_object": {"response_cost": 0.001},
|
||||
"response_cost": 9.9, # ignored when SL payload present
|
||||
},
|
||||
response_obj=object(),
|
||||
start_time=datetime.now(UTC),
|
||||
end_time=datetime.now(UTC),
|
||||
)
|
||||
|
||||
assert reconcile.await_args.kwargs["actual_cost"] == 0.001
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_noop_when_no_reservation(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
reconcile = AsyncMock(return_value=())
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
|
||||
reconcile,
|
||||
)
|
||||
logger = BudgetReservationReconcileLogger()
|
||||
|
||||
await logger.async_log_success_event(
|
||||
{"litellm_params": {"metadata": {}}},
|
||||
response_obj=object(),
|
||||
start_time=datetime.now(UTC),
|
||||
end_time=datetime.now(UTC),
|
||||
)
|
||||
|
||||
reconcile.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_noop_when_already_finalized(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
reservation = _reservation()
|
||||
reservation["finalized"] = True
|
||||
reconcile = AsyncMock(return_value=())
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
|
||||
reconcile,
|
||||
)
|
||||
logger = BudgetReservationReconcileLogger()
|
||||
|
||||
await logger.async_log_success_event(
|
||||
{
|
||||
"litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}},
|
||||
"response_cost": 0.0,
|
||||
},
|
||||
response_obj=object(),
|
||||
start_time=datetime.now(UTC),
|
||||
end_time=datetime.now(UTC),
|
||||
)
|
||||
|
||||
reconcile.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_falls_back_to_release_when_reconcile_fails(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
reservation = _reservation()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
|
||||
AsyncMock(side_effect=RuntimeError("counter down")),
|
||||
)
|
||||
release = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.release_or_invalidate_budget_reservation",
|
||||
release,
|
||||
)
|
||||
logger = BudgetReservationReconcileLogger()
|
||||
|
||||
await logger.async_log_success_event(
|
||||
{
|
||||
"litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}},
|
||||
"response_cost": 0.0,
|
||||
},
|
||||
response_obj=object(),
|
||||
start_time=datetime.now(UTC),
|
||||
end_time=datetime.now(UTC),
|
||||
)
|
||||
|
||||
release.assert_awaited_once_with(budget_reservation=reservation)
|
||||
|
|
@ -776,18 +776,23 @@ async def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set(
|
|||
assert event.response_cost > 0
|
||||
|
||||
|
||||
def test_cost_tracking_no_op_when_prisma_missing(monkeypatch):
|
||||
"""Without a prisma_client cost_tracking is a no-op — not an error."""
|
||||
def test_cost_tracking_registers_budget_reconciler_when_prisma_missing(monkeypatch):
|
||||
"""Without a prisma_client, still register a success reconciler for in-memory reservations (#45895)."""
|
||||
import litellm
|
||||
from litellm.proxy.hooks.budget_reservation_reconcile_logger import (
|
||||
BudgetReservationReconcileLogger,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
||||
monkeypatch.setattr(litellm, "callbacks", [], raising=False)
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
|
||||
|
||||
cost_tracking()
|
||||
cost_tracking() # dedupe by logger class key
|
||||
|
||||
assert litellm.callbacks == []
|
||||
assert litellm._async_success_callback == []
|
||||
assert len(litellm._async_success_callback) == 1
|
||||
assert isinstance(litellm._async_success_callback[0], BudgetReservationReconcileLogger)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue