scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-15 04:19:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 121

fix(proxy): evict the cached user on every proxy for every /user/update and /user/bulk_update (#46500)

Browse source 
* fix(proxy): evict the cached user on every proxy for every /user/update and /user/bulk_update

Previously the cached user row was only evicted and broadcast when the
update touched max_budget, model_max_budget, tpm_limit, rpm_limit or
metadata, so a models or user_role change kept being served by peer
proxies until the cache TTL expired

* test(proxy): cover user cache eviction past the first bulk invalidation batch
This commit is contained in:
yuneng-jiang 2026-10-14 13:25:20 -07:00 • committed by GitHub
parent 493ed616bb
commit 47c40586e1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 52 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -121,8 +121,7 @@ if TYPE_CHECKING:
router: Final = APIRouter()
_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig])
_USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
_USER_LIMIT_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget", "tpm_limit", "rpm_limit"})
_USER_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
def _user_table(
@ -1634,11 +1633,10 @@ async def _update_single_user_helper(
await _invalidate_user_spend_counter_if_changed(non_default_values)
if not _USER_LIMIT_CACHE_FIELDS.isdisjoint(non_default_values) or "metadata" in data_json:
await evict_and_broadcast(
cache_keys=(non_default_values["user_id"],),
user_api_key_cache=user_api_key_cache,
)
await evict_and_broadcast(
cache_keys=(non_default_values["user_id"],),
user_api_key_cache=user_api_key_cache,
)
if "object_permission_id" in non_default_values:
await _invalidate_cached_user_entitlement(
@ -1992,14 +1990,13 @@ async def bulk_user_update(
),
)
if not _USER_LIMIT_CACHE_FIELDS.isdisjoint(non_default_values):
for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE):
await asyncio.gather(
*(
evict_and_broadcast(cache_keys=(user.user_id,), user_api_key_cache=user_api_key_cache)
for user in all_users_in_db[start : start + _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE]
)
for start in range(0, len(all_users_in_db), _USER_CACHE_INVALIDATION_BATCH_SIZE):
await asyncio.gather(
*(
evict_and_broadcast(cache_keys=(user.user_id,), user_api_key_cache=user_api_key_cache)
for user in all_users_in_db[start : start + _USER_CACHE_INVALIDATION_BATCH_SIZE]
)
)
# Create individual success results
for user in all_users_in_db:

1
tests/integration/management/test_user_updates_wedged_coordination_redis.py
View file

@ -634,7 +634,6 @@ def test_user_models_only_update_is_enforced_on_the_gateway(gateway: Gateway) ->
def test_user_models_update_without_metadata_propagates_to_gateway_and_peer(gateway: Gateway, peer: Gateway) -> None:
pytest.skip("BUG: models-only /user/update leaves warmed peer serving the revoked model with 200")
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,

59
tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -2496,13 +2496,27 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker:
@pytest.mark.asyncio
@pytest.mark.parametrize(
("field", "new_limit"),
[("tpm_limit", 100), ("rpm_limit", 1), ("tpm_limit", None), ("rpm_limit", None)],
ids=["tpm_limit", "rpm_limit", "tpm_limit-cleared", "rpm_limit-cleared"],
("field", "new_value"),
[
("tpm_limit", 100),
("rpm_limit", 1),
("tpm_limit", None),
("rpm_limit", None),
("models", ["model-fir"]),
("user_role", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value),
],
ids=[
"tpm_limit",
"rpm_limit",
"tpm_limit-cleared",
"rpm_limit-cleared",
"models",
"user_role",
],
)
@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"])
async def test_user_rate_limit_update_reaches_cached_user_on_every_worker(
mocker: MockerFixture, field: str, new_limit: int | None, all_users: bool
async def test_user_update_reaches_cached_user_on_every_worker(
mocker: MockerFixture, field: str, new_value: object, all_users: bool
) -> None:
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.auth.auth_checks import get_user_object
@ -2560,16 +2574,25 @@ async def test_user_rate_limit_update_reaches_cached_user_on_every_worker(
self.update_data_payload = data
return {"user_id": user_id, "data": self.updated_user}
saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", tpm_limit=100000, rpm_limit=1000)
updated_user: Final = saved_user.model_copy(update={field: new_limit})
old_limit: Final = 100000 if field == "tpm_limit" else 1000
saved_user: Final = LiteLLM_UserTable(
user_id="user-spruce",
tpm_limit=100000,
rpm_limit=1000,
models=["model-spruce"],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
updated_user: Final = saved_user.model_copy(update={field: new_value})
old_value: Final = saved_user.model_dump()[field]
assert old_value != new_value
prisma_client: Final = _PrismaClientMock(
db=_DatabaseMocks(
litellm_usertable=_UserTableMocks(
find_first=mocker.AsyncMock(return_value=saved_user),
find_unique=mocker.AsyncMock(return_value=updated_user),
find_many=mocker.AsyncMock(return_value=[saved_user]),
find_many=mocker.AsyncMock(
return_value=[*(LiteLLM_UserTable(user_id=f"user-pine-{index}") for index in range(50)), saved_user]
),
update_many=mocker.AsyncMock(return_value=1),
)
),
@ -2614,28 +2637,28 @@ async def test_user_rate_limit_update_reaches_cached_user_on_every_worker(
user_id_upsert=False,
)
assert handling_user_before is not None
assert handling_user_before.model_dump()[field] == old_limit
assert handling_user_before.model_dump()[field] == old_value
assert other_user_before is not None
assert other_user_before.model_dump()[field] == old_limit
assert other_user_before.model_dump()[field] == old_value
admin: Final = UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN)
if all_users:
await bulk_user_update(
data=BulkUpdateUserRequest(
all_users=True,
user_updates=UpdateUserRequestNoUserIDorEmail.model_validate({field: new_limit}),
user_updates=UpdateUserRequestNoUserIDorEmail.model_validate({field: new_value}),
),
user_api_key_dict=admin,
litellm_changed_by=None,
)
prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={field: new_limit})
prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={field: new_value})
else:
await user_update(
data=UpdateUserRequest.model_validate({"user_id": saved_user.user_id, field: new_limit}),
data=UpdateUserRequest.model_validate({"user_id": saved_user.user_id, field: new_value}),
user_api_key_dict=admin,
)
assert prisma_client.update_data_payload is not None
assert prisma_client.update_data_payload[field] == new_limit
assert prisma_client.update_data_payload[field] == new_value
remote_subscriber: Final = AuthCacheInvalidationSubscriber(
redis_cache=cast("RedisCache", _FakeRedisCache()),
@ -2659,10 +2682,10 @@ async def test_user_rate_limit_update_reaches_cached_user_on_every_worker(
user_id_upsert=False,
)
assert handling_user_after is not None
assert handling_user_after.model_dump()[field] == new_limit
assert handling_user_after.model_dump()[field] == new_value
assert other_user_after is not None
assert other_user_after.model_dump()[field] == new_limit, (
"another worker still enforces the old limit; the update was never broadcast"
assert other_user_after.model_dump()[field] == new_value, (
"another worker still serves the old user row; the update was never broadcast"
)