scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

test(proxy): await background spend counters before asserting (#51249)

Browse source 
* test(proxy): wait for child readiness before supervisor cancellation

* test(proxy): sync supervisor cancellation test on a socketpair handshake

Replace the readiness marker poll and timed sleeps with one-byte socketpair
messages. The child reports its SIGTERM handler is installed, reports the
stop, and holds its exit until the parent releases it, so the assertion only
passes when the cancelled supervisor waits for the child. A pytest-timeout
marker bounds a deadlock.

* test(proxy): drive supervisor cancellation cleanup with a controlled child

Replace the real child and socket handshake with a scripted process returned
by asyncio.create_subprocess_exec and a stop observed through os.killpg. The
child only exits after the stop and a checkpoint, so returncode is 0 only when
the cancelled supervisor stops the child and waits for it without being
cancelled.

* test(proxy): start the supervisor cancellation only once the reporter waits

The scripted child signals when a wait begins, and the test cancels only after
the exit reporter is blocked on the stop, so cancellation reaches it without
relying on task start order.

* test(proxy): await background failure spend counters before asserting

Failure spend counters finish in a background task when the counter store is
slower than the wait budget, so the tests read stale values on a slow runner.
Hold the counter store on a release event with a zero wait budget, capture the
background task, and await it before asserting the counters.

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri 2026-11-11 16:26:58 -08:00 • committed by GitHub
parent 0608a41a07
commit 072290e78f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 110 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

135
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -1,11 +1,11 @@
import asyncio
import json
import logging
from collections.abc import Mapping
from collections.abc import Awaitable, Callable, Mapping
from datetime import datetime, timedelta, timezone
from pathlib import Path
from types import MappingProxyType, SimpleNamespace
from typing import Final, cast
from typing import Final, ParamSpec, cast
from unittest.mock import AsyncMock, MagicMock, create_autospec, patch
import httpx
@ -29,6 +29,7 @@ from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction
from litellm.proxy.db.spend_log_tool_index import response_tool_call_names
from litellm.proxy.hooks.proxy_track_cost_callback import (
ProxyDBLogger,
_IncrementSpendCounters,
_charge_request_cost,
_get_budget_reservation_from_metadata,
_should_track_cost_callback,
@ -2970,7 +2971,42 @@ async def _counter_values(counters: SpendCounterStore) -> dict[str, float | None
return {key: counters.get_local(key) for key in _LIVE_COUNTERS}
_P = ParamSpec("_P")
def _held_until(
release: asyncio.Event, counting: asyncio.Future[asyncio.Task[object]], store: Callable[_P, Awaitable[None]]
) -> Callable[_P, Awaitable[None]]:
async def held(*args: _P.args, **kwargs: _P.kwargs) -> None:
if (task := asyncio.current_task()) is not None and not counting.done():
counting.set_result(task)
await release.wait()
await store(*args, **kwargs)
return held
async def _fail_request_and_finish_counting_in_background(
store: _IncrementSpendCounters, failure: Mapping[str, object], committed: CommittedBudgetSpend | None = None
) -> None:
release: Final = asyncio.Event()
counting: Final[asyncio.Future[asyncio.Task[object]]] = asyncio.get_running_loop().create_future()
held: Final = _held_until(release, counting, store)
try:
await _fail_request(
ProxyDBLogger(spend_counters=lambda: held, failure_counter_wait_seconds=0), failure, committed
)
background: Final = await counting
release.set()
await background
finally:
release.set()
if counting.done():
await asyncio.wait((counting.result(),))
@pytest.mark.asyncio
@pytest.mark.timeout(30)
@pytest.mark.parametrize(
("failure", "charged"),
[
@ -2982,7 +3018,6 @@ async def _counter_values(counters: SpendCounterStore) -> dict[str, float | None
0.125,
),
(_failure_with_billed_attempts(), 1.25),
({}, 0.0),
],
)
async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
@ -2990,10 +3025,12 @@ async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
):
"""#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."""
from litellm.proxy.proxy_server import increment_spend_counters
for key, spend in _LIVE_COUNTERS.items():
live_spend_counters.set_local(key, spend)
await _fail_request(ProxyDBLogger(), failure)
await _fail_request_and_finish_counting_in_background(increment_spend_counters, failure)
assert await _counter_values(live_spend_counters) == pytest.approx(
{key: spend + charged for key, spend in _LIVE_COUNTERS.items()}
@ -3001,39 +3038,83 @@ async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
@pytest.mark.asyncio
async def test_a_failure_that_cost_nothing_leaves_the_live_spend_counters_alone(live_spend_counters: SpendCounterStore):
for key, spend in _LIVE_COUNTERS.items():
live_spend_counters.set_local(key, spend)
await _fail_request(ProxyDBLogger(), {})
assert await _counter_values(live_spend_counters) == pytest.approx(_LIVE_COUNTERS)
_COMMITTED_COUNTERS: Final = MappingProxyType(
{
"spend:committed:key:hashed-key": 0.2,
"spend:committed:team:team-1": 0.3,
"spend:committed:user:user-1": 0.4,
}
)
_BUDGETED_KEY_AND_TEAM: Final = CommittedBudgetSpend(
counter_keys=EntitySpendCounterKeys.of(
token="hashed-key", team_id="team-1", user_id=None, org_id=None, end_user_id=None
)
)
def _committed_counter_values(counters: SpendCounterStore) -> dict[str, float | None]:
return {key: counters.get_local(key) for key in _COMMITTED_COUNTERS}
def _committed_counters_after_charging_key_and_team(charged: float) -> dict[str, float]:
return {
"spend:committed:key:hashed-key": 0.2 + charged,
"spend:committed:team:team-1": 0.3 + charged,
"spend:committed:user:user-1": 0.4,
}
@pytest.mark.asyncio
@pytest.mark.timeout(30)
@pytest.mark.parametrize(
("failure", "charged"),
[
({"combined_usage_object": Usage(prompt_tokens=30), "response_cost": 0.125}, 0.125),
(_failure_with_billed_attempts(), 1.25),
({"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}}, 0.125),
],
ids=["recovered_stream_cost", "billed_retries", "completed_cost"],
ids=["recovered_stream_cost", "billed_retries"],
)
async def test_a_failed_charge_moves_only_the_committed_counters_auth_named_as_budgeted(
live_spend_counters: SpendCounterStore, failure: Mapping[str, object], charged: float
):
"""The remaining budget gauges read the committed counters, so a charge missing from them shows as budget the
request never left behind, until the next reseed from the database."""
budgeted: Final = EntitySpendCounterKeys.of(
token="hashed-key", team_id="team-1", user_id=None, org_id=None, end_user_id=None
)
starting: Final = {
"spend:committed:key:hashed-key": 0.2,
"spend:committed:team:team-1": 0.3,
"spend:committed:user:user-1": 0.4,
}
for key, spend in starting.items():
from litellm.proxy.proxy_server import increment_spend_counters
for key, spend in _COMMITTED_COUNTERS.items():
live_spend_counters.set_local(key, spend)
await _fail_request(ProxyDBLogger(), failure, CommittedBudgetSpend(counter_keys=budgeted))
await _fail_request_and_finish_counting_in_background(increment_spend_counters, failure, _BUDGETED_KEY_AND_TEAM)
assert {key: live_spend_counters.get_local(key) for key in starting} == pytest.approx(
{
"spend:committed:key:hashed-key": 0.2 + charged,
"spend:committed:team:team-1": 0.3 + charged,
"spend:committed:user:user-1": 0.4,
}
assert _committed_counter_values(live_spend_counters) == pytest.approx(
_committed_counters_after_charging_key_and_team(charged)
)
@pytest.mark.asyncio
async def test_a_failed_completed_charge_moves_only_the_committed_counters_auth_named_as_budgeted(
live_spend_counters: SpendCounterStore,
):
for key, spend in _COMMITTED_COUNTERS.items():
live_spend_counters.set_local(key, spend)
await _fail_request(
ProxyDBLogger(),
{"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}},
_BUDGETED_KEY_AND_TEAM,
)
assert _committed_counter_values(live_spend_counters) == pytest.approx(
_committed_counters_after_charging_key_and_team(0.125)
)
@ -3067,10 +3148,11 @@ async def test_failure_spend_rows_leave_the_live_spend_counters_alone_while_spen
@pytest.mark.asyncio
@pytest.mark.timeout(30)
async def test_a_failing_spend_counter_store_is_still_sent_every_failure_spend_row():
counters: Final = AsyncMock(side_effect=RuntimeError("redis unavailable"))
await _fail_request(ProxyDBLogger(spend_counters=lambda: counters), _failure_with_billed_attempts())
await _fail_request_and_finish_counting_in_background(counters, _failure_with_billed_attempts())
assert [call.kwargs["response_cost"] for call in counters.call_args_list] == [0.5, 0.25, 0.5]
@ -3141,6 +3223,7 @@ async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_requ
@pytest.mark.asyncio
@pytest.mark.timeout(30)
async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_attributed_access_group(
live_spend_counters: SpendCounterStore, monkeypatch: pytest.MonkeyPatch
):
@ -3159,8 +3242,10 @@ async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_
),
)
await _fail_request(
ProxyDBLogger(),
from litellm.proxy.proxy_server import increment_spend_counters
await _fail_request_and_finish_counting_in_background(
increment_spend_counters,
{
**_failure_with_billed_attempts(),
"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool", "premium"], "model_info": {"id": "dep-b"}},