mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-24 10:27:13 +00:00
fix: address review feedback on default tpm/rpm limits
- Use min() across all matching deployments instead of first-wins when resolving default_api_key_tpm/rpm_limit for a model group, so load-balanced setups with different per-deployment limits always apply the most conservative value - Replace the global SensitiveDataMasker non_sensitive_overrides change with a targeted excluded_keys set at the remove_sensitive_info_from_deployment call site, avoiding unintended suppression of other fields - Update the v1 parallel request limiter to pass model_name to get_key_model_tpm/rpm_limit so deployment defaults apply there too - Add 4 tests covering multi-deployment min semantics Co-Authored-By: Claude (claude-sonnet-4-6) <noreply@anthropic.com>
This commit is contained in:
@@ -30,9 +30,7 @@ class SensitiveDataMasker:
|
||||
# If any key segment matches one of these, the key is not considered sensitive
|
||||
# even if it also matches a sensitive pattern. For example, "input_cost_per_token"
|
||||
# contains "token" but "cost" overrides that — it's a pricing field, not a secret.
|
||||
# Similarly, "*_limit" fields (tpm_limit, rpm_limit, etc.) are rate/budget caps,
|
||||
# not credentials, even though their names may contain "key" (e.g. default_api_key_tpm_limit).
|
||||
self.non_sensitive_overrides = non_sensitive_overrides or {"cost", "limit"}
|
||||
self.non_sensitive_overrides = non_sensitive_overrides or {"cost"}
|
||||
|
||||
self.visible_prefix = visible_prefix
|
||||
self.visible_suffix = visible_suffix
|
||||
|
||||
@@ -541,8 +541,13 @@ def bytes_to_mb(bytes_value: int):
|
||||
# helpers used by parallel request limiter to handle model rpm/tpm limits for a given api key
|
||||
def _get_deployment_default_rpm_limit(model_name: str) -> Optional[int]:
|
||||
"""
|
||||
Return the default_api_key_rpm_limit configured on the deployment for model_name,
|
||||
or None if not set.
|
||||
Return the default_api_key_rpm_limit for model_name.
|
||||
|
||||
When multiple deployments share the same model name, returns the minimum
|
||||
across all deployments that have the field set. This is the safest choice
|
||||
for load-balanced setups: it ensures no deployment is over-consumed
|
||||
regardless of which one actually serves a given request.
|
||||
Returns None if no deployment has the field set.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
@@ -551,18 +556,24 @@ def _get_deployment_default_rpm_limit(model_name: str) -> Optional[int]:
|
||||
deployments = llm_router.get_model_list(model_name=model_name)
|
||||
if not deployments:
|
||||
return None
|
||||
for deployment in deployments:
|
||||
litellm_params = deployment.get("litellm_params", {})
|
||||
limit = litellm_params.get("default_api_key_rpm_limit")
|
||||
if limit is not None:
|
||||
return int(limit)
|
||||
return None
|
||||
limits = [
|
||||
int(deployment.get("litellm_params", {}).get("default_api_key_rpm_limit"))
|
||||
for deployment in deployments
|
||||
if deployment.get("litellm_params", {}).get("default_api_key_rpm_limit")
|
||||
is not None
|
||||
]
|
||||
return min(limits) if limits else None
|
||||
|
||||
|
||||
def _get_deployment_default_tpm_limit(model_name: str) -> Optional[int]:
|
||||
"""
|
||||
Return the default_api_key_tpm_limit configured on the deployment for model_name,
|
||||
or None if not set.
|
||||
Return the default_api_key_tpm_limit for model_name.
|
||||
|
||||
When multiple deployments share the same model name, returns the minimum
|
||||
across all deployments that have the field set. This is the safest choice
|
||||
for load-balanced setups: it ensures no deployment is over-consumed
|
||||
regardless of which one actually serves a given request.
|
||||
Returns None if no deployment has the field set.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
@@ -571,12 +582,13 @@ def _get_deployment_default_tpm_limit(model_name: str) -> Optional[int]:
|
||||
deployments = llm_router.get_model_list(model_name=model_name)
|
||||
if not deployments:
|
||||
return None
|
||||
for deployment in deployments:
|
||||
litellm_params = deployment.get("litellm_params", {})
|
||||
limit = litellm_params.get("default_api_key_tpm_limit")
|
||||
if limit is not None:
|
||||
return int(limit)
|
||||
return None
|
||||
limits = [
|
||||
int(deployment.get("litellm_params", {}).get("default_api_key_tpm_limit"))
|
||||
for deployment in deployments
|
||||
if deployment.get("litellm_params", {}).get("default_api_key_tpm_limit")
|
||||
is not None
|
||||
]
|
||||
return min(limits) if limits else None
|
||||
|
||||
|
||||
def get_key_model_rpm_limit(
|
||||
|
||||
@@ -32,8 +32,17 @@ def remove_sensitive_info_from_deployment(
|
||||
deployment_dict["litellm_params"].pop("aws_access_key_id", None)
|
||||
deployment_dict["litellm_params"].pop("aws_secret_access_key", None)
|
||||
|
||||
# Rate-limit config fields must never be masked — they are integers, not credentials.
|
||||
# The field names contain "key" which matches the masker's sensitive pattern, so we
|
||||
# explicitly exclude them here rather than widening the global non_sensitive_overrides.
|
||||
_rate_limit_config_keys = {
|
||||
"default_api_key_tpm_limit",
|
||||
"default_api_key_rpm_limit",
|
||||
}
|
||||
_excluded = (excluded_keys or set()) | _rate_limit_config_keys
|
||||
|
||||
deployment_dict["litellm_params"] = SENSITIVE_DATA_MASKER.mask_dict(
|
||||
deployment_dict["litellm_params"], excluded_keys=excluded_keys
|
||||
deployment_dict["litellm_params"], excluded_keys=_excluded
|
||||
)
|
||||
|
||||
return deployment_dict
|
||||
|
||||
@@ -295,16 +295,20 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
|
||||
)
|
||||
|
||||
# Check if request under RPM/TPM per model for a given API Key
|
||||
_model = data.get("model", None)
|
||||
if (
|
||||
get_key_model_tpm_limit(user_api_key_dict) is not None
|
||||
or get_key_model_rpm_limit(user_api_key_dict) is not None
|
||||
get_key_model_tpm_limit(user_api_key_dict, model_name=_model) is not None
|
||||
or get_key_model_rpm_limit(user_api_key_dict, model_name=_model) is not None
|
||||
):
|
||||
_model = data.get("model", None)
|
||||
request_count_api_key = (
|
||||
f"{api_key}::{_model}::{precise_minute}::request_count"
|
||||
)
|
||||
_tpm_limit_for_key_model = get_key_model_tpm_limit(user_api_key_dict)
|
||||
_rpm_limit_for_key_model = get_key_model_rpm_limit(user_api_key_dict)
|
||||
_tpm_limit_for_key_model = get_key_model_tpm_limit(
|
||||
user_api_key_dict, model_name=_model
|
||||
)
|
||||
_rpm_limit_for_key_model = get_key_model_rpm_limit(
|
||||
user_api_key_dict, model_name=_model
|
||||
)
|
||||
tpm_limit_for_model = None
|
||||
rpm_limit_for_model = None
|
||||
|
||||
|
||||
@@ -387,6 +387,31 @@ class TestDeploymentDefaultRpmLimit:
|
||||
result = get_key_model_rpm_limit(user_api_key_dict, model_name="model1")
|
||||
assert result is None
|
||||
|
||||
def test_returns_minimum_across_multiple_deployments(self):
|
||||
"""When multiple deployments share a model name, the minimum rpm limit is used."""
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-123")
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_list.return_value = [
|
||||
_make_deployment_dict("model1", rpm=200),
|
||||
_make_deployment_dict("model1", rpm=50),
|
||||
_make_deployment_dict("model1", rpm=150),
|
||||
]
|
||||
with patch(_ROUTER_PATCH, mock_router):
|
||||
result = get_key_model_rpm_limit(user_api_key_dict, model_name="model1")
|
||||
assert result == {"model1": 50}
|
||||
|
||||
def test_ignores_deployments_without_default_when_others_have_it(self):
|
||||
"""Deployments missing the field are skipped; min is taken over those that have it."""
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-123")
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_list.return_value = [
|
||||
_make_deployment_dict("model1"), # no rpm default
|
||||
_make_deployment_dict("model1", rpm=75),
|
||||
]
|
||||
with patch(_ROUTER_PATCH, mock_router):
|
||||
result = get_key_model_rpm_limit(user_api_key_dict, model_name="model1")
|
||||
assert result == {"model1": 75}
|
||||
|
||||
|
||||
class TestDeploymentDefaultTpmLimit:
|
||||
"""Tests for deployment default_api_key_tpm_limit fallback in get_key_model_tpm_limit."""
|
||||
@@ -444,3 +469,28 @@ class TestDeploymentDefaultTpmLimit:
|
||||
with patch(_ROUTER_PATCH, None):
|
||||
result = get_key_model_tpm_limit(user_api_key_dict, model_name="model1")
|
||||
assert result is None
|
||||
|
||||
def test_returns_minimum_across_multiple_deployments(self):
|
||||
"""When multiple deployments share a model name, the minimum tpm limit is used."""
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-123")
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_list.return_value = [
|
||||
_make_deployment_dict("model1", tpm=1000),
|
||||
_make_deployment_dict("model1", tpm=300),
|
||||
_make_deployment_dict("model1", tpm=700),
|
||||
]
|
||||
with patch(_ROUTER_PATCH, mock_router):
|
||||
result = get_key_model_tpm_limit(user_api_key_dict, model_name="model1")
|
||||
assert result == {"model1": 300}
|
||||
|
||||
def test_ignores_deployments_without_default_when_others_have_it(self):
|
||||
"""Deployments missing the field are skipped; min is taken over those that have it."""
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-123")
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_list.return_value = [
|
||||
_make_deployment_dict("model1"), # no tpm default
|
||||
_make_deployment_dict("model1", tpm=400),
|
||||
]
|
||||
with patch(_ROUTER_PATCH, mock_router):
|
||||
result = get_key_model_tpm_limit(user_api_key_dict, model_name="model1")
|
||||
assert result == {"model1": 400}
|
||||
|
||||
@@ -82,6 +82,9 @@ class TestModelInfoDefaultLimitsInResponse:
|
||||
"""
|
||||
default_api_key_tpm_limit / default_api_key_rpm_limit must not be
|
||||
treated as sensitive and must survive remove_sensitive_info_from_deployment.
|
||||
They contain "key" which normally triggers masking; the call site explicitly
|
||||
excludes these two fields via excluded_keys rather than widening the global
|
||||
non_sensitive_overrides.
|
||||
"""
|
||||
deployment = _make_deployment("model1", default_tpm=100, default_rpm=200)
|
||||
model_dict = deployment.model_dump(exclude_none=True)
|
||||
|
||||
Reference in New Issue
Block a user