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

fix(auth): load team membership once per request and skip prisma on an L1 hit

Browse source 
common_checks was querying get_team_membership twice, and DualCache awaited Redis SET on the auth path, so LRU eviction plus a hung Redis write showed up as two postgres spans
This commit is contained in:
Shivi Jain 2026-09-14 22:05:17 +05:30 • committed by yassin
parent 70ddc7e492
commit 53ba8b9866
5 changed files with 93 additions and 57 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -2197,30 +2197,51 @@ def _membership_from_cached_payload(
return cached_membership if cached_membership is not None else TEAM_MEMBERSHIP_CACHE_MISS
async def _set_team_membership_cache_entry(
async def _set_team_membership_l1(
user_api_key_cache: UserApiKeyCache,
key: str,
value: object,
*,
local_only: bool,
model_type: type[LiteLLM_TeamMembership] | None,
ttl: float | None,
) -> None:
match (model_type is not None, ttl is not None):
case (False, False):
await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only)
await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True)
case (False, True):
await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only, ttl=ttl)
await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, ttl=ttl)
case (True, False):
await user_api_key_cache.async_set_cache(
key=key, value=value, local_only=local_only, model_type=model_type
)
await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, model_type=model_type)
case (True, True):
await user_api_key_cache.async_set_cache(
key=key, value=value, local_only=local_only, model_type=model_type, ttl=ttl
key=key, value=value, local_only=True, model_type=model_type, ttl=ttl
)
async def _replicate_team_membership_to_redis(
user_api_key_cache: UserApiKeyCache,
key: str,
value: object,
*,
model_type: type[LiteLLM_TeamMembership] | None,
ttl: float | None,
write_epoch: int,
) -> None:
redis_cache: Final = user_api_key_cache.redis_cache
if redis_cache is None or _membership_write_epoch(key) != write_epoch:
return
payload: Final[object] = CacheCodec.serialize(value, model_type=model_type)
try:
if ttl is None:
await redis_cache.async_set_cache(key, payload)
else:
await redis_cache.async_set_cache(key, payload, ttl=ttl)
if _membership_write_epoch(key) != write_epoch:
await redis_cache.async_delete_cache(key)
except Exception:
return
async def _populate_team_membership_cache(
user_api_key_cache: UserApiKeyCache,
key: str,
@ -2230,36 +2251,27 @@ async def _populate_team_membership_cache(
ttl: float | None = None,
) -> None:
write_epoch: Final = _membership_write_epoch(key)
await _set_team_membership_cache_entry(
await _set_team_membership_l1(
user_api_key_cache,
key,
value,
local_only=True,
model_type=model_type,
ttl=ttl,
)
if _membership_write_epoch(key) != write_epoch:
await user_api_key_cache.async_delete_cache(key)
user_api_key_cache.in_memory_cache_for(key).delete_cache(key)
return
async def _replicate_to_redis() -> None:
try:
if _membership_write_epoch(key) != write_epoch:
return
await _set_team_membership_cache_entry(
user_api_key_cache,
key,
value,
local_only=False,
model_type=model_type,
ttl=ttl,
)
if _membership_write_epoch(key) != write_epoch:
await user_api_key_cache.async_delete_cache(key)
except Exception:
return
asyncio.create_task(_replicate_to_redis())
asyncio.create_task(
_replicate_team_membership_to_redis(
user_api_key_cache,
key,
value,
model_type=model_type,
ttl=ttl,
write_epoch=write_epoch,
)
)
@log_db_metrics
@ -2363,6 +2375,9 @@ async def get_team_membership(
if prisma_client is None:
raise Exception("No db connected")
prisma: Final[object] = prisma_client
if isinstance(prisma, str):
return None
task: Final = asyncio.ensure_future(
_load_team_membership_on_cache_miss(

5
litellm/proxy/auth/auth_object_prefetch.py
View file

@ -14,7 +14,6 @@ from pydantic import BaseModel, TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_cache import RedisCache
from litellm.constants import DEFAULT_IN_MEMORY_TTL
from litellm.models.organization import LiteLLM_OrganizationTable
from litellm.models.team import LiteLLM_TeamTableCachedObj
from litellm.models.team_membership import LiteLLM_TeamMembership
@ -191,13 +190,13 @@ def _iter_entries(refs: AuthObjectRefs, management_ttl: float) -> Iterator[_Cach
)
if refs.organization_id is not None:
yield _CacheEntry(
f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, DEFAULT_IN_MEMORY_TTL
f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, management_ttl
)
yield _CacheEntry(
f"org_id:{refs.organization_id}:with_budget",
"organization_row",
LiteLLM_OrganizationTable,
DEFAULT_IN_MEMORY_TTL,
management_ttl,
)
if refs.project_id is not None:
yield _CacheEntry(f"project_id:{refs.project_id}", "project_row", LiteLLM_ProjectTableCachedObj, management_ttl)

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

@ -5527,7 +5527,9 @@ async def _run_internal_user_budget_alert(
with (
patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam
patch("litellm.proxy.proxy_server.get_current_spend", _get_spend), # test-quality-ok: common_checks imports it locally
patch(
"litellm.proxy.proxy_server.get_current_spend", _get_spend
), # test-quality-ok: common_checks imports it locally
patch.object(slack_alerting, "send_alert", send_alert),
):
error: Final = await _check_for_error()
@ -6419,9 +6421,7 @@ async def test_get_team_membership_negative_caches_a_missing_row():
assert first is None
assert second is None
mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once()
cached = await cache.async_get_cache(
key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")
)
cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1"))
assert cached == NO_TEAM_MEMBERSHIP_SENTINEL
@ -6601,13 +6601,25 @@ async def test_get_team_membership_db_error_raises_503_not_none():
user_api_key_cache=cache,
)
cached = await cache.async_get_cache(
key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail")
)
cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail"))
assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE
assert cached is None
@pytest.mark.asyncio
async def test_get_team_membership_string_prisma_client_returns_none():
"""Unit tests stub prisma_client as a string; that is not a lookup failure and must not 503."""
from litellm.proxy.auth.auth_checks import get_team_membership
result = await get_team_membership(
user_id="u-str",
team_id="t-str",
prisma_client="hello-world",
user_api_key_cache=UserApiKeyCache(),
)
assert result is None
@pytest.mark.asyncio
async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_fails():
"""common_checks must not mark membership loaded-absent after a lookup error."""
@ -6699,7 +6711,7 @@ async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load():
@pytest.mark.asyncio
async def test_stale_membership_redis_replicate_does_not_restore_after_invalidate():
"""A delayed DualCache Redis SET must not resurrect membership after invalidation."""
"""A delayed Redis SET must not rewrite L1 or leave Redis holding membership after invalidation."""
from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state
from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key
@ -6739,6 +6751,7 @@ async def test_stale_membership_redis_replicate_does_not_restore_after_invalidat
assert loaded is not None
assert after_invalidate is None
assert after_replicate is None
redis_cache.async_delete_cache.assert_awaited()
@pytest.mark.asyncio
@ -6764,10 +6777,7 @@ async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sent
assert before is None
await invalidate_team_member_spend_state(user_id="u-1", team_id="t-1", user_api_key_cache=cache)
assert (
await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1"))
is None
)
assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) is None
after = await get_team_membership(
user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache

4
tests/test_litellm/proxy/auth/test_auth_object_prefetch.py
View file

@ -181,8 +181,8 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_
assert sets == sorted(
[
f"SET {TEAM_ID}_{USER_ID} ttl=5",
f"SET org_id:{ORG_ID} ttl=5",
f"SET org_id:{ORG_ID}:with_budget ttl=5",
f"SET org_id:{ORG_ID} ttl=60",
f"SET org_id:{ORG_ID}:with_budget ttl=60",
f"SET {USER_ID} ttl=60",
f"SET team_id:{TEAM_ID} ttl=60",
f"SET team_membership:{USER_ID}:{TEAM_ID} ttl=None",

34
tests/test_litellm/proxy/auth/test_model_access_group_budgets.py
View file

@ -31,6 +31,7 @@ from litellm.proxy.auth.auth_checks import (
)
from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key
from litellm.proxy.common_utils.user_api_key_cache import (
NO_TEAM_MEMBERSHIP_SENTINEL,
UserApiKeyCache,
model_access_group_registry_cache_key,
model_access_group_spend_counter_key,
@ -95,6 +96,11 @@ async def _cache(
),
model_type=LiteLLM_TeamMembership,
)
else:
await cache.async_set_cache(
key=team_membership_reservation_cache_key(user_id=USER_ID, team_id=TEAM_ID),
value=NO_TEAM_MEMBERSHIP_SENTINEL,
)
if org_models:
await cache.async_set_cache(
key=f"org_id:{ORG_ID}",
@ -314,9 +320,7 @@ class _RecordingPrismaClient:
def __init__(self, *rows: _MagBudgetRow) -> None:
self.rows = {row.access_group_name: row for row in rows}
self.batches: list[list[str]] = []
self.db = SimpleNamespace(
litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many)
)
self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many))
async def _find_many(self, **kwargs):
requested = list(kwargs["where"]["access_group_name"]["in"])
@ -345,7 +349,9 @@ async def _enforce(
read, seen = _spend_reader(spend_by_counter_key or {})
# The check takes its client and cache as arguments, injected just below. get_current_spend is the
# one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter.
with patch("litellm.proxy.proxy_server.get_current_spend", read): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point
with patch(
"litellm.proxy.proxy_server.get_current_spend", read
): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point
await _model_access_group_max_budget_check(
matched_model_access_groups=matched,
prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows),
@ -491,9 +497,7 @@ async def test_a_second_request_serves_the_budget_row_from_cache():
async def test_a_database_error_does_not_block_the_request():
class _FailingPrismaClient:
def __init__(self) -> None:
self.db = SimpleNamespace(
litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom)
)
self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom))
async def _boom(self, **kwargs):
raise RuntimeError("database unavailable")
@ -509,9 +513,15 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) ->
with (
# common_checks resolves all three off the proxy_server module at call time; its signature
# has no client, cache or spend-reader parameter to pass them through instead.
patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter
patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter
patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point
patch(
"litellm.proxy.proxy_server.prisma_client", prisma_client
), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter
patch(
"litellm.proxy.proxy_server.user_api_key_cache", cache
), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter
patch(
"litellm.proxy.proxy_server.get_current_spend", read
), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point
):
return await common_checks(
request_body={"model": "gpt-4o", "messages": []},
@ -524,7 +534,9 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) ->
llm_router=Router(model_list=MODEL_LIST),
proxy_logging_obj=ProxyLogging(user_api_key_cache=cache),
valid_token=UserAPIKeyAuth(api_key="hashed", models=["tier-a"], user_id=USER_ID),
request=SimpleNamespace(method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")),
request=SimpleNamespace(
method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")
),
skip_budget_checks=skip_budget_checks,
)