mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(proxy): avoid duplicate model_per_team descriptors in v3 limiter
Team model TPM/RPM from metadata was merged into descriptors twice (_create_rate_limit_descriptors plus async_pre_call_hook), doubling increments on the same model_per_team counter. Remove redundant path and add regression tests. Made-with: Cursor
This commit is contained in:
parent
609454d0f1
commit
6909fd248a
2 changed files with 63 additions and 60 deletions
|
|
@ -1122,59 +1122,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
return batch_limiter
|
||||
return None
|
||||
|
||||
def _add_team_model_rate_limit_descriptor_from_metadata(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: Optional[str],
|
||||
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 = (
|
||||
get_model_rate_limit_from_metadata(
|
||||
user_api_key_dict, "team_metadata", "model_tpm_limit"
|
||||
)
|
||||
or {}
|
||||
)
|
||||
_rpm_limit_for_team_model = (
|
||||
get_model_rate_limit_from_metadata(
|
||||
user_api_key_dict, "team_metadata", "model_rpm_limit"
|
||||
)
|
||||
or {}
|
||||
)
|
||||
should_check_rate_limit = (
|
||||
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 = _tpm_limit_for_team_model.get(
|
||||
requested_model
|
||||
)
|
||||
model_specific_rpm_limit = _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,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -1332,13 +1279,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
model_has_failures=model_has_failures,
|
||||
)
|
||||
|
||||
# Add team model rate limits from team_metadata
|
||||
self._add_team_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
|
||||
# Project Level Rate Limits
|
||||
self._add_project_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
|
|||
|
|
@ -2685,3 +2685,66 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3():
|
|||
assert (
|
||||
"model_per_project" not in descriptor_keys
|
||||
), f"model_per_project should not be added for unrelated model, got: {descriptor_keys}"
|
||||
|
||||
|
||||
def test_team_metadata_model_tpm_limit_single_model_per_team_descriptor():
|
||||
"""
|
||||
team_metadata.model_tpm_limit must not produce duplicate model_per_team descriptors.
|
||||
Duplicates caused the same counter to increment twice per request (effective limit halved).
|
||||
"""
|
||||
local_cache = DualCache()
|
||||
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache)
|
||||
)
|
||||
user = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-team-model-tpm-dedupe"),
|
||||
team_id="team-dedupe-test",
|
||||
team_metadata={"model_tpm_limit": {"gpt-4o-mini": 100}},
|
||||
)
|
||||
data = {"model": "gpt-4o-mini"}
|
||||
descriptors = parallel_request_handler._create_rate_limit_descriptors(
|
||||
user_api_key_dict=user,
|
||||
data=data,
|
||||
rpm_limit_type=None,
|
||||
tpm_limit_type=None,
|
||||
model_has_failures=False,
|
||||
)
|
||||
mpt = [d for d in descriptors if d["key"] == "model_per_team"]
|
||||
assert len(mpt) == 1, (
|
||||
f"expected exactly one model_per_team descriptor, got {len(mpt)}: {mpt}"
|
||||
)
|
||||
assert mpt[0]["value"] == "team-dedupe-test:gpt-4o-mini"
|
||||
assert mpt[0]["rate_limit"]["tokens_per_unit"] == 100
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_metadata_model_tpm_pre_call_not_double_incremented():
|
||||
"""
|
||||
With team model TPM = N, pre-call should allow N increments on model_per_team
|
||||
(one per request), not N/2 from duplicate descriptors.
|
||||
"""
|
||||
local_cache = DualCache()
|
||||
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache)
|
||||
)
|
||||
user = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-team-model-tpm-pre"),
|
||||
team_id="team-pre-dedupe",
|
||||
team_metadata={"model_tpm_limit": {"gpt-4o-mini": 10}},
|
||||
)
|
||||
data = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "x"}]}
|
||||
allowed = 0
|
||||
for _ in range(20):
|
||||
try:
|
||||
await parallel_request_handler.async_pre_call_hook(
|
||||
user_api_key_dict=user,
|
||||
cache=local_cache,
|
||||
data=dict(data),
|
||||
call_type="completion",
|
||||
)
|
||||
allowed += 1
|
||||
except HTTPException:
|
||||
break
|
||||
assert allowed == 10, (
|
||||
f"expected 10 pre-calls before model_per_team TPM block, got {allowed}"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue