From 0ffab87da418604cfa1ec46c80a8b926f2da842e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Sat, 30 May 2026 08:13:14 +0530 Subject: [PATCH] fix(router): enforce deployment budgets for dynamically added models (#29273) * fix(router): enforce deployment budgets for dynamically added models Register deployment max_budget/budget_duration when models are added via upsert_deployment (e.g. /model/new) so RouterBudgetLimiting matches startup model_list behavior. Co-authored-by: Cursor * fix(router): address CI lint and router coverage for budget sync helpers Remove unused RouterBudgetLimiting import and add router unit tests for deployment budget helper methods required by router_code_coverage. Co-authored-by: Cursor * fix(router): clear stale deployment budget on upsert without limits Unregister deployment budget config when max_budget/budget_duration are removed, including upsert replace paths. Hoist provider budget logger lookup outside the provider loop. Co-authored-by: Cursor * Fix mypy * fix black --------- Co-authored-by: Cursor --- .../proxy/hooks/model_max_budget_limiter.py | 1 + .../spend_management_endpoints.py | 15 +-- litellm/router.py | 56 ++++++++ litellm/router_strategy/budget_limiter.py | 24 +++- .../test_router_helper_utils.py | 126 +++++++++++++++++- .../test_budget_limiter_hotpath.py | 69 ++++++++++ 6 files changed, 270 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 9286424878..3c96067da8 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -28,6 +28,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): def __init__(self, dual_cache: DualCache): self.dual_cache = dual_cache self.redis_increment_operation_queue = [] + self.deployment_budget_config = None async def is_key_within_model_budget( self, diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 36beb5e9ab..c7ecd64c0f 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -21,7 +21,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team_and_customer, ) from litellm.proxy.utils import handle_exception_on_proxy -from litellm.router_strategy.budget_limiter import RouterBudgetLimiting if TYPE_CHECKING: from litellm.proxy.proxy_server import PrismaClient @@ -3149,18 +3148,12 @@ async def provider_budgets() -> ProviderBudgetResponse: "No provider budget config found. Please set a provider budget config in the router settings. https://docs.litellm.ai/docs/proxy/provider_budget_routing" ) + router_budget_logger = llm_router._get_router_deployment_budget_limiter() + if router_budget_logger is None: + raise ValueError("No router budget logger found") + provider_budget_response_dict: Dict[str, ProviderBudgetResponseObject] = {} for _provider, _budget_info in provider_budget_config.items(): - router_budget_logger = next( - ( - cb - for cb in (llm_router.optional_callbacks or []) - if isinstance(cb, RouterBudgetLimiting) - ), - None, - ) - if router_budget_logger is None: - raise ValueError("No router budget logger found") _provider_spend = ( await router_budget_logger._get_current_provider_spend(_provider) or 0.0 ) diff --git a/litellm/router.py b/litellm/router.py index 2c61031da3..d60c39ca40 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1753,11 +1753,14 @@ class Router: if pre_call_check == "prompt_caching": _callback = PromptCachingDeploymentCheck(cache=self.cache) elif pre_call_check == "router_budget_limiting": + if self._get_router_deployment_budget_limiter() is not None: + continue _callback = RouterBudgetLimiting( dual_cache=self.cache, provider_budget_config=self.provider_budget_config, model_list=self.model_list, ) + self.router_budget_logger = _callback elif pre_call_check == "enforce_model_rate_limits": _callback = ModelRateLimitingCheck(dual_cache=self.cache) @@ -8529,6 +8532,7 @@ class Router: model=_deployment, model_id=deployment.model_info.id ) self.model_names.add(deployment.model_name) + self._sync_deployment_budget_config(deployment=deployment) return deployment def _update_deployment_indices_after_removal( @@ -8717,12 +8721,64 @@ class Router: self._update_deployment_indices_after_removal( model_id=id, removal_idx=deployment_idx ) + _budget_limiter = self._get_router_deployment_budget_limiter() + if _budget_limiter is not None: + _budget_limiter.unregister_deployment_budget(model_id=id) return item else: return None except Exception: return None + def _get_router_deployment_budget_limiter( + self, + ) -> Optional[RouterBudgetLimiting]: + """ + Return the router's deployment-budget callback. + + Uses exact-type matching so proxy subclasses (e.g. virtual-key model budgets) + registered on litellm.callbacks are not mistaken for router deployment budgets. + """ + if self.router_budget_logger is not None: + return self.router_budget_logger + + if self.optional_callbacks: + for _cb in self.optional_callbacks: + if type(_cb) is RouterBudgetLimiting: + self.router_budget_logger = _cb + return _cb + return None + + def _deployment_has_budget_limits(self, deployment: Deployment) -> bool: + return ( + deployment.litellm_params.get("max_budget") is not None + and deployment.litellm_params.get("budget_duration") is not None + and deployment.model_info.id is not None + ) + + def _sync_deployment_budget_config(self, deployment: Deployment) -> None: + model_id = deployment.model_info.id + if model_id is None: + return + + _budget_limiter = self._get_router_deployment_budget_limiter() + + if not self._deployment_has_budget_limits(deployment=deployment): + if _budget_limiter is not None: + _budget_limiter.unregister_deployment_budget(model_id=model_id) + return + + if _budget_limiter is None: + self.add_optional_pre_call_checks( + optional_pre_call_checks=["router_budget_limiting"] + ) + _budget_limiter = self._get_router_deployment_budget_limiter() + + if _budget_limiter is not None: + _budget_limiter.register_deployment_budget( + deployment=deployment.to_json(exclude_none=True) + ) + def get_deployment(self, model_id: str) -> Optional[Deployment]: """ Returns -> Deployment or None diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index da41577e99..f677e40b93 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -96,9 +96,7 @@ class RouterBudgetLimiting(CustomLogger): self, dual_cache: DualCache, provider_budget_config: Optional[dict], - model_list: Optional[ - Union[List[DeploymentTypedDict], List[Dict[str, Any]]] - ] = None, + model_list: Optional[List[Union[DeploymentTypedDict, Dict[str, Any]]]] = None, ): self.dual_cache = dual_cache self.redis_increment_operation_queue: List[RedisPipelineIncrementOperation] = [] @@ -854,9 +852,7 @@ class RouterBudgetLimiting(CustomLogger): def _init_deployment_budgets( self, - model_list: Optional[ - Union[List[DeploymentTypedDict], List[Dict[str, Any]]] - ] = None, + model_list: Optional[List[Union[DeploymentTypedDict, Dict[str, Any]]]] = None, ): if model_list is None: return @@ -887,6 +883,22 @@ class RouterBudgetLimiting(CustomLogger): f"Initialized Deployment Budget Config: {self.deployment_budget_config}" ) + def register_deployment_budget( + self, + deployment: Union[Dict[str, Any], DeploymentTypedDict], + ) -> None: + """ + Register or refresh deployment-level budget config for a runtime-added deployment. + """ + self._init_deployment_budgets(model_list=[deployment]) + + def unregister_deployment_budget(self, model_id: str) -> None: + if self.deployment_budget_config is None: + return + self.deployment_budget_config.pop(model_id, None) + if len(self.deployment_budget_config) == 0: + self.deployment_budget_config = None + def _init_tag_budgets(self): if litellm.tag_budget_config is None: return diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 65d9d6b925..83d6d56df4 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -14,7 +14,7 @@ import litellm from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload from litellm.types.utils import StandardLoggingPayload -from litellm.types.router import Deployment, LiteLLM_Params +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @pytest.fixture @@ -997,9 +997,7 @@ def test_filter_cooldown_deployments(model_list): healthy_deployments=router._get_all_deployments(model_name="gpt-5-mini"), # type: ignore cooldown_deployments=[], ) - assert len(deployments) == len( - router._get_all_deployments(model_name="gpt-5-mini") - ) + assert len(deployments) == len(router._get_all_deployments(model_name="gpt-5-mini")) def test_track_deployment_metrics(model_list): @@ -2379,3 +2377,123 @@ def test_get_router_model_info_with_deployment_object(): # Verify we got valid model info back assert model_info is not None assert isinstance(model_info, dict) + + +def test_deployment_has_budget_limits(): + router = Router(model_list=[]) + + with_budget = Deployment( + model_name="budgeted-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o-mini", + max_budget=0.001, + budget_duration="1d", + ), + model_info=ModelInfo(id="budget-deployment-id"), + ) + without_budget = Deployment( + model_name="unbudgeted-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"), + model_info=ModelInfo(id="no-budget-deployment-id"), + ) + + assert router._deployment_has_budget_limits(deployment=with_budget) is True + assert router._deployment_has_budget_limits(deployment=without_budget) is False + + +def test_sync_deployment_budget_config(monkeypatch): + import asyncio + + monkeypatch.setattr(asyncio, "create_task", lambda coro: None) + + router = Router(model_list=[], optional_pre_call_checks=[]) + deployment = Deployment( + model_name="dynamic-budget-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o-mini", + api_key="fake-key", + max_budget=0.000000000001, + budget_duration="1d", + ), + model_info=ModelInfo(id="runtime-budget-deployment"), + ) + + router._sync_deployment_budget_config(deployment=deployment) + + budget_limiter = router._get_router_deployment_budget_limiter() + assert budget_limiter is not None + config = budget_limiter._get_budget_config_for_deployment( + "runtime-budget-deployment" + ) + assert config is not None + assert config.max_budget == 0.000000000001 + + +def test_sync_deployment_budget_config_clears_removed_limits(monkeypatch): + import asyncio + + monkeypatch.setattr(asyncio, "create_task", lambda coro: None) + + router = Router(model_list=[], optional_pre_call_checks=[]) + model_id = "runtime-budget-deployment" + budgeted = Deployment( + model_name="dynamic-budget-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o-mini", + api_key="fake-key", + max_budget=0.000000000001, + budget_duration="1d", + ), + model_info=ModelInfo(id=model_id), + ) + unbudgeted = Deployment( + model_name="dynamic-budget-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o-mini", + api_key="fake-key", + ), + model_info=ModelInfo(id=model_id), + ) + + router._sync_deployment_budget_config(deployment=budgeted) + budget_limiter = router._get_router_deployment_budget_limiter() + assert budget_limiter is not None + assert budget_limiter._get_budget_config_for_deployment(model_id) is not None + + router._sync_deployment_budget_config(deployment=unbudgeted) + assert budget_limiter._get_budget_config_for_deployment(model_id) is None + + +def test_upsert_deployment_clears_stale_budget_config(monkeypatch): + import asyncio + + monkeypatch.setattr(asyncio, "create_task", lambda coro: None) + + router = Router(model_list=[], optional_pre_call_checks=[]) + model_id = "upsert-budget-deployment" + budgeted = Deployment( + model_name="dynamic-budget-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o-mini", + api_key="fake-key", + max_budget=0.000000000001, + budget_duration="1d", + ), + model_info=ModelInfo(id=model_id), + ) + unbudgeted = Deployment( + model_name="dynamic-budget-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o-mini", + api_key="fake-key", + ), + model_info=ModelInfo(id=model_id), + ) + + router.upsert_deployment(deployment=budgeted) + budget_limiter = router._get_router_deployment_budget_limiter() + assert budget_limiter is not None + assert budget_limiter._get_budget_config_for_deployment(model_id) is not None + + router.upsert_deployment(deployment=unbudgeted) + assert budget_limiter._get_budget_config_for_deployment(model_id) is None diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py index d4fb9084c8..36fa38bacb 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -234,3 +234,72 @@ async def test_get_llm_provider_for_deployment_matches_legacy_behavior( legacy_provider = _legacy_provider_resolution(deployment) assert current_provider == legacy_provider + + +def test_register_deployment_budget_for_runtime_added_deployment( + disable_budget_sync, monkeypatch +): + import asyncio + + monkeypatch.setattr(asyncio, "create_task", lambda coro: None) + budget_limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={}, + ) + model_id = "dynamic-deployment-id" + budget_limiter.register_deployment_budget( + deployment={ + "model_name": "dynamic-budget-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "max_budget": 0.000000000001, + "budget_duration": "1d", + }, + "model_info": {"id": model_id}, + } + ) + + config = budget_limiter._get_budget_config_for_deployment(model_id) + assert config is not None + assert config.max_budget == 0.000000000001 + assert config.budget_duration == "1d" + + budget_limiter.unregister_deployment_budget(model_id=model_id) + assert budget_limiter._get_budget_config_for_deployment(model_id) is None + + +def test_router_add_deployment_registers_deployment_budget( + disable_budget_sync, monkeypatch +): + import asyncio + + from litellm import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + monkeypatch.setattr(asyncio, "create_task", lambda coro: None) + + router = Router( + model_list=[], + optional_pre_call_checks=[], + ) + + router.add_deployment( + deployment=Deployment( + model_name="dynamic-budget-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o-mini", + api_key="fake-key", + max_budget=0.000000000001, + budget_duration="1d", + ), + model_info=ModelInfo(id="runtime-budget-deployment"), + ) + ) + + budget_limiter = router._get_router_deployment_budget_limiter() + assert budget_limiter is not None + config = budget_limiter._get_budget_config_for_deployment( + "runtime-budget-deployment" + ) + assert config is not None + assert config.max_budget == 0.000000000001