mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
fix(proxy): raise live budget counters for failure spend rows
A stream that breaks mid-flight writes its recovered cost to the spend log, but the failure hook had already released the reservation and never touched the live counters, so an over-budget key kept passing until the counter was next repaired from the database Every charged failure row with a cost now raises the same counters the success path raises. The error waits at most a second for them, so a hung Redis cannot hold it back Fixes #35568
This commit is contained in:
yucheng-berri
parent
e333ccabdb
commit
5485c90ffc
4 changed files with 540 additions and 9 deletions
|
|
@ -12,6 +12,7 @@ from pydantic import TypeAdapter, ValidationError
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.batches.batch_utils import batch_cost_is_final
|
||||
from litellm.caching.redis_batch import flush_post_call_redis_batches
|
||||
from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
|
|
@ -126,6 +127,8 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
|||
)
|
||||
)
|
||||
|
||||
FAILURE_SPEND_COUNTER_WAIT_SECONDS: Final = 1.0
|
||||
|
||||
|
||||
def _proxy_stamped_used_client_oauth_token(
|
||||
request_data: Mapping[str, object], request_route: str | None
|
||||
|
|
@ -189,6 +192,12 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter:
|
|||
return proxy_logging_obj.db_spend_update_writer
|
||||
|
||||
|
||||
def _proxy_spend_counters() -> "_IncrementSpendCounters":
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
return increment_spend_counters
|
||||
|
||||
|
||||
DeploymentLookup: TypeAlias = Callable[[str], Deployment | None]
|
||||
|
||||
|
||||
|
|
@ -315,11 +324,16 @@ class _ProxyDBLogger(CustomLogger):
|
|||
message_logging: bool = True,
|
||||
spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer,
|
||||
deployment_lookup: DeploymentLookup = _proxy_deployment,
|
||||
spend_counters: Callable[[], "_IncrementSpendCounters"] = _proxy_spend_counters,
|
||||
failure_counter_wait_seconds: float = FAILURE_SPEND_COUNTER_WAIT_SECONDS,
|
||||
) -> None:
|
||||
super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging)
|
||||
self.spend_event_producer = spend_event_producer
|
||||
self._spend_writer: Final = spend_writer
|
||||
self._deployment_lookup: Final = deployment_lookup
|
||||
self._spend_counters: Final = spend_counters
|
||||
self._failure_counter_wait_seconds: Final = failure_counter_wait_seconds
|
||||
self._pending_increments: Final[set[asyncio.Task[None]]] = set() # mutable-ok: tasks leave when done
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime
|
||||
|
|
@ -504,13 +518,16 @@ class _ProxyDBLogger(CustomLogger):
|
|||
start_time: datetime,
|
||||
spend_rows: Sequence[tuple[float, Mapping[str, object]]],
|
||||
) -> None:
|
||||
"""One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair."""
|
||||
for row_cost, row_kwargs in spend_rows:
|
||||
"""One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair, then the
|
||||
live spend counters raised by every charged row that cost something."""
|
||||
token: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
|
||||
end_user_id: Final = None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id
|
||||
charged: Final = [
|
||||
await self._spend_writer().update_database(
|
||||
token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
|
||||
token=token,
|
||||
response_cost=row_cost,
|
||||
user_id=user_api_key_dict.user_id,
|
||||
end_user_id=None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id,
|
||||
end_user_id=end_user_id,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
kwargs=dict(row_kwargs),
|
||||
completion_response=original_exception,
|
||||
|
|
@ -519,6 +536,91 @@ class _ProxyDBLogger(CustomLogger):
|
|||
org_id=user_api_key_dict.org_id,
|
||||
project_id=user_api_key_dict.project_id,
|
||||
)
|
||||
for row_cost, row_kwargs in spend_rows
|
||||
]
|
||||
counted_rows: Final = tuple(
|
||||
row for row, row_charged in zip(spend_rows, charged, strict=True) if row_charged and row[0] > 0
|
||||
)
|
||||
if counted_rows and not ProxyUpdateSpend.disable_spend_updates():
|
||||
await self._increment_failure_spend_counters(
|
||||
user_api_key_dict, token, end_user_id, counted_rows, start_time
|
||||
)
|
||||
|
||||
async def _increment_failure_spend_counters(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
token: str | None,
|
||||
end_user_id: str | None,
|
||||
counted_rows: Sequence[tuple[float, Mapping[str, object]]],
|
||||
start_time: datetime,
|
||||
) -> None:
|
||||
"""The counters the success path raises for the same spend, sent before the client sees the error. A
|
||||
counter store slower than ``failure_counter_wait_seconds`` finishes them in the background instead of
|
||||
holding the error back. The reservation was already released, so every counter takes the full row cost."""
|
||||
increment: Final = asyncio.create_task(
|
||||
self._raise_failure_spend_counters(user_api_key_dict, token, end_user_id, counted_rows, start_time)
|
||||
)
|
||||
self._pending_increments.add(increment)
|
||||
increment.add_done_callback(self._pending_increments.discard)
|
||||
finished, _ = await asyncio.wait((increment,), timeout=self._failure_counter_wait_seconds)
|
||||
if not finished:
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend counters for failure spend rows took over %ss, finishing them in the background",
|
||||
self._failure_counter_wait_seconds,
|
||||
)
|
||||
|
||||
async def _raise_failure_spend_counters(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
token: str | None,
|
||||
end_user_id: str | None,
|
||||
counted_rows: Sequence[tuple[float, Mapping[str, object]]],
|
||||
start_time: datetime,
|
||||
) -> None:
|
||||
for row_cost, row_kwargs in counted_rows:
|
||||
await self._raise_row_spend_counters(
|
||||
user_api_key_dict, token, end_user_id, row_cost, row_kwargs, start_time
|
||||
)
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
async def _raise_row_spend_counters(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
token: str | None,
|
||||
end_user_id: str | None,
|
||||
row_cost: float,
|
||||
row_kwargs: Mapping[str, object],
|
||||
start_time: datetime,
|
||||
) -> None:
|
||||
standard_logging_object: Final = row_kwargs.get("standard_logging_object")
|
||||
sl_object: Final = (
|
||||
cast("StandardLoggingPayload", standard_logging_object) # cast-ok: the failure snapshot dict
|
||||
if isinstance(standard_logging_object, dict)
|
||||
else None
|
||||
)
|
||||
metadata: Final = _str_keyed_fields(_str_keyed_fields(row_kwargs.get("litellm_params")).get("metadata"))
|
||||
served_model_id: Final = _str_keyed_fields(metadata.get("model_info")).get("id") or _str_keyed_fields(
|
||||
standard_logging_object
|
||||
).get("model_id")
|
||||
try:
|
||||
await self._spend_counters()(
|
||||
token=token,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
user_id=user_api_key_dict.user_id,
|
||||
response_cost=row_cost,
|
||||
org_id=user_api_key_dict.org_id,
|
||||
end_user_id=end_user_id,
|
||||
tags=_get_request_tags_for_cost_tracking(sl_object=sl_object, metadata=dict(metadata)),
|
||||
request_started_at=start_time,
|
||||
model_access_groups=debitable_model_access_groups(
|
||||
attributed=get_request_model_access_groups(row_kwargs),
|
||||
served_model_id=served_model_id if isinstance(served_model_id, str) else None,
|
||||
router=get_llm_router(),
|
||||
),
|
||||
project_id=user_api_key_dict.project_id,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing
|
||||
verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row")
|
||||
|
||||
async def _PROXY_track_cost_callback(
|
||||
self,
|
||||
|
|
@ -915,7 +1017,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]:
|
|||
return tuple(sorted(str(key) for key in metadata))
|
||||
|
||||
|
||||
def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None:
|
||||
def _get_budget_reservation_from_metadata(metadata: dict) -> dict[str, object] | None:
|
||||
return budget_reservation_from_metadata(metadata)
|
||||
|
||||
|
||||
|
|
@ -950,6 +1052,7 @@ class _IncrementSpendCounters(Protocol):
|
|||
tags: list[str] | None = None,
|
||||
request_started_at: datetime | None = None,
|
||||
model_access_groups: Sequence[str] | None = None,
|
||||
project_id: str | None = None,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
|
|
@ -966,7 +1069,7 @@ async def _update_database_and_spend_counters(
|
|||
start_time: datetime | None,
|
||||
end_time: datetime | None,
|
||||
response_cost: float,
|
||||
budget_reservation: dict | None,
|
||||
budget_reservation: dict[str, object] | None,
|
||||
request_tags: list[str] | None = None,
|
||||
model_access_groups: Sequence[str] | None = None,
|
||||
project_id: str | None = None,
|
||||
|
|
@ -1027,7 +1130,7 @@ async def _update_database_and_spend_counters_in_batch(
|
|||
start_time: datetime | None,
|
||||
end_time: datetime | None,
|
||||
response_cost: float,
|
||||
budget_reservation: dict | None,
|
||||
budget_reservation: dict[str, object] | None,
|
||||
request_tags: list[str] | None,
|
||||
model_access_groups: Sequence[str] | None,
|
||||
project_id: str | None,
|
||||
|
|
|
|||
|
|
@ -3414,7 +3414,7 @@ async def increment_spend_counters(
|
|||
user_id: str | None,
|
||||
response_cost: float | None,
|
||||
org_id: str | None = None,
|
||||
budget_reservation: dict | None = None,
|
||||
budget_reservation: dict[str, object] | None = None,
|
||||
end_user_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
request_started_at: datetime | None = None,
|
||||
|
|
|
|||
195
tests/integration/spend/test_partial_stream_live_budget.py
Normal file
195
tests/integration/spend/test_partial_stream_live_budget.py
Normal file
|
|
@ -0,0 +1,195 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterable, Iterator
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, eventually, held
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai.types.chat import ChatCompletionChunk
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_BACKEND: Final = "integration-partial-stream"
|
||||
_PIECE: Final = "partial stream billed tokens "
|
||||
_DELIVERED_PIECES: Final = 12
|
||||
_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"'
|
||||
|
||||
|
||||
class _SpendRow(BaseModel):
|
||||
spend: float
|
||||
|
||||
|
||||
class _BudgetError(BaseModel):
|
||||
type: str
|
||||
|
||||
|
||||
class _ErrorResponse(BaseModel):
|
||||
error: _BudgetError
|
||||
|
||||
|
||||
def _chunk(delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes:
|
||||
body: Final = {
|
||||
"id": "chatcmpl-partial",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": _BACKEND,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
||||
}
|
||||
return f"data: {json.dumps(body)}\n\n".encode()
|
||||
|
||||
|
||||
def _completion() -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": 1700000000,
|
||||
"model": _BACKEND,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
def _peer(delivered: int) -> Callable[[Request], Reply]:
|
||||
frames: Final = (
|
||||
*(_chunk({"role": "assistant", "content": _PIECE}) for _ in range(_DELIVERED_PIECES)),
|
||||
_chunk({}, "stop"),
|
||||
b"data: [DONE]\n\n",
|
||||
)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method != "POST":
|
||||
return Reply(status=404)
|
||||
if _JSON_OBJECT.validate_json(request.body).get("stream") is True:
|
||||
return Reply(
|
||||
content_type="text/event-stream", chunks=frames, abort_after=delivered, pause_between_chunks=0.02
|
||||
)
|
||||
return Reply(body=_completion())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _priced_model(scenario: Scenario, wire: Wire) -> str:
|
||||
return scenario.model(
|
||||
model=f"openai/{_BACKEND}",
|
||||
api_base=f"{wire.url}/v1",
|
||||
api_key="synthetic-openai-key",
|
||||
input_cost_per_token=0.001,
|
||||
output_cost_per_token=0.002,
|
||||
)
|
||||
|
||||
|
||||
def _stream(gateway: Gateway, model: str, key: str) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"max_tokens": 5,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": "partial stream"}],
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _broken_stream(gateway: Gateway, model: str, key: str) -> str:
|
||||
response: Final = _stream(gateway, model, key)
|
||||
assert response.status_code == 200, response.text
|
||||
events: Final = tuple(
|
||||
line.removeprefix("data: ")
|
||||
for line in response.text.splitlines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
assert "error" in _JSON_OBJECT.validate_json(events[-1]), response.text
|
||||
return "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1]))
|
||||
|
||||
|
||||
def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]:
|
||||
for chunk in chunks:
|
||||
yield from (choice.delta.content or "" for choice in chunk.choices)
|
||||
|
||||
|
||||
def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}]},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]:
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70)
|
||||
found: Final = held(
|
||||
lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2
|
||||
)
|
||||
return tuple(_SpendRow.model_validate(row).spend for row in found)
|
||||
|
||||
|
||||
def _chats_sent(wire: Wire) -> int:
|
||||
return sum(1 for request in wire.drain() if request.method == "POST")
|
||||
|
||||
|
||||
def _recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float:
|
||||
key: Final = scenario.key(models=[model], max_budget=100)
|
||||
assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES
|
||||
(spend,) = _failure_spend(gateway, key)
|
||||
assert spend > 0
|
||||
return spend
|
||||
|
||||
|
||||
def _assert_budget_refusal(response: httpx.Response) -> None:
|
||||
assert response.status_code == 422, response.text
|
||||
assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(_DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario, wire)
|
||||
cost: Final = _recovered_cost(gateway, scenario, model)
|
||||
key: Final = scenario.key(models=[model], max_budget=cost / 2)
|
||||
_chats_sent(wire)
|
||||
|
||||
assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES
|
||||
after: Final = _chat(gateway, model, key)
|
||||
|
||||
_assert_budget_refusal(after)
|
||||
assert _failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0))
|
||||
assert _chats_sent(wire) == 1
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(_DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario, wire)
|
||||
cost: Final = _recovered_cost(gateway, scenario, model)
|
||||
key: Final = scenario.key(models=[model], max_budget=cost * 1.5)
|
||||
|
||||
assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES
|
||||
assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES
|
||||
after: Final = _chat(gateway, model, key)
|
||||
|
||||
_assert_budget_refusal(after)
|
||||
assert _failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0))
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(0)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario, wire)
|
||||
key: Final = scenario.key(models=[model], max_budget=0.05)
|
||||
|
||||
refused: Final = _stream(gateway, model, key)
|
||||
after: Final = _chat(gateway, model, key)
|
||||
|
||||
assert refused.status_code == 500, refused.text
|
||||
|
||||
assert after.status_code == 200, after.text
|
||||
assert _failure_spend(gateway, key) == (0.0,)
|
||||
|
|
@ -11,6 +11,7 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
|
||||
from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth
|
||||
from litellm.proxy.collector import SpendEventConsumer
|
||||
|
|
@ -26,7 +27,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import (
|
|||
run_spend_event,
|
||||
)
|
||||
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
|
||||
from litellm.proxy.utils import ProxyUpdateSpend
|
||||
from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend
|
||||
from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event
|
||||
from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
|
|
@ -2550,6 +2551,238 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend():
|
|||
assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def live_spend_counters(monkeypatch: pytest.MonkeyPatch) -> DualCache:
|
||||
counters: Final = DualCache()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", counters)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
return counters
|
||||
|
||||
|
||||
def _charging_spend_writer() -> DBSpendUpdateWriter:
|
||||
writer: Final = create_autospec(DBSpendUpdateWriter, instance=True)
|
||||
writer.update_database.return_value = True
|
||||
return writer
|
||||
|
||||
|
||||
def _failure_with_billed_attempts() -> dict:
|
||||
usage: Final = Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
|
||||
return {
|
||||
"combined_usage_object": usage,
|
||||
"response_cost": 0.5,
|
||||
"billed_failed_attempts": (
|
||||
BilledFailedAttempt(attempted_retries=0, deployment_id="dep-a", usage=usage, response_cost=0.25),
|
||||
BilledFailedAttempt(attempted_retries=1, deployment_id="dep-b", usage=usage, response_cost=0.5),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
_LIVE_COUNTERS: Final = {
|
||||
"spend:key:hashed-key": 0.2,
|
||||
"spend:team:team-1": 0.3,
|
||||
"spend:team_member:user-1:team-1": 0.35,
|
||||
"spend:user:user-1": 0.4,
|
||||
"spend:end_user:end-user-1": 0.45,
|
||||
"spend:org:org-1": 0.5,
|
||||
"spend:project:project-1": 0.55,
|
||||
"spend:tag:batch-jobs": 0.6,
|
||||
"spend:model_access_group:premium": 0.65,
|
||||
}
|
||||
|
||||
|
||||
async def _fail_request(logger: ProxyDBLogger, failure: dict) -> None:
|
||||
await logger.async_post_call_failure_hook(
|
||||
request_data={
|
||||
"model": "gpt-4.1",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"metadata": {"tags": ["batch-jobs"], MODEL_ACCESS_GROUP_METADATA_KEY: ["premium"]},
|
||||
"proxy_server_request": {"request_id": "rid"},
|
||||
**failure,
|
||||
},
|
||||
original_exception=Exception("stream broke mid-flight"),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="hashed-key",
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
end_user_id="end-user-1",
|
||||
org_id="org-1",
|
||||
project_id="project-1",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _counter_values(counters: DualCache) -> dict[str, float]:
|
||||
return {key: await counters.async_get_cache(key) for key in _LIVE_COUNTERS}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("failure", "charged"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70),
|
||||
"response_cost": 0.125,
|
||||
},
|
||||
0.125,
|
||||
),
|
||||
(_failure_with_billed_attempts(), 1.25),
|
||||
({}, 0.0),
|
||||
],
|
||||
)
|
||||
async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
|
||||
live_spend_counters: DualCache, failure: dict, charged: float
|
||||
):
|
||||
"""#35568: budget checks read the live counters, so a failure row's cost has to land there as well as in the
|
||||
spend log, or an over-budget key keeps passing until the counter is next repaired from the database."""
|
||||
for key, spend in _LIVE_COUNTERS.items():
|
||||
await live_spend_counters.async_set_cache(key, spend)
|
||||
|
||||
await _fail_request(ProxyDBLogger(spend_writer=_charging_spend_writer), failure)
|
||||
|
||||
assert await _counter_values(live_spend_counters) == pytest.approx(
|
||||
{key: spend + charged for key, spend in _LIVE_COUNTERS.items()}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("spend_updates_disabled", "writer_charged"), [(True, True), (False, False)])
|
||||
async def test_failure_spend_rows_leave_the_live_spend_counters_alone_when_nothing_was_charged(
|
||||
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch, spend_updates_disabled: bool, writer_charged: bool
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings", {"disable_spend_updates": True} if spend_updates_disabled else {}
|
||||
)
|
||||
writer: Final = create_autospec(DBSpendUpdateWriter, instance=True)
|
||||
writer.update_database.return_value = writer_charged
|
||||
for key, spend in _LIVE_COUNTERS.items():
|
||||
await live_spend_counters.async_set_cache(key, spend)
|
||||
|
||||
await _fail_request(ProxyDBLogger(spend_writer=lambda: writer), _failure_with_billed_attempts())
|
||||
|
||||
assert await _counter_values(live_spend_counters) == pytest.approx(_LIVE_COUNTERS)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_spend_counter_outage_still_writes_every_failure_spend_row():
|
||||
writer: Final = _charging_spend_writer()
|
||||
counters: Final = AsyncMock(side_effect=RuntimeError("redis unavailable"))
|
||||
|
||||
await _fail_request(
|
||||
ProxyDBLogger(spend_writer=lambda: writer, spend_counters=lambda: counters), _failure_with_billed_attempts()
|
||||
)
|
||||
|
||||
assert [call.kwargs["response_cost"] for call in writer.update_database.call_args_list] == [0.5, 0.25, 0.5]
|
||||
assert [call.kwargs["response_cost"] for call in counters.call_args_list] == [0.5, 0.25, 0.5]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_slow_spend_counter_store_lets_the_failure_out_and_raises_the_counters_once_it_answers():
|
||||
store_answers: Final = asyncio.Event()
|
||||
landed: Final = asyncio.Event()
|
||||
raised_costs: Final[list[object]] = []
|
||||
|
||||
async def slow_counter_store(**counter_args: object) -> None:
|
||||
await store_answers.wait()
|
||||
raised_costs.append(counter_args["response_cost"])
|
||||
landed.set()
|
||||
|
||||
logger: Final = ProxyDBLogger(
|
||||
spend_writer=_charging_spend_writer,
|
||||
spend_counters=lambda: slow_counter_store,
|
||||
failure_counter_wait_seconds=0.05,
|
||||
)
|
||||
|
||||
await asyncio.wait_for(
|
||||
_fail_request(
|
||||
logger,
|
||||
{
|
||||
"combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70),
|
||||
"response_cost": 0.125,
|
||||
},
|
||||
),
|
||||
timeout=1,
|
||||
)
|
||||
assert raised_costs == []
|
||||
|
||||
store_answers.set()
|
||||
await asyncio.wait_for(landed.wait(), timeout=1)
|
||||
assert raised_costs == [0.125]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_request_start(
|
||||
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""The window row's one-time seed from the spend log stops at the request start, which keeps the failure row
|
||||
it already logged from being counted twice."""
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
window_queue: Final = proxy_logging_obj.db_spend_update_writer.window_spend_update_queue
|
||||
await window_queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
key_cache: Final = DualCache()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", key_cache)
|
||||
reset_at: Final = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
await key_cache.async_set_cache(
|
||||
"hashed-key",
|
||||
UserAPIKeyAuth(
|
||||
token="hashed-key",
|
||||
budget_limits=[{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}],
|
||||
),
|
||||
)
|
||||
request_start: Final = datetime(2026, 10, 20, 12, 0, tzinfo=timezone.utc)
|
||||
|
||||
await _fail_request(
|
||||
ProxyDBLogger(spend_writer=_charging_spend_writer),
|
||||
{
|
||||
"combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70),
|
||||
"response_cost": 0.125,
|
||||
"start_time": request_start,
|
||||
},
|
||||
)
|
||||
|
||||
assert [
|
||||
(row["entity_type"], row["entity_id"], row["spend"], row["started_at"])
|
||||
for row in await window_queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
] == [("key", "hashed-key", pytest.approx(0.125), "2026-10-20T12:00:00.000000")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_attributed_access_group(
|
||||
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""Its spend row names the dropped deployment, which narrows nothing, so the counters must not narrow it to the
|
||||
final attempt's groups either."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.llm_router",
|
||||
litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-4.1",
|
||||
"litellm_params": {"model": "openai/gpt-4.1", "api_key": "sk-test"},
|
||||
"model_info": {"id": "dep-b", "access_groups": ["pool"]},
|
||||
}
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
await _fail_request(
|
||||
ProxyDBLogger(spend_writer=_charging_spend_writer),
|
||||
{
|
||||
**_failure_with_billed_attempts(),
|
||||
"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool", "premium"]},
|
||||
"standard_logging_object": {"id": "rid", "model_id": "dep-b"},
|
||||
},
|
||||
)
|
||||
|
||||
assert {
|
||||
group: await live_spend_counters.async_get_cache(f"spend:model_access_group:{group}")
|
||||
for group in ("premium", "pool")
|
||||
} == pytest.approx({"premium": 0.25, "pool": 1.25})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata():
|
||||
"""MCP tool calls may only carry user_api_key; user/team rollups still need user_id."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue