mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): enforce a user's max_parallel_requests and allow clearing it with null
/user/new and /user/update store max_parallel_requests on a user, but nothing ever read it back out. UserAPIKeyAuth carried the user's tpm/rpm limits but had no slot for max_parallel_requests, so none of the three auth paths that build it (key auth, JWT auth, keyless MCP auth) could pass it to the rate limiter. The v3 limiter's user-scope descriptor builder never looked for it either, so the concurrency gauge was never created for that scope even though the same gauge mechanism already works for API keys. The legacy limiter (used behind LEGACY_MULTI_INSTANCE_RATE_LIMITING) had the identical gap, hardcoding sys.maxsize with a TODO instead of reading the field. Separately, /user/update with max_parallel_requests: null silently kept the old value, unlike tpm_limit/rpm_limit/max_budget/budget_duration, which null already clears. Adds the missing field, wires it through every path that already carries user_tpm_limit/user_rpm_limit, lets both limiters enforce it the same way they already enforce the API key's own limit, and adds max_parallel_requests to the set of fields a null can clear.
This commit is contained in:
parent
c98641fd5a
commit
01c1426abf
11 changed files with 111 additions and 6 deletions
|
|
@ -3056,6 +3056,7 @@ class JWTAuthManager:
|
|||
agent_id=result.get("agent_id"),
|
||||
user_tpm_limit=user.tpm_limit if user is not None and not admin else None,
|
||||
user_rpm_limit=user.rpm_limit if user is not None and not admin else None,
|
||||
user_max_parallel_requests=user.max_parallel_requests if user is not None and not admin else None,
|
||||
user_model_max_budget=user.model_max_budget if user is not None and not admin else None,
|
||||
**team_grants(
|
||||
team_object=result["team_object"],
|
||||
|
|
|
|||
|
|
@ -4464,6 +4464,7 @@ async def _return_user_api_key_auth_obj(
|
|||
user_api_key_kwargs.update(
|
||||
user_tpm_limit=user_obj.tpm_limit,
|
||||
user_rpm_limit=user_obj.rpm_limit,
|
||||
user_max_parallel_requests=user_obj.max_parallel_requests,
|
||||
user_email=user_obj.user_email,
|
||||
user_spend=getattr(user_obj, "spend", None),
|
||||
user_max_budget=getattr(user_obj, "max_budget", None),
|
||||
|
|
|
|||
|
|
@ -422,10 +422,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
|
|||
if user_id is not None:
|
||||
user_tpm_limit = user_api_key_dict.user_tpm_limit
|
||||
user_rpm_limit = user_api_key_dict.user_rpm_limit
|
||||
user_max_parallel_requests = user_api_key_dict.user_max_parallel_requests
|
||||
if user_tpm_limit is None:
|
||||
user_tpm_limit = sys.maxsize
|
||||
if user_rpm_limit is None:
|
||||
user_rpm_limit = sys.maxsize
|
||||
if user_max_parallel_requests is None:
|
||||
user_max_parallel_requests = sys.maxsize
|
||||
|
||||
request_count_api_key = f"{user_id}::{precise_minute}::request_count"
|
||||
# print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}")
|
||||
|
|
@ -434,7 +437,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
|
|||
cache=cache,
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a user
|
||||
max_parallel_requests=user_max_parallel_requests,
|
||||
current=cache_objects["request_count_user_id"],
|
||||
request_count_api_key=request_count_api_key,
|
||||
tpm_limit=user_tpm_limit,
|
||||
|
|
|
|||
|
|
@ -3435,7 +3435,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
|
||||
# User rate limits
|
||||
if user_api_key_dict.user_id and (
|
||||
user_api_key_dict.user_rpm_limit is not None or user_api_key_dict.user_tpm_limit is not None
|
||||
user_api_key_dict.user_rpm_limit is not None
|
||||
or user_api_key_dict.user_tpm_limit is not None
|
||||
or user_api_key_dict.user_max_parallel_requests is not None
|
||||
):
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
|
|
@ -3444,6 +3446,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
rate_limit={
|
||||
"requests_per_unit": user_api_key_dict.user_rpm_limit,
|
||||
"tokens_per_unit": user_api_key_dict.user_tpm_limit,
|
||||
"max_parallel_requests": user_api_key_dict.user_max_parallel_requests,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1286,7 +1286,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
|
|||
fields_set: Final = data.fields_set() if hasattr(data, "fields_set") else set()
|
||||
|
||||
for k, v in data_json.items():
|
||||
if k in ("max_budget", "budget_duration", "tpm_limit", "rpm_limit"):
|
||||
if k in ("max_budget", "budget_duration", "tpm_limit", "rpm_limit", "max_parallel_requests"):
|
||||
if k in fields_set:
|
||||
non_default_values[k] = v
|
||||
elif k == "model_max_budget":
|
||||
|
|
|
|||
|
|
@ -1308,6 +1308,7 @@ class MCPRequestHandler:
|
|||
# per-KEY limits do not apply, there being no key.)
|
||||
user_tpm_limit=user_object.tpm_limit,
|
||||
user_rpm_limit=user_object.rpm_limit,
|
||||
user_max_parallel_requests=user_object.max_parallel_requests,
|
||||
)
|
||||
# Server-only marker, set AFTER construction: the before-validator strips it from any validated
|
||||
# input, so caller-supplied data (key metadata, JWT claims) can never forge it.
|
||||
|
|
|
|||
|
|
@ -139,6 +139,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
tpm_limit_per_model: dict[str, int] | None = None
|
||||
user_tpm_limit: int | None = None
|
||||
user_rpm_limit: int | None = None
|
||||
user_max_parallel_requests: int | None = None
|
||||
user_email: str | None = None
|
||||
user_spend: float | None = None
|
||||
user_max_budget: float | None = None
|
||||
|
|
|
|||
|
|
@ -1924,6 +1924,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget():
|
|||
{
|
||||
"tpm_limit": 1000,
|
||||
"rpm_limit": 100,
|
||||
"max_parallel_requests": 5,
|
||||
"user_email": "test@example.com",
|
||||
"spend": 250.0,
|
||||
"max_budget": 1000.0,
|
||||
|
|
@ -1961,6 +1962,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget():
|
|||
assert result.user_max_budget == 1000.0
|
||||
assert result.user_tpm_limit == 1000
|
||||
assert result.user_rpm_limit == 100
|
||||
assert result.user_max_parallel_requests == 5
|
||||
assert result.user_email == "test@example.com"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -97,6 +97,35 @@ async def test_pre_call_hook_counts_a_cli_session_under_the_per_user_alias_not_t
|
|||
assert counted == {"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, counted
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_max_parallel_requests_is_enforced():
|
||||
"""
|
||||
A user's own `max_parallel_requests` must be enforced on the legacy
|
||||
(LEGACY_MULTI_INSTANCE_RATE_LIMITING) limiter too, the same as it already
|
||||
is for an API key: a second request from the same user, under a
|
||||
different API key, in the same minute bucket, is refused.
|
||||
"""
|
||||
handler: Final = PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock
|
||||
)
|
||||
user_id: Final = "pia"
|
||||
first_key: Final = UserAPIKeyAuth(
|
||||
api_key="sk-pia-key-one", user_id=user_id, user_max_parallel_requests=1
|
||||
)
|
||||
second_key: Final = UserAPIKeyAuth(
|
||||
api_key="sk-pia-key-two", user_id=user_id, user_max_parallel_requests=1
|
||||
)
|
||||
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=first_key, cache=DualCache(), data={"model": "gpt-4o-mini"}, call_type="completion"
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=second_key, cache=DualCache(), data={"model": "gpt-4o-mini"}, call_type="completion"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"response_obj",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -1029,6 +1029,47 @@ async def test_rejected_request_does_not_consume_parallel_slot_v3():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_max_parallel_requests_is_enforced_v3():
|
||||
"""
|
||||
A user's own `max_parallel_requests` (set via /user/new or /user/update,
|
||||
independent of any key-level limit) must be enforced: a second concurrent
|
||||
request from the same user, admitted under a different API key, is
|
||||
refused with 429 once the user's slot is full.
|
||||
"""
|
||||
local_cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache)
|
||||
)
|
||||
user_id = "pia"
|
||||
first_key = hash_token("sk-pia-key-one")
|
||||
second_key = hash_token("sk-pia-key-two")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key=first_key, user_id=user_id, user_max_parallel_requests=1
|
||||
)
|
||||
other_key_same_user = UserAPIKeyAuth(
|
||||
api_key=second_key, user_id=user_id, user_max_parallel_requests=1
|
||||
)
|
||||
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data={"model": "gpt-3.5-turbo"},
|
||||
call_type="",
|
||||
)
|
||||
acquisition = get_request_stash().parallel_slot
|
||||
assert acquisition is not None
|
||||
assert acquisition["counter_keys"] == [f"{{user:{user_id}}}:max_parallel_requests"]
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=other_key_same_user,
|
||||
cache=local_cache,
|
||||
data={"model": "gpt-3.5-turbo"},
|
||||
call_type="",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parallel_gauge_uses_atomic_redis_script_v3():
|
||||
"""
|
||||
|
|
@ -1242,6 +1283,25 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3():
|
|||
should_rate_limit_called
|
||||
), "should_rate_limit should be called when max parallel requests are configured"
|
||||
|
||||
# Test 7: User-level max parallel requests configured - should_rate_limit SHOULD be called
|
||||
should_rate_limit_called = False
|
||||
user_api_key_dict_with_user_parallel_limits = UserAPIKeyAuth(
|
||||
api_key=_api_key,
|
||||
user_id="test_user",
|
||||
user_max_parallel_requests=5, # User-level max parallel requests configured
|
||||
)
|
||||
|
||||
await parallel_request_handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict_with_user_parallel_limits,
|
||||
cache=local_cache,
|
||||
data={"model": "gpt-3.5-turbo"},
|
||||
call_type="",
|
||||
)
|
||||
|
||||
assert (
|
||||
should_rate_limit_called
|
||||
), "should_rate_limit should be called when user-level max parallel requests are configured"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_specific_rate_limits_only_called_when_configured_v3():
|
||||
|
|
|
|||
|
|
@ -2255,16 +2255,20 @@ def test_update_internal_user_params_ignores_other_nones():
|
|||
assert non_default_values["max_budget"] == 100.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["tpm_limit", "rpm_limit"], ids=["tpm_limit", "rpm_limit"])
|
||||
@pytest.mark.parametrize(
|
||||
"field",
|
||||
["tpm_limit", "rpm_limit", "max_parallel_requests"],
|
||||
ids=["tpm_limit", "rpm_limit", "max_parallel_requests"],
|
||||
)
|
||||
def test_update_internal_user_params_explicit_null_clears_rate_limit_but_omitted_is_untouched(
|
||||
field: str,
|
||||
) -> None:
|
||||
data: Final = UpdateUserRequest(user_id="limit-clear", **{field: None})
|
||||
result: Final = _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data)
|
||||
other_field: Final = "rpm_limit" if field == "tpm_limit" else "tpm_limit"
|
||||
other_fields: Final = {"tpm_limit", "rpm_limit", "max_parallel_requests"} - {field}
|
||||
|
||||
assert result[field] is None
|
||||
assert other_field not in result
|
||||
assert other_fields.isdisjoint(result)
|
||||
|
||||
|
||||
def test_update_internal_user_params_keeps_original_max_budget_when_not_provided():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue