diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 86f7849bb1..56b471c290 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -489,11 +489,12 @@ async def _update_database_and_spend_counters( await _invalidate_budget_reservation_counters( budget_reservation=budget_reservation ) - budget_reservation["finalized"] = True except Exception: verbose_proxy_logger.exception( "Failed to invalidate budget reservation counters after spend counter update failed" ) + finally: + budget_reservation["finalized"] = True raise diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2f2b59db67..dd46f09fbc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1993,8 +1993,9 @@ async def _increment_end_user_and_tag_spend_counters( reserved_counter_keys: Set[str], ) -> None: if end_user_id is not None: - await _increment_warm_unreserved_spend_counter( + await _init_and_increment_unreserved_spend_counter( counter_key=f"spend:end_user:{end_user_id}", + source_cache_key=f"end_user_id:{end_user_id}", increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) @@ -2007,26 +2008,14 @@ async def _increment_end_user_and_tag_spend_counters( if not tag_name or not isinstance(tag_name, str) or tag_name in seen_tags: continue seen_tags.add(tag_name) - await _increment_warm_unreserved_spend_counter( + await _init_and_increment_unreserved_spend_counter( counter_key=f"spend:tag:{tag_name}", + source_cache_key=f"tag:{tag_name}", increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) -async def _increment_warm_unreserved_spend_counter( - counter_key: str, - increment: float, - reserved_counter_keys: Set[str], -) -> None: - if counter_key in reserved_counter_keys: - return - if await spend_counter_cache.async_get_cache(key=counter_key) is None: - return - - await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) - - async def _increment_org_spend_counter( org_id: Optional[str], response_cost: float, diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index b17251f4b5..3cc532def6 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -97,7 +97,10 @@ async def reserve_budget_for_request( applied_entries: List[Dict[str, Any]] = [] try: for counter in counters: - entry = _counter_to_reservation_entry(counter) + entry = _counter_to_reservation_entry( + counter=counter, + reserved_cost=reservation_cost, + ) reserved_value = await _reserve_counter( counter=counter, reservation_cost=reservation_cost, @@ -135,9 +138,10 @@ async def reserve_budget_for_request( ), ) except Exception: - await _set_reserved_entries_adjustment( + await _set_reserved_entries_actual_cost( entries=applied_entries, - target_adjustment=-reservation_cost, + actual_cost=0.0, + default_reserved_cost=reservation_cost, ) raise @@ -157,10 +161,10 @@ async def reconcile_budget_reservation( reserved_cost = float(budget_reservation.get("reserved_cost") or 0.0) actual = float(actual_cost or 0.0) - adjustment = actual - reserved_cost - await _set_reserved_entries_adjustment( + await _set_reserved_entries_actual_cost( entries=budget_reservation.get("entries") or [], - target_adjustment=adjustment, + actual_cost=actual, + default_reserved_cost=reserved_cost, ) budget_reservation["finalized"] = True @@ -624,9 +628,10 @@ async def _ensure_malformed_window_counter_initialized( ) -async def _set_reserved_entries_adjustment( +async def _set_reserved_entries_actual_cost( entries: List[dict], - target_adjustment: float, + actual_cost: float, + default_reserved_cost: float, ) -> None: from litellm.proxy.proxy_server import _increment_spend_counter_cache @@ -634,6 +639,11 @@ async def _set_reserved_entries_adjustment( counter_key = entry.get("counter_key") if counter_key is None: continue + 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: @@ -650,23 +660,33 @@ async def _resize_applied_reservation( current_reserved_cost: float, new_reserved_cost: float, ) -> None: - await _set_reserved_entries_adjustment( + await _set_reserved_entries_actual_cost( entries=entries, - target_adjustment=new_reserved_cost - current_reserved_cost, + actual_cost=new_reserved_cost, + default_reserved_cost=current_reserved_cost, ) - for entry in entries: - entry["applied_adjustment"] = 0.0 -def _counter_to_reservation_entry(counter: _BudgetCounter) -> Dict[str, Any]: +def _counter_to_reservation_entry( + counter: _BudgetCounter, + reserved_cost: float, +) -> Dict[str, Any]: return { "counter_key": counter.counter_key, "entity_type": counter.entity_type, "entity_id": counter.entity_id, + "reserved_cost": reserved_cost, "applied_adjustment": 0.0, } +def _get_entry_reserved_cost(entry: dict, default_reserved_cost: float) -> float: + try: + return float(entry.get("reserved_cost", default_reserved_cost) or 0.0) + except (TypeError, ValueError): + return default_reserved_cost + + def get_budget_window_start(window: Any) -> Optional[datetime]: window_dict = _coerce_window(window) budget_duration = window_dict.get("budget_duration") diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 351d3059d4..01cfd1708a 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -500,6 +500,7 @@ async def test_update_database_and_spend_counters_preserves_counter_exception_wh mock_log_exception.assert_called_once_with( "Failed to invalidate budget reservation counters after spend counter update failed" ) + assert budget_reservation["finalized"] is True proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 51016126d3..1ae47d268f 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -245,6 +245,7 @@ async def test_should_prevent_second_tag_reservation_over_budget( "counter_key": "spend:tag:tag-budget-race", "entity_type": "Tag", "entity_id": "tag-budget-race", + "reserved_cost": 0.6, "applied_adjustment": 0.0, } ] @@ -290,16 +291,33 @@ async def test_should_prevent_second_tag_reservation_over_budget( @pytest.mark.asyncio -async def test_should_update_warm_end_user_and_tag_counters_without_reservation( +async def test_should_seed_and_update_end_user_and_tag_counters_without_reservation( spend_counter_state, ): - counter_cache, _ = spend_counter_state - counter_cache.in_memory_cache.set_cache( - key="spend:end_user:customer-1", - value=4.0, + counter_cache, key_cache = spend_counter_state + await key_cache.async_set_cache( + key="end_user_id:customer-1", + value=LiteLLM_EndUserTable( + user_id="customer-1", + blocked=False, + spend=4.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=10.0), + ).model_dump(), + ) + await key_cache.async_set_cache( + key="tag:paid-tag", + value=LiteLLM_TagTable( + tag_name="paid-tag", + spend=7.0, + ).model_dump(), + ) + await key_cache.async_set_cache( + key="tag:other-tag", + value=LiteLLM_TagTable( + tag_name="other-tag", + spend=2.0, + ).model_dump(), ) - counter_cache.in_memory_cache.set_cache(key="spend:tag:paid-tag", value=7.0) - counter_cache.in_memory_cache.set_cache(key="spend:tag:other-tag", value=2.0) from litellm.proxy.proxy_server import increment_spend_counters @@ -539,6 +557,82 @@ async def test_should_shrink_uncapped_reservation_when_counter_advances( ) == pytest.approx(0.3) +@pytest.mark.asyncio +async def test_should_shrink_uncapped_reservation_multiple_times( + 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-double-resize", + spend=0.2, + max_budget=1.0, + team_id="team-budget-double-resize", + ) + team_object = LiteLLM_TeamTable( + team_id="team-budget-double-resize", + spend=0.2, + max_budget=1.0, + ) + request_body = _request_body() + request_body.pop("max_tokens") + + from litellm.proxy.spend_tracking import budget_reservation + + stale_spend_by_counter_key = { + "spend:key:key-budget-double-resize": 0.3, + "spend:team:team-budget-double-resize": 0.4, + } + + async def stale_counter_read(counter): + await counter_cache.async_increment_cache( + key=counter.counter_key, + value=stale_spend_by_counter_key[counter.counter_key], + ) + return 0.2 + + monkeypatch.setattr( + budget_reservation, + "_get_current_counter_value", + stale_counter_read, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=None, + ): + reservation = 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 reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.6) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-double-resize" + ) == pytest.approx(0.9) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-budget-double-resize" + ) == pytest.approx(1.0) + + await release_budget_reservation(reservation) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-double-resize" + ) == pytest.approx(0.3) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-budget-double-resize" + ) == pytest.approx(0.4) + + def test_should_start_window_without_reset_at_at_duration_boundary(): before = datetime.now(timezone.utc) - timedelta(hours=1)