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

fix(user_update): evict cached user on max_budget change so the personal key ceiling refreshes on every worker

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-20 00:50:18 +00:00
parent ee7d2b5094
commit 2e23c2d653
2 changed files with 46 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -101,6 +101,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_BUDGET_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget"})
def _user_table(
@ -1561,7 +1562,7 @@ async def _update_single_user_helper(
await _invalidate_user_spend_counter_if_changed(non_default_values)
if "model_max_budget" in non_default_values:
if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values):
await evict_and_broadcast(
cache_keys=(non_default_values["user_id"],),
user_api_key_cache=user_api_key_cache,
@ -1892,7 +1893,7 @@ async def bulk_user_update(
),
)
if "model_max_budget" in non_default_values:
if not _USER_BUDGET_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(
*(

43
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -2269,6 +2269,49 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
broadcast.assert_awaited_once_with(cache_key=saved_user.user_id)
@pytest.mark.asyncio
@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"])
async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> None:
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper, bulk_user_update
from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkUpdateUserRequest
saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", max_budget=500.0)
prisma_client: Final = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user])
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
)
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={"max_budget": 50.0}),
user_api_key_dict=admin,
litellm_changed_by=None,
)
prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={"max_budget": 50.0})
else:
await _update_single_user_helper(
user_request=UpdateUserRequest(user_id=saved_user.user_id, max_budget=50.0),
user_api_key_dict=admin,
)
assert prisma_client.update_data.call_args.kwargs["data"]["max_budget"] == 50.0
assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None
broadcast.assert_awaited_once_with(cache_key=saved_user.user_id)
def test_generate_request_base_validator():
"""
Test that GenerateRequestBase validator converts empty string to None for max_budget