From 47c40586e196c035b1dc3da87f62f3fb8459ad07 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 14 Oct 2026 13:25:20 -0700 Subject: [PATCH] fix(proxy): evict the cached user on every proxy for every /user/update and /user/bulk_update (#46500) * 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 --- .../internal_user_endpoints.py | 25 ++++---- ..._user_updates_wedged_coordination_redis.py | 1 - .../test_internal_user_endpoints.py | 59 +++++++++++++------ 3 files changed, 52 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 2374f7d7aae..4e1001a4430 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -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: diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index 09b16e7d5ef..5ec5f85eef2 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -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, diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 39265b37b95..84c4608d29b 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -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" )