From 36dc893770fa69cfcc014e5c535b87d58ca12c89 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Thu, 19 Mar 2026 01:43:27 -0400 Subject: [PATCH] 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) --- .../sensitive_data_masker.py | 4 +- litellm/proxy/auth/auth_utils.py | 44 ++++++++++------ .../common_utils/openai_endpoint_utils.py | 11 +++- .../proxy/hooks/parallel_request_limiter.py | 14 ++++-- .../proxy/auth/test_auth_utils.py | 50 +++++++++++++++++++ .../proxy/test_model_info_default_limits.py | 3 ++ 6 files changed, 101 insertions(+), 25 deletions(-) diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index f22cfa11a3..663c3fac80 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -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 diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 235b217610..ace39c05ff 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -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( diff --git a/litellm/proxy/common_utils/openai_endpoint_utils.py b/litellm/proxy/common_utils/openai_endpoint_utils.py index 6df5491f37..7e5c83500a 100644 --- a/litellm/proxy/common_utils/openai_endpoint_utils.py +++ b/litellm/proxy/common_utils/openai_endpoint_utils.py @@ -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 diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index c7bfc27d6b..48bf255ac1 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index be4db666a0..d64f17e70f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -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} diff --git a/tests/test_litellm/proxy/test_model_info_default_limits.py b/tests/test_litellm/proxy/test_model_info_default_limits.py index e749c84dfb..907f2390a8 100644 --- a/tests/test_litellm/proxy/test_model_info_default_limits.py +++ b/tests/test_litellm/proxy/test_model_info_default_limits.py @@ -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)