diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 3cc532def63..b5bafb1cc35 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -138,9 +138,8 @@ async def reserve_budget_for_request( ), ) except Exception: - await _set_reserved_entries_actual_cost( + await _release_applied_entries_best_effort( entries=applied_entries, - actual_cost=0.0, default_reserved_cost=reservation_cost, ) raise @@ -633,26 +632,101 @@ async def _set_reserved_entries_actual_cost( actual_cost: float, default_reserved_cost: float, ) -> None: - from litellm.proxy.proxy_server import _increment_spend_counter_cache - for entry in entries: - counter_key = entry.get("counter_key") - if counter_key is None: - continue - reserved_cost = _get_entry_reserved_cost( + await _set_reserved_entry_actual_cost( entry=entry, + actual_cost=actual_cost, default_reserved_cost=default_reserved_cost, ) - target_adjustment = actual_cost - reserved_cost - applied_adjustment = float(entry.get("applied_adjustment") or 0.0) - adjustment = target_adjustment - applied_adjustment - if adjustment == 0: - continue - await _increment_spend_counter_cache( - counter_key=counter_key, - increment=adjustment, + + +async def _set_reserved_entry_actual_cost( + entry: dict, + actual_cost: float, + default_reserved_cost: float, +) -> None: + from litellm.proxy.proxy_server import _increment_spend_counter_cache + + counter_key = entry.get("counter_key") + if counter_key is None: + return + reserved_cost = _get_entry_reserved_cost( + entry=entry, + default_reserved_cost=default_reserved_cost, + ) + target_adjustment = actual_cost - reserved_cost + applied_adjustment = float(entry.get("applied_adjustment") or 0.0) + adjustment = target_adjustment - applied_adjustment + if adjustment == 0: + return + await _ensure_counter_can_apply_adjustment( + counter_key=counter_key, + adjustment=adjustment, + ) + await _increment_spend_counter_cache( + counter_key=counter_key, + increment=adjustment, + ) + entry["applied_adjustment"] = target_adjustment + + +async def _ensure_counter_can_apply_adjustment( + counter_key: str, + adjustment: float, +) -> None: + from litellm.proxy.proxy_server import ( + _invalidate_spend_counter, + spend_counter_cache, + ) + + current_value = await spend_counter_cache.async_get_cache(key=counter_key) + if current_value is None: + await _invalidate_spend_counter(counter_key=counter_key) + raise RuntimeError( + f"Cannot apply budget reservation adjustment to missing counter {counter_key}" ) - entry["applied_adjustment"] = target_adjustment + + try: + current_float = float(current_value) + except (TypeError, ValueError): + await _invalidate_spend_counter(counter_key=counter_key) + raise RuntimeError( + f"Cannot apply budget reservation adjustment to non-numeric counter {counter_key}" + ) + + if adjustment < 0 and current_float + adjustment < -1e-12: + await _invalidate_spend_counter(counter_key=counter_key) + raise RuntimeError( + f"Budget reservation adjustment would make counter negative {counter_key}" + ) + + +async def _release_applied_entries_best_effort( + entries: List[dict], + default_reserved_cost: float, +) -> None: + for entry in entries: + try: + await _set_reserved_entry_actual_cost( + entry=entry, + actual_cost=0.0, + default_reserved_cost=default_reserved_cost, + ) + except Exception: + counter_key = entry.get("counter_key") + verbose_proxy_logger.exception( + "Failed to release partial budget reservation during exception cleanup" + ) + if counter_key is None: + continue + try: + from litellm.proxy.proxy_server import _invalidate_spend_counter + + await _invalidate_spend_counter(counter_key=counter_key) + except Exception: + verbose_proxy_logger.exception( + "Failed to invalidate partial budget reservation counter during exception cleanup" + ) async def _resize_applied_reservation( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1ae47d268f8..3b9454d5d80 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -978,6 +978,167 @@ async def test_should_retry_partial_release_without_double_decrement( ) == pytest.approx(0.0) +@pytest.mark.asyncio +async def test_should_preserve_budget_error_and_continue_partial_cleanup( + spend_counter_state, + monkeypatch, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-cleanup-failure", + spend=0.0, + max_budget=1.0, + team_id="team-budget-cleanup-failure", + ) + team_object = LiteLLM_TeamTable( + team_id="team-budget-cleanup-failure", + spend=0.0, + max_budget=0.3, + ) + + original_increment_cache = counter_cache.async_increment_cache + fail_key_cleanup = True + + async def flaky_increment_cache(key, value, *args, **kwargs): + nonlocal fail_key_cleanup + if key == "spend:key:key-budget-cleanup-failure" and value < 0: + if fail_key_cleanup: + fail_key_cleanup = False + raise RuntimeError("simulated cleanup failure") + return await original_increment_cache(key=key, value=value, *args, **kwargs) + + monkeypatch.setattr(counter_cache, "async_increment_cache", flaky_increment_cache) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.4, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.verbose_proxy_logger.exception" + ) as mock_log_exception, + ): + with pytest.raises(litellm.BudgetExceededError): + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-cleanup-failure" + ) + is None + ) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-budget-cleanup-failure" + ) == pytest.approx(0.0) + mock_log_exception.assert_called() + + +@pytest.mark.asyncio +async def test_should_not_create_negative_counter_when_release_counter_is_missing( + spend_counter_state, +): + counter_cache, _ = spend_counter_state + reservation = { + "reserved_cost": 0.4, + "entries": [ + { + "counter_key": "spend:key:key-budget-missing-release", + "reserved_cost": 0.4, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + with pytest.raises(RuntimeError, match="missing counter"): + await release_budget_reservation(reservation) + + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-missing-release" + ) + is None + ) + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_should_invalidate_counter_when_release_would_underflow( + spend_counter_state, +): + counter_cache, _ = spend_counter_state + await counter_cache.async_increment_cache( + key="spend:key:key-budget-underflow-release", + value=0.1, + ) + reservation = { + "reserved_cost": 0.4, + "entries": [ + { + "counter_key": "spend:key:key-budget-underflow-release", + "reserved_cost": 0.4, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + with pytest.raises(RuntimeError, match="negative"): + await release_budget_reservation(reservation) + + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-underflow-release" + ) + is None + ) + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_should_invalidate_non_numeric_counter_during_release( + spend_counter_state, +): + counter_cache, _ = spend_counter_state + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-nonnumeric-release", + value="stale", + ) + reservation = { + "reserved_cost": 0.4, + "entries": [ + { + "counter_key": "spend:key:key-budget-nonnumeric-release", + "reserved_cost": 0.4, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + with pytest.raises(RuntimeError, match="non-numeric"): + await release_budget_reservation(reservation) + + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-nonnumeric-release" + ) + is None + ) + assert reservation["finalized"] is False + + @pytest.mark.asyncio async def test_should_invalidate_reserved_counters_after_persisted_spend_failure( spend_counter_state,