mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
harden partial budget reservation cleanup
This commit is contained in:
user
parent
694fadd175
commit
38ebd4de3d
2 changed files with 252 additions and 17 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue