fix(router): enforce deployment budgets for dynamically added models by Sameerlite · Pull Request #29273 · BerriAI/litellm

Sameerlite · 2026-05-29T09:08:06Z

Summary

Register max_budget / budget_duration on RouterBudgetLimiting when deployments are added via add_deployment / upsert_deployment (e.g. /model/new), not only at router startup from model_list.
Lazily enable router_budget_limiting when the first budgeted deployment is added at runtime.
Use exact-type lookup for the router budget callback so _PROXY_VirtualKeyModelMaxBudgetLimiter is not mistaken for deployment budget enforcement.

Test plan

poetry run pytest tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py::test_register_deployment_budget_for_runtime_added_deployment -v
poetry run pytest tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py::test_router_add_deployment_registers_deployment_budget -v
Add model via /model/new with tiny max_budget; verify second request fails with budget error
Confirm startup model_list budget behavior unchanged

Fixes LIT-3043
Fixes #25799

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 <[email protected]>

codecov · 2026-05-29T09:11:17Z

Codecov Report

❌ Patch coverage is 77.27273% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/router.py	75.00%	8 Missing ⚠️
...proxy/spend_tracking/spend_management_endpoints.py	66.66%	1 Missing ⚠️
litellm/router_strategy/budget_limiter.py	87.50%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

Remove unused RouterBudgetLimiting import and add router unit tests for deployment budget helper methods required by router_code_coverage. Co-authored-by: Cursor <[email protected]>

greptile-apps · 2026-05-29T09:12:18Z

Greptile Summary

This PR fixes a gap where deployments added at runtime via add_deployment/upsert_deployment (e.g., /model/new) never had their max_budget/budget_duration registered with the RouterBudgetLimiting callback, causing budget limits to be silently ignored for dynamically-added models.

router.py: Adds _sync_deployment_budget_config (called from add_deployment) to register or unregister a deployment's budget config on every add/upsert, and _get_router_deployment_budget_limiter using exact-type matching so _PROXY_VirtualKeyModelMaxBudgetLimiter (a subclass) is never mistakenly treated as the router's deployment budget enforcer. Lazy-initialises RouterBudgetLimiting when the first budgeted deployment arrives if it wasn't already present.
budget_limiter.py: Adds register_deployment_budget and unregister_deployment_budget public methods; the latter nulls deployment_budget_config when the map becomes empty, keeping the filtering hot-path O(1).
spend_management_endpoints.py: Hoists the RouterBudgetLimiting lookup outside the provider loop and replaces the per-iteration isinstance scan with _get_router_deployment_budget_limiter(), eliminating a redundant call on each loop iteration.

Confidence Score: 5/5

Safe to merge — the change is additive, all new code paths are covered by unit tests, and no existing behaviour is altered for routers that don't use runtime deployment additions.

The fix is well-scoped: _sync_deployment_budget_config is called at the end of add_deployment after all data structures are updated, so there is no window where the deployment is live but untracked. The upsert path correctly cleans up stale budget entries via the same method (verified by test_upsert_deployment_clears_stale_budget_config). The exact-type guard in _get_router_deployment_budget_limiter correctly excludes the _PROXY_VirtualKeyModelMaxBudgetLimiter subclass, which never initialises provider_budget_config. New tests are mock-only and follow the existing patterns in both test files.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/router.py	Adds `_sync_deployment_budget_config`, `_deployment_has_budget_limits`, and `_get_router_deployment_budget_limiter` to wire budget registration into `add_deployment` and `delete_deployment`. Lazy init of `RouterBudgetLimiting` when the first budgeted deployment arrives at runtime is correctly handled. Upsert path also cleans up stale budget entries via `_sync_deployment_budget_config`.
litellm/router_strategy/budget_limiter.py	Adds `register_deployment_budget` (delegates to `_init_deployment_budgets`) and `unregister_deployment_budget` (removes entry, nulls map when empty). Minor type-hint cleanup of `model_list` signature. Logic is sound and consistent with existing `_init_deployment_budgets` behavior.
litellm/proxy/spend_tracking/spend_management_endpoints.py	Hoists `_get_router_deployment_budget_limiter()` lookup outside the provider loop and removes the redundant `RouterBudgetLimiting` import. Cleaner than the previous per-iteration `isinstance` scan.
litellm/proxy/hooks/model_max_budget_limiter.py	Adds explicit `self.deployment_budget_config = None` in `_PROXY_VirtualKeyModelMaxBudgetLimiter.__init__` so the attribute exists when accessed, preventing `AttributeError` since the subclass never calls `super().__init__()`.
tests/router_unit_tests/test_router_helper_utils.py	Adds unit tests for `_deployment_has_budget_limits`, `_sync_deployment_budget_config` (add and clear paths), and the upsert stale-budget scenario. One existing assertion reformatted (no semantic change). Tests are well-structured and cover the new code paths.
tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py	Adds two tests directly exercising `register_deployment_budget`/`unregister_deployment_budget` on `RouterBudgetLimiting` and the full `Router.add_deployment` → budget config path. Both tests use the `disable_budget_sync` fixture and mock `asyncio.create_task`, matching the file's existing test style.

_{Reviews (3): Last reviewed commit: "fix black" | Re-trigger Greptile}

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 <[email protected]>

Sameerlite · 2026-05-29T09:15:58Z

@greptileai

Sameerlite · 2026-05-29T14:18:13Z

@greptileai

mateo-berri

LGTM; thanks

…erriAI#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 <[email protected]> * 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 <[email protected]> * 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 <[email protected]> * Fix mypy * fix black --------- Co-authored-by: Cursor <[email protected]>

fix(router): enforce deployment budgets for dynamically added models

3dd591e

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 <[email protected]>

fix(router): address CI lint and router coverage for budget sync helpers

6911053

Remove unused RouterBudgetLimiting import and add router unit tests for deployment budget helper methods required by router_code_coverage. Co-authored-by: Cursor <[email protected]>

fix(router): clear stale deployment budget on upsert without limits

ee2f79f

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 <[email protected]>

Sameerlite requested a review from mateo-berri May 29, 2026 13:31

Sameerlite added 2 commits May 29, 2026 19:19

Fix mypy

a705731

fix black

e1a476d

mateo-berri approved these changes May 30, 2026

View reviewed changes

mateo-berri merged commit 0ffab87 into litellm_internal_staging May 30, 2026
116 of 118 checks passed

mateo-berri deleted the litellm_fix_dynamic_model_budget branch May 30, 2026 02:43

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

fix(router): enforce deployment budgets for dynamically added models#29273

fix(router): enforce deployment budgets for dynamically added models#29273
mateo-berri merged 5 commits into
litellm_internal_stagingfrom
litellm_fix_dynamic_model_budget

Sameerlite commented May 29, 2026 •

edited

Loading

Uh oh!

codecov Bot commented May 29, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented May 29, 2026 •

edited

Loading

Important Files Changed

Uh oh!

Sameerlite commented May 29, 2026

Uh oh!

Sameerlite commented May 29, 2026

Uh oh!

mateo-berri left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Uh oh!

Conversation

Sameerlite commented May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

codecov Bot commented May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

greptile-apps Bot commented May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Uh oh!

Sameerlite commented May 29, 2026

Uh oh!

Sameerlite commented May 29, 2026

Uh oh!

mateo-berri left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Sameerlite commented May 29, 2026 •

edited

Loading

codecov Bot commented May 29, 2026 •

edited

Loading

greptile-apps Bot commented May 29, 2026 •

edited

Loading