mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
fix(proxy): key model rpm/tpm override takes precedence over team model limit
A key inside a team with model_rpm_limit / model_tpm_limit in team metadata could not override those limits for itself: the v3 limiter always added the team's per-model descriptor next to the key's, so the tighter team limit won. The docs already say the resolution order is key metadata > key model_max_budget > team metadata get_key_own_model_rate_limit returns only what the key sets on itself, and the team descriptor now carries only the metrics the key does not override, so an rpm-only override still leaves the team tpm pool enforced Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
79d4d4d8f5
commit
51a4cb9fdd
4 changed files with 165 additions and 57 deletions
|
|
@ -976,6 +976,32 @@ def _get_deployment_default_tpm_limit(model_name: str) -> int | None:
|
|||
return _get_deployment_default_limit(model_name, "default_api_key_tpm_limit")
|
||||
|
||||
|
||||
def get_key_own_model_rate_limit(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
|
||||
) -> dict[str, int] | None:
|
||||
"""
|
||||
Per-model limit the key sets on itself: key metadata first, then model_max_budget.
|
||||
|
||||
Unlike get_key_model_rpm_limit / get_key_model_tpm_limit this never falls back to the
|
||||
team, so callers can tell a key override apart from an inherited team limit.
|
||||
"""
|
||||
if user_api_key_dict.metadata:
|
||||
result: Final = user_api_key_dict.metadata.get(rate_limit_key)
|
||||
if result:
|
||||
return result
|
||||
|
||||
if not user_api_key_dict.model_max_budget:
|
||||
return None
|
||||
budget_key: Final = "rpm_limit" if rate_limit_key == "model_rpm_limit" else "tpm_limit"
|
||||
model_limit: Final = {
|
||||
model: budget[budget_key]
|
||||
for model, budget in user_api_key_dict.model_max_budget.items()
|
||||
if isinstance(budget, dict) and budget.get(budget_key) is not None
|
||||
}
|
||||
return model_limit or None
|
||||
|
||||
|
||||
def get_key_model_rpm_limit(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
model_name: str | None = None,
|
||||
|
|
@ -989,20 +1015,9 @@ def get_key_model_rpm_limit(
|
|||
3. Team metadata (model_rpm_limit)
|
||||
4. Deployment default_api_key_rpm_limit (when model_name is provided)
|
||||
"""
|
||||
# 1. Check key metadata first (takes priority)
|
||||
if user_api_key_dict.metadata:
|
||||
result: Final = user_api_key_dict.metadata.get("model_rpm_limit")
|
||||
if result:
|
||||
return result
|
||||
|
||||
# 2. Check model_max_budget
|
||||
if user_api_key_dict.model_max_budget:
|
||||
model_rpm_limit: Final[dict[str, int]] = {}
|
||||
for model, budget in user_api_key_dict.model_max_budget.items():
|
||||
if isinstance(budget, dict) and budget.get("rpm_limit") is not None:
|
||||
model_rpm_limit[model] = budget["rpm_limit"]
|
||||
if model_rpm_limit:
|
||||
return model_rpm_limit
|
||||
key_own_limit: Final = get_key_own_model_rate_limit(user_api_key_dict, "model_rpm_limit")
|
||||
if key_own_limit is not None:
|
||||
return key_own_limit
|
||||
|
||||
# 3. Fallback to team metadata
|
||||
if user_api_key_dict.team_metadata:
|
||||
|
|
@ -1032,20 +1047,9 @@ def get_key_model_tpm_limit(
|
|||
3. Team metadata (model_tpm_limit)
|
||||
4. Deployment default_api_key_tpm_limit (when model_name is provided)
|
||||
"""
|
||||
# 1. Check key metadata first (takes priority)
|
||||
if user_api_key_dict.metadata:
|
||||
result: Final = user_api_key_dict.metadata.get("model_tpm_limit")
|
||||
if result:
|
||||
return result
|
||||
|
||||
# 2. Check model_max_budget (iterate per-model like RPM does)
|
||||
if user_api_key_dict.model_max_budget:
|
||||
model_tpm_limit: Final[dict[str, int]] = {}
|
||||
for model, budget in user_api_key_dict.model_max_budget.items():
|
||||
if isinstance(budget, dict) and budget.get("tpm_limit") is not None:
|
||||
model_tpm_limit[model] = budget["tpm_limit"]
|
||||
if model_tpm_limit:
|
||||
return model_tpm_limit
|
||||
key_own_limit: Final = get_key_own_model_rate_limit(user_api_key_dict, "model_tpm_limit")
|
||||
if key_own_limit is not None:
|
||||
return key_own_limit
|
||||
|
||||
# 3. Fallback to team metadata
|
||||
if user_api_key_dict.team_metadata:
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.auth.auth_utils import (
|
||||
ESTIMATED_OUTPUT_TOKENS_FIELD,
|
||||
get_estimated_output_tokens,
|
||||
get_key_own_model_rate_limit,
|
||||
get_key_tag_rpm_limit,
|
||||
get_model_rate_limit_from_metadata,
|
||||
)
|
||||
|
|
@ -2892,41 +2893,46 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
return batch_limiter
|
||||
return None
|
||||
|
||||
def _inherited_team_model_limit(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: str,
|
||||
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
|
||||
) -> int | None:
|
||||
"""Team per-model limit this key inherits: None when the key sets its own limit for the model."""
|
||||
team_limits: Final = get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", rate_limit_key)
|
||||
team_limit: Final = team_limits.get(requested_model) if team_limits else None
|
||||
if team_limit is None:
|
||||
return None
|
||||
key_own_limits: Final = get_key_own_model_rate_limit(user_api_key_dict, rate_limit_key)
|
||||
if key_own_limits and key_own_limits.get(requested_model) is not None:
|
||||
return None
|
||||
return team_limit
|
||||
|
||||
def _add_team_model_rate_limit_descriptor_from_metadata(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: str | None,
|
||||
descriptors: list[RateLimitDescriptor],
|
||||
) -> None:
|
||||
"""Add team model rate limit descriptor from team_metadata if applicable."""
|
||||
if (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") is not None
|
||||
or get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") is not None
|
||||
):
|
||||
_tpm_limit_for_team_model: Final = (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") or {}
|
||||
"""Add the team's per-model descriptor for the metrics the key does not override itself."""
|
||||
if requested_model is None:
|
||||
return
|
||||
team_rpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, requested_model, "model_rpm_limit")
|
||||
team_tpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, requested_model, "model_tpm_limit")
|
||||
if team_rpm_limit is None and team_tpm_limit is None:
|
||||
return
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_team",
|
||||
value=f"{user_api_key_dict.team_id}:{requested_model}",
|
||||
rate_limit={
|
||||
"requests_per_unit": team_rpm_limit,
|
||||
"tokens_per_unit": team_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
_rpm_limit_for_team_model: Final = (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") or {}
|
||||
)
|
||||
should_check_rate_limit: Final = (
|
||||
requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model
|
||||
)
|
||||
|
||||
if should_check_rate_limit and requested_model is not None:
|
||||
model_specific_tpm_limit: Final = _tpm_limit_for_team_model.get(requested_model)
|
||||
model_specific_rpm_limit: Final = _rpm_limit_for_team_model.get(requested_model)
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_team",
|
||||
value=f"{user_api_key_dict.team_id}:{requested_model}",
|
||||
rate_limit={
|
||||
"requests_per_unit": model_specific_rpm_limit,
|
||||
"tokens_per_unit": model_specific_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def _add_project_model_rate_limit_descriptor_from_metadata(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
get_key_mcp_rpm_limit,
|
||||
get_key_model_rpm_limit,
|
||||
get_key_model_tpm_limit,
|
||||
get_key_own_model_rate_limit,
|
||||
get_key_tag_rpm_limit,
|
||||
get_model_from_request,
|
||||
get_project_model_rpm_limit,
|
||||
|
|
@ -141,6 +142,35 @@ class TestLogOnceIfBudgetReservationDisabled:
|
|||
class TestGetKeyModelRpmLimit:
|
||||
"""Tests for get_key_model_rpm_limit function."""
|
||||
|
||||
def test_own_limit_excludes_team_metadata(self):
|
||||
"""A team-only limit is inherited, not owned: the key resolves it but does not override it."""
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="sk-123",
|
||||
metadata={"some_other_key": "value"},
|
||||
team_metadata={"model_rpm_limit": {"gpt-4": 50}, "model_tpm_limit": {"gpt-4": 500}},
|
||||
)
|
||||
assert get_key_model_rpm_limit(user_api_key_dict) == {"gpt-4": 50}
|
||||
assert get_key_own_model_rate_limit(user_api_key_dict, "model_rpm_limit") is None
|
||||
assert get_key_own_model_rate_limit(user_api_key_dict, "model_tpm_limit") is None
|
||||
|
||||
def test_own_limit_resolves_metadata_then_model_max_budget(self):
|
||||
from_metadata = UserAPIKeyAuth(
|
||||
api_key="sk-123",
|
||||
metadata={"model_rpm_limit": {"gpt-4": 100}},
|
||||
model_max_budget={"gpt-4": {"rpm_limit": 10, "tpm_limit": 1000}},
|
||||
team_metadata={"model_rpm_limit": {"gpt-4": 50}},
|
||||
)
|
||||
assert get_key_own_model_rate_limit(from_metadata, "model_rpm_limit") == {"gpt-4": 100}
|
||||
assert get_key_own_model_rate_limit(from_metadata, "model_tpm_limit") == {"gpt-4": 1000}
|
||||
|
||||
from_budget = UserAPIKeyAuth(
|
||||
api_key="sk-123",
|
||||
model_max_budget={"gpt-4": {"rpm_limit": 10}, "gpt-3.5-turbo": {"tpm_limit": 1000}},
|
||||
team_metadata={"model_rpm_limit": {"gpt-4": 50}},
|
||||
)
|
||||
assert get_key_own_model_rate_limit(from_budget, "model_rpm_limit") == {"gpt-4": 10}
|
||||
assert get_key_own_model_rate_limit(from_budget, "model_tpm_limit") == {"gpt-3.5-turbo": 1000}
|
||||
|
||||
def test_returns_key_metadata_when_present(self):
|
||||
"""Key metadata takes priority over team metadata."""
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
|
|
|
|||
|
|
@ -6311,7 +6311,7 @@ async def test_an_open_circuit_breaker_reads_the_sliding_window_locally_without_
|
|||
(
|
||||
{
|
||||
"team_id": "t",
|
||||
"metadata": {"model_rpm_limit": {"test-model": 100}},
|
||||
"metadata": {"model_rpm_limit": {"other-model": 100}},
|
||||
"team_metadata": {"model_rpm_limit": {"test-model": 1}},
|
||||
},
|
||||
{},
|
||||
|
|
@ -6529,3 +6529,71 @@ async def test_request_capacity_rejection_keeps_existing_redis_mirror():
|
|||
pytest.fail("rejection released another request's mirrored slot")
|
||||
assert exc.value.status_code == 429
|
||||
assert await cache.async_get_cache(counter_key, local_only=True) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key_limits",
|
||||
[
|
||||
{"metadata": {"model_rpm_limit": {"test-model": 3}}},
|
||||
{"model_max_budget": {"test-model": {"rpm_limit": 3}}},
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_model_rpm_override_takes_precedence_over_team_model_rpm_limit(key_limits):
|
||||
cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-key-override"),
|
||||
team_id="t",
|
||||
team_metadata={"model_rpm_limit": {"test-model": 1}},
|
||||
**key_limits,
|
||||
)
|
||||
|
||||
async def request():
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=auth, cache=cache, data={"model": "test-model"}, call_type="acompletion"
|
||||
)
|
||||
|
||||
for _ in range(3):
|
||||
await request()
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await request()
|
||||
assert exc.value.status_code == 429
|
||||
assert "model_per_key" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key_limits, override_key_gets_through",
|
||||
[
|
||||
({"model_rpm_limit": {"test-model": 10}}, False),
|
||||
({"model_rpm_limit": {"test-model": 10}, "model_tpm_limit": {"test-model": 5000}}, True),
|
||||
],
|
||||
ids=["rpm_only_override_still_shares_team_tpm", "rpm_and_tpm_override_leaves_team_tpm"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_model_rpm_override_keeps_team_model_tpm_limit(key_limits, override_key_gets_through):
|
||||
cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
|
||||
team_metadata = {"model_rpm_limit": {"test-model": 5}, "model_tpm_limit": {"test-model": 500}}
|
||||
sibling_key = UserAPIKeyAuth(api_key=hash_token("sk-sibling"), team_id="t", team_metadata=team_metadata)
|
||||
override_key = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-key-override"), team_id="t", metadata=key_limits, team_metadata=team_metadata
|
||||
)
|
||||
|
||||
async def request(auth):
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
cache=cache,
|
||||
data={"model": "test-model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 300},
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
await request(sibling_key)
|
||||
if override_key_gets_through:
|
||||
await request(override_key)
|
||||
return
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await request(override_key)
|
||||
assert exc.value.status_code == 429
|
||||
assert "model_per_team" in str(exc.value.detail)
|
||||
assert exc.value.headers["rate_limit_type"] == "tokens"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue