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

perf(auth): gather independent pre-call budget-enforcement reads (#31604)

Browse source 
increment_spend_counters was parallelized in #31578, but the dominant
per-request cost under high concurrency is the pre-call budget enforcement
in common_checks, which still ran a Redis-first get_current_spend per scope
(team, team windows, key windows, org, tag, user, team member, end user)
one sequential await after another inside the auth span.

The per-scope reads target distinct counter keys with no cross-scope
ordering dependency, so they now run concurrently under asyncio.gather.
Key metadata.tags injection still runs before the gather so the tag budget
check sees it, and every scope settles before the first error in
scope-priority order propagates, preserving the previous rejection semantics.

Resolves LIT-4090
This commit is contained in:
Yassin Kortam 2026-06-30 00:44:35 +03:00 • committed by GitHub
parent 7a1ba958f8
commit 829bfebe0f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 259 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

147
litellm/proxy/auth/auth_checks.py
View file

@ -603,41 +603,8 @@ async def common_checks(
# If this is a free model, skip all budget checks
if not skip_budget_checks:
# 3. If team is in budget
with tracer.trace("litellm.proxy.auth.common_checks.team_max_budget_check"):
await _team_max_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
# 3.1. Multi-window budget check for team
with tracer.trace("litellm.proxy.auth.common_checks.team_multi_budget_check"):
await _team_multi_budget_check(team_object=team_object)
# 3.2. Multi-window budget check for key
with tracer.trace("litellm.proxy.auth.common_checks.virtual_key_multi_budget_check"):
if valid_token is not None:
await _virtual_key_multi_budget_check(valid_token=valid_token)
# 3.0.5. If team is over soft budget (alert only, doesn't block)
with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"):
await _team_soft_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
# 3.1. If organization is in budget
with tracer.trace("litellm.proxy.auth.common_checks.organization_max_budget_check"):
await _organization_max_budget_check(
valid_token=valid_token,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# Key metadata.tags are injected into request_body here so the tag budget
# check can read them; this mutation must run before the gathered checks.
if valid_token is not None:
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
@ -651,51 +618,83 @@ async def common_checks(
user_api_key_dict=valid_token,
)
with tracer.trace("litellm.proxy.auth.common_checks.tag_max_budget_check"):
await _tag_max_budget_check(
request_body=request_body,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
async def _user_max_budget_check() -> None:
# 4.1 personal budget, if personal key
if (
(team_object is None or team_object.team_id is None)
and user_object is not None
and user_object.max_budget is not None
):
from litellm.proxy.proxy_server import get_current_spend
# 4. If user is in budget
## 4.1 check personal budget, if personal key
if (
(team_object is None or team_object.team_id is None)
and user_object is not None
and user_object.max_budget is not None
):
user_budget = user_object.max_budget
from litellm.proxy.proxy_server import get_current_spend
user_spend = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
)
if math.isfinite(user_budget) and user_spend >= user_budget:
raise litellm.BudgetExceededError(
current_cost=user_spend,
user_budget = user_object.max_budget
user_spend = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
)
if math.isfinite(user_budget) and user_spend >= user_budget:
raise litellm.BudgetExceededError(
current_cost=user_spend,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
)
## 4.2 check team member budget, if team key
with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_budget"):
await _check_team_member_budget(
team_object=team_object,
user_object=user_object,
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
# Each scope reads a distinct counter key with no cross-scope ordering
# dependency, so the per-scope Redis-first reads run concurrently instead
# of one sequential await per scope. return_exceptions lets every scope
# settle, then the first error in scope-priority order propagates exactly
# as the sequential path raised.
budget_check_coros = tuple(
coro
for coro in (
_team_max_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
),
_team_multi_budget_check(team_object=team_object),
_virtual_key_multi_budget_check(valid_token=valid_token) if valid_token is not None else None,
_team_soft_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
),
_organization_max_budget_check(
valid_token=valid_token,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
),
_tag_max_budget_check(
request_body=request_body,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
),
_user_max_budget_check(),
_check_team_member_budget(
team_object=team_object,
user_object=user_object,
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
),
_check_end_user_budget(end_user_obj=end_user_object, route=route)
if end_user_object is not None and end_user_object.litellm_budget_table is not None
else None,
)
if coro is not None
)
# 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
if end_user_object is not None and end_user_object.litellm_budget_table is not None:
await _check_end_user_budget(end_user_obj=end_user_object, route=route)
with tracer.trace("litellm.proxy.auth.common_checks.budget_checks"):
budget_results = await asyncio.gather(*budget_check_coros, return_exceptions=True)
budget_error = next((r for r in budget_results if isinstance(r, BaseException)), None)
if budget_error is not None:
raise budget_error
_enforce_user_param_check(general_settings, request, request_body, route)
_global_proxy_budget_check(global_proxy_spend, skip_budget_checks, route)

186
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -3990,3 +3990,189 @@ class TestManagementObjectTTLHonored:
)
assert mem.last_ttl == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
class _BudgetSpendConcurrencyProbe:
"""Stand-in for get_current_spend that pins how many scope checks are in flight.
Each call registers itself, records the peak simultaneous count, and blocks on
``release`` until the test lets it proceed. ``all_arrived`` only fires once
``expected`` distinct scope reads are suspended here at the same time, which can
happen only if common_checks gathers the per-scope reads instead of awaiting
them one after another.
"""
def __init__(self, expected: int):
self.expected = expected
self.in_flight = 0
self.max_in_flight = 0
self.all_arrived = asyncio.Event()
self.release = asyncio.Event()
async def __call__(self, *args, **kwargs) -> float:
self.in_flight += 1
self.max_in_flight = max(self.max_in_flight, self.in_flight)
if self.in_flight >= self.expected:
self.all_arrived.set()
try:
await self.release.wait()
finally:
self.in_flight -= 1
return 0.0
@pytest.mark.asyncio
async def test_common_checks_budget_reads_run_concurrently():
"""Independent per-scope budget reads in common_checks must run concurrently.
team max, team window, key window, and end-user each read a distinct spend
counter with no cross-scope dependency. With the gather they are all suspended
in get_current_spend simultaneously; reverting to sequential awaits leaves only
one in flight at a time, so ``all_arrived`` never fires and this test times out.
"""
from fastapi import Request
from litellm.proxy.auth.auth_checks import common_checks
team = LiteLLM_TeamTable(
team_id="t1",
spend=0.0,
max_budget=100.0,
budget_limits=[{"budget_duration": "1d", "max_budget": 100.0}],
)
token = UserAPIKeyAuth(
token="k1",
budget_limits=[{"budget_duration": "1d", "max_budget": 100.0}],
)
end_user = LiteLLM_EndUserTable(
user_id="eu1",
blocked=False,
spend=0.0,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0),
)
probe = _BudgetSpendConcurrencyProbe(expected=4)
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
"litellm.proxy.proxy_server.get_current_spend", probe
):
task = asyncio.create_task(
common_checks(
request_body={"messages": [{"role": "user", "content": "hi"}]},
team_object=team,
user_object=None,
end_user_object=end_user,
global_proxy_spend=None,
general_settings={},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=MagicMock(spec=Request),
)
)
try:
await asyncio.wait_for(probe.all_arrived.wait(), timeout=3.0)
assert probe.max_in_flight == 4
finally:
probe.release.set()
assert await task is True
@pytest.mark.asyncio
async def test_common_checks_budget_gather_raises_highest_priority_scope():
"""A gathered scope that is over budget must still raise BudgetExceededError.
When more than one scope is over budget the error from the highest-priority
scope (team, matching the previous sequential order) propagates; when only a
lower-priority scope (end-user) is over budget its error still surfaces. This
fails if any scope is dropped from the gather or if errors are swallowed.
"""
from fastapi import Request
from litellm.proxy.auth.auth_checks import common_checks
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
if counter_key == "spend:team:t1":
return _spend_by_counter.team
if counter_key == "spend:end_user:eu1":
return _spend_by_counter.end_user
return 0.0
team = LiteLLM_TeamTable(team_id="t1", spend=0.0, max_budget=100.0)
end_user = LiteLLM_EndUserTable(
user_id="eu1",
blocked=False,
spend=0.0,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0),
)
async def _run():
return await common_checks(
request_body={"messages": [{"role": "user", "content": "hi"}]},
team_object=team,
user_object=None,
end_user_object=end_user,
global_proxy_spend=None,
general_settings={},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=None,
request=MagicMock(spec=Request),
)
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
"litellm.proxy.proxy_server.get_current_spend", _spend_by_counter
):
# Both team and end-user over budget: team wins on priority.
_spend_by_counter.team = 999.0
_spend_by_counter.end_user = 999.0
with pytest.raises(litellm.BudgetExceededError) as both_over:
await _run()
assert "Team=t1" in str(both_over.value)
# Only the lower-priority end-user scope over budget: its error still raises.
_spend_by_counter.team = 0.0
_spend_by_counter.end_user = 999.0
with pytest.raises(litellm.BudgetExceededError) as end_user_over:
await _run()
assert "End User=eu1" in str(end_user_over.value)
@pytest.mark.asyncio
async def test_common_checks_personal_user_budget_blocks_in_gather():
"""The personal-key user budget scope is enforced inside the gather.
For a personal key (no team) whose user is over budget, the gathered user
check must raise BudgetExceededError. This guards the relocated personal
user-budget read and fails if that scope is dropped from the gather.
"""
from fastapi import Request
from litellm.proxy.auth.auth_checks import common_checks
user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0)
token = UserAPIKeyAuth(token="k1", user_id="u1")
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
return 999.0 if counter_key == "spend:user:u1" else 0.0
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
"litellm.proxy.proxy_server.get_current_spend", _spend_by_counter
):
with pytest.raises(litellm.BudgetExceededError) as over:
await common_checks(
request_body={"messages": [{"role": "user", "content": "hi"}]},
team_object=None,
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=MagicMock(spec=Request),
)
assert "User=u1" in str(over.value)