tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-07 22:21:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix budget reservation greptile findings

Browse Source
This commit is contained in:
user
2026-04-30 17:08:45 -07:00
parent f8d187785d
commit fce86d1334
5 changed files with 141 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/hooks/proxy_track_cost_callback.py
+2 -1
View File
@@ -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
litellm/proxy/proxy_server.py
+4 -15
View File
@@ -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,
litellm/proxy/spend_tracking/budget_reservation.py
+33 -13
View File
@@ -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")
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
+1
View File
@@ -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()
tests/test_litellm/proxy/test_budget_reservation.py
+101 -7
View File
@@ -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)