diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ac85e432d8f..10c11119006 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -22,7 +22,7 @@ from typing import Any, Final, Literal, Protocol, cast, overload import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -88,6 +88,7 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( SCIM_ENTITLEMENTS_METADATA_KEY, SCIM_ROLES_METADATA_KEY, ) +from litellm.types.utils import BudgetConfig if TYPE_CHECKING: from prisma import models as prisma_models @@ -98,7 +99,8 @@ if TYPE_CHECKING: from litellm.proxy.utils import ProxyLogging router: Final = APIRouter() -_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(GenericBudgetConfigType) +_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) +_USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 def _user_table( @@ -1257,7 +1259,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda non_default_values[k] = v elif k == "model_max_budget": if k in fields_set: - _USER_MODEL_BUDGET_ADAPTER.validate_python({} if v is None else v) + try: + _USER_MODEL_BUDGET_ADAPTER.validate_python({} if v is None else v) + except ValidationError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc non_default_values[k] = {} if v is None else v elif ( v is not None @@ -1888,10 +1893,13 @@ async def bulk_user_update( ) if "model_max_budget" in non_default_values: - await evict_and_broadcast( - cache_keys=tuple(user.user_id for user in all_users_in_db), - user_api_key_cache=user_api_key_cache, - ) + 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] + ) + ) # Create individual success results for user in all_users_in_db: diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 2f665f4acd3..92b1ab1586d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -5,7 +5,7 @@ from typing import Final import pytest from fastapi.testclient import TestClient -from pydantic import ValidationError +from fastapi import HTTPException from pytest_mock import MockerFixture @@ -2147,7 +2147,7 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo "user_alias": "Spruce", } - replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}} + replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0} request: Final = UpdateUserRequest( user_id="user-spruce", model_max_budget=replacement, @@ -2168,8 +2168,9 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo def test_update_internal_user_params_rejects_invalid_model_budget(invalid_budget: dict[str, object]) -> None: request: Final = UpdateUserRequest(user_id="user-spruce", model_max_budget=invalid_budget) - with pytest.raises(ValidationError): + with pytest.raises(HTTPException) as exc: _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request) + assert exc.value.status_code == 400 @pytest.mark.asyncio @@ -2228,6 +2229,16 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke new_callable=mocker.AsyncMock, ) + with pytest.raises(HTTPException) as exc: + await bulk_user_update( + data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": {"model-spruce": "invalid"}}), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_changed_by=None, + ) + assert exc.value.status_code == 400 + prisma_client.db.litellm_usertable.update_many.assert_not_called() + assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) == saved_user + response: Final = await bulk_user_update( data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": None}), user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN),