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:
Ephrim Stanley
2026-03-19 01:43:27 -04:00
co-authored by Claude
parent cac685014f
commit 36dc893770
6 changed files with 101 additions and 25 deletions
@@ -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
+28 -16
View File
@@ -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)