scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix: move the budget gate off the success callback into the detached task

Browse source 
greptile: _key_or_team_is_over_budget awaited get_current_spend inline in
async_log_success_event, and that read can fall back to an authoritative
DB query — the same request-path DB touch the job-snapshot rework just
removed. The gate now runs first inside _run_shadow_eval, where a
detached task absorbs the latency and the paid shadow/judge calls still
never fire for an over-budget key. The callback is back to zero awaits
beyond its own bookkeeping; a regression test asserts get_current_spend
is never awaited before the callback returns.

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Abhimanyu Kapur 2026-08-10 20:02:15 -07:00
parent f95f1c7751
commit 876bbecf3f
2 changed files with 49 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/integrations/shadow_eval_logger.py
View file

@ -330,8 +330,6 @@ class ShadowEvalLogger(CustomLogger):
return # only chat-shaped traffic is comparable
if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS:
return
if await _key_or_team_is_over_budget(metadata):
return # the shadowed key/team has no budget left for the extra calls
raw_messages: Final = kwargs.get("messages")
self._inflight_shadow_tasks += 1
task: Final = asyncio.create_task(
@ -347,6 +345,7 @@ class ShadowEvalLogger(CustomLogger):
dict(payload.get("model_parameters") or {}) # mutable-ok: frozen snapshot
),
parent_metadata=MappingProxyType(dict(request_metadata)), # mutable-ok: frozen snapshot
budget_metadata=MappingProxyType(dict(metadata)), # mutable-ok: frozen snapshot
)
)
task.add_done_callback(lambda _: setattr(self, "_inflight_shadow_tasks", self._inflight_shadow_tasks - 1))
@ -451,13 +450,21 @@ class ShadowEvalLogger(CustomLogger):
real_model: str,
model_parameters: Mapping[str, object],
parent_metadata: Mapping[str, object],
budget_metadata: Mapping[str, object] = _EMPTY_METADATA,
) -> None:
"""Detached background task: shadow call -> blind judge -> verdict row."""
"""Detached background task: budget gate -> shadow call -> blind judge -> verdict row.
The budget read lives here, not in the success hook: get_current_spend can fall
back to an authoritative DB read, which a detached task absorbs and the
production callback must not.
"""
prisma: Final = self._prisma_provider()
try:
real_text: Final = self._extract_response_text(response_obj)
if not real_text or not messages:
return
if await _key_or_team_is_over_budget(budget_metadata):
return
shadow: Final = await self._call_router_shadow(job.router_name, messages, model_parameters, parent_metadata)
if shadow is None:

70
tests/test_litellm/integrations/test_shadow_eval_logger.py
View file

@ -528,12 +528,13 @@ class TestKeyOrTeamIsOverBudget:
@pytest.mark.asyncio
class TestSuccessHookSkipsWhenOverBudget:
async def test_over_budget_key_is_skipped_before_scheduling_the_shadow_task(self):
job = ActiveShadowEvalJob(id="j1", router_name="r", shadow_percentage=100.0, judge_model="m", status="running")
logger, _, router = _logger_with_mocks(job)
logger._run_shadow_eval = AsyncMock()
kwargs = {
class TestBudgetGateRunsInTheDetachedTaskNotTheCallback:
"""get_current_spend can fall back to an authoritative DB read; the production
success callback must never await that, so the gate lives in the detached task."""
@staticmethod
def _kwargs(spend: float):
return {
"standard_logging_object": {
"id": "req-1",
"model": "gpt-4o",
@ -541,44 +542,51 @@ class TestSuccessHookSkipsWhenOverBudget:
"metadata": {
"user_api_key_hash": "key-hash",
"user_api_key_max_budget": 10.0,
"user_api_key_spend": 10.0,
"user_api_key_spend": spend,
},
},
"litellm_params": {"metadata": {}},
"messages": [{"role": "user", "content": "hi"}],
}
@staticmethod
def _job():
return ActiveShadowEvalJob(id="j1", router_name="r", shadow_percentage=100.0, judge_model="m", status="running")
async def test_over_budget_key_never_fires_a_shadow_call(self):
logger, _, router = _logger_with_mocks(self._job())
response = {"choices": [{"message": {"content": "real answer"}}]}
with patch("litellm.proxy.proxy_server.get_current_spend", AsyncMock(return_value=10.0)):
await logger.async_log_success_event(kwargs, MagicMock(), None, None)
await asyncio.sleep(0)
await logger.async_log_success_event(self._kwargs(spend=10.0), response, None, None)
await asyncio.sleep(0.05)
logger._run_shadow_eval.assert_not_awaited()
router.acompletion.assert_not_called()
async def test_under_budget_key_still_schedules_the_shadow_task(self):
job = ActiveShadowEvalJob(id="j1", router_name="r", shadow_percentage=100.0, judge_model="m", status="running")
logger, _, router = _logger_with_mocks(job)
logger._run_shadow_eval = AsyncMock()
kwargs = {
"standard_logging_object": {
"id": "req-1",
"model": "gpt-4o",
"call_type": "acompletion",
"metadata": {
"user_api_key_hash": "key-hash",
"user_api_key_max_budget": 10.0,
"user_api_key_spend": 4.0,
},
},
"litellm_params": {"metadata": {}},
"messages": [{"role": "user", "content": "hi"}],
}
async def test_under_budget_key_fires_the_shadow_call(self):
logger, _, router = _logger_with_mocks(self._job())
router.acompletion = AsyncMock(return_value={"choices": [{"message": {"content": "shadow"}}]})
response = {"choices": [{"message": {"content": "real answer"}}]}
with patch("litellm.proxy.proxy_server.get_current_spend", AsyncMock(return_value=4.0)):
await logger.async_log_success_event(kwargs, MagicMock(), None, None)
await asyncio.sleep(0)
await logger.async_log_success_event(self._kwargs(spend=4.0), response, None, None)
await asyncio.sleep(0.05)
logger._run_shadow_eval.assert_awaited_once()
router.acompletion.assert_awaited_once()
async def test_the_callback_itself_never_awaits_the_spend_read(self):
"""The spend read must happen after the callback returns, inside the task."""
logger, _, _ = _logger_with_mocks(self._job())
get_current_spend = AsyncMock(return_value=0.0)
response = {"choices": [{"message": {"content": "real answer"}}]}
with patch("litellm.proxy.proxy_server.get_current_spend", get_current_spend):
await logger.async_log_success_event(self._kwargs(spend=0.0), response, None, None)
spend_reads_when_callback_returned = get_current_spend.await_count
await asyncio.sleep(0.05)
assert spend_reads_when_callback_returned == 0
assert get_current_spend.await_count > 0
@pytest.mark.asyncio