scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(router): guard registry cleanup to strategy deployments; retire adaptive hooks

Browse source 
Review follow-ups:
- Only deregister when the removed deployment's litellm_params.model
  carries the auto_router/ prefix — a regular deployment that merely
  shares a strategy router's public model_name can no longer evict the
  shared router when deleted.
- Removing an adaptive router now also retires its
  AdaptiveRouterPostCallHook from the global callback lists (previously
  each upsert/delete cycle leaked one hook).
- Adaptive-router init is deferred out of _add_deployment, so
  upsert_deployment now re-registers the adaptive router for the updated
  deployment explicitly — without this an upserted adaptive deployment
  would stop routing until the next full set_model_list().
- Tests: registry helper exercised directly (also satisfies
  router_code_coverage), colliding-name delete guard, adaptive delete
  retires hook, adaptive upsert applies new config without leaking hooks.
This commit is contained in:
Mihidum Hettiyahandi 2026-07-15 09:29:43 +10:00 • committed by Devin AI
parent 4218ce0e19
commit 2f1393365f
2 changed files with 235 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

72
litellm/router.py
View file

@ -7711,7 +7711,7 @@ class Router:
len(config.available_models),
)
def _remove_deployment_from_strategy_registries(self, model_name: str) -> None:
def _remove_deployment_from_strategy_registries(self, model_name: str, litellm_model: Optional[str]) -> None:
"""
Drop a model_name's strategy-router registration (auto / complexity /
adaptive / quality) when its deployment is removed from the router.
@ -7724,16 +7724,39 @@ class Router:
`/v1/models` & `/model/info`, and the stale registry object keeps
serving requests with its old config. See issue #33168.
Safe for regular deployments: they are never in these registries, so
the pops are no-ops. Strategy deployments are unique per model_name
(the init functions enforce this), so popping by name cannot orphan a
sibling deployment.
`litellm_model` is the removed deployment's `litellm_params.model`;
registries are only touched when it carries the `auto_router/` prefix.
This guards against a REGULAR deployment that merely shares a strategy
router's public model_name evicting the strategy router when deleted.
For adaptive routers, the deployment's `AdaptiveRouterPostCallHook` is
also removed from the global callback lists — otherwise every
upsert/delete cycle leaks one hook (and a deleted router's hook keeps
firing forever).
"""
if not (litellm_model or "").startswith("auto_router/"):
return
self.auto_routers.pop(model_name, None)
self.complexity_routers.pop(model_name, None)
self.adaptive_routers.pop(model_name, None)
self.quality_routers.pop(model_name, None)
adaptive_router = self.adaptive_routers.pop(model_name, None)
if adaptive_router is not None:
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
)
for _cb_list in (
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
for _cb in list(_cb_list):
if isinstance(_cb, AdaptiveRouterPostCallHook) and _cb.adaptive_router is adaptive_router:
litellm.logging_callback_manager.remove_callback_from_all_lists(_cb)
def _is_quality_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
"""
Check if the deployment is a quality-router deployment.
@ -8272,10 +8295,24 @@ class Router:
# of dying on the init functions' "already exists"
# check (which silently delists the model under
# ignore_invalid_deployments). See issue #33168.
self._remove_deployment_from_strategy_registries(model_name=_deployment_on_router.model_name)
self._remove_deployment_from_strategy_registries(
model_name=_deployment_on_router.model_name,
litellm_model=_deployment_on_router.litellm_params.model,
)
# if the model_id is not in router
self.add_deployment(deployment=deployment)
added = self.add_deployment(deployment=deployment)
# adaptive-router init is deferred out of _add_deployment (it needs
# the finalized model_list), so an upserted adaptive deployment must
# be re-registered here — its previous registration (and post-call
# hook) were removed above, and without this the deployment would
# stop routing until the next full set_model_list().
if (
added is not None
and self._is_adaptive_router_deployment(litellm_params=deployment.litellm_params)
and deployment.model_name not in self.adaptive_routers
):
self.init_adaptive_router_deployment(deployment=deployment)
return deployment
except Exception as e:
if self.ignore_invalid_deployments:
@ -8309,11 +8346,22 @@ class Router:
# de-register any strategy router tied to this deployment —
# otherwise it keeps serving the deleted model as a ghost and
# blocks any future re-add under the same name (issue #33168)
_removed_model_name = (
item.get("model_name") if isinstance(item, dict) else getattr(item, "model_name", None)
)
if isinstance(item, dict):
_removed_model_name = item.get("model_name")
_removed_lp = item.get("litellm_params") or {}
_removed_litellm_model = (
_removed_lp.get("model")
if isinstance(_removed_lp, dict)
else getattr(_removed_lp, "model", None)
)
else:
_removed_model_name = getattr(item, "model_name", None)
_removed_litellm_model = getattr(getattr(item, "litellm_params", None), "model", None)
if _removed_model_name:
self._remove_deployment_from_strategy_registries(model_name=_removed_model_name)
self._remove_deployment_from_strategy_registries(
model_name=_removed_model_name,
litellm_model=_removed_litellm_model,
)
_budget_limiter = self._get_router_deployment_budget_limiter()
if _budget_limiter is not None:
_budget_limiter.unregister_deployment_budget(model_id=id)

175
tests/test_litellm/test_router.py
View file

@ -5381,3 +5381,178 @@ class TestStrategyRegistryCleanupOnRemoval:
assert deleted is not None
assert "gpt-4o-mini" not in router.get_model_names()
class TestStrategyRegistryCleanupGuards:
"""Guards on _remove_deployment_from_strategy_registries: only strategy
deployments (litellm_params.model starting with auto_router/) may evict a
registry entry, and adaptive-router removal must also retire the
deployment's AdaptiveRouterPostCallHook. Issue #33168 follow-ups."""
def test_direct_call_is_noop_for_non_strategy_model(self):
router = litellm.Router(
model_list=[
{
"model_name": "backing-model",
"litellm_params": {"model": "gpt-4o-mini", "api_key": "fake-key"},
}
]
)
from litellm.types.router import Deployment, LiteLLM_Params
router.add_deployment(
Deployment(
model_name="shared-name",
litellm_params=LiteLLM_Params(
model="auto_router/complexity_router",
complexity_router_config={
"tiers": {
"SIMPLE": "backing-model",
"MEDIUM": "backing-model",
"COMPLEX": "backing-model",
"REASONING": "backing-model",
}
},
complexity_router_default_model="backing-model",
),
model_info={"id": "strategy-row-id"},
)
)
assert "shared-name" in router.complexity_routers
# a NON-strategy litellm_model must not evict the registry entry
router._remove_deployment_from_strategy_registries(
model_name="shared-name", litellm_model="azure/gpt-4o"
)
assert "shared-name" in router.complexity_routers
# a strategy litellm_model does
router._remove_deployment_from_strategy_registries(
model_name="shared-name", litellm_model="auto_router/complexity_router"
)
assert "shared-name" not in router.complexity_routers
def test_deleting_regular_deployment_with_colliding_name_keeps_strategy_router(self):
"""A regular deployment that merely SHARES a strategy router's public
model_name must not evict the strategy router when deleted."""
from litellm.types.router import Deployment, LiteLLM_Params
router = litellm.Router(
model_list=[
{
"model_name": "backing-model",
"litellm_params": {"model": "gpt-4o-mini", "api_key": "fake-key"},
}
]
)
router.add_deployment(
Deployment(
model_name="collide-auto-latest",
litellm_params=LiteLLM_Params(
model="auto_router/complexity_router",
complexity_router_config={
"tiers": {
"SIMPLE": "backing-model",
"MEDIUM": "backing-model",
"COMPLEX": "backing-model",
"REASONING": "backing-model",
}
},
complexity_router_default_model="backing-model",
),
model_info={"id": "strategy-row"},
)
)
# regular deployment under the SAME public name (different tenant/id)
router.add_deployment(
Deployment(
model_name="collide-auto-latest",
litellm_params=LiteLLM_Params(model="gpt-4o-mini", api_key="fake-key"),
model_info={"id": "regular-row"},
)
)
router.delete_deployment(id="regular-row")
assert "collide-auto-latest" in router.complexity_routers, (
"deleting a regular deployment with a colliding name must not evict "
"the strategy router registration"
)
@staticmethod
def _count_adaptive_hooks():
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
)
seen = set()
for cb_list in (
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
for cb in cb_list:
if isinstance(cb, AdaptiveRouterPostCallHook):
seen.add(id(cb))
return len(seen)
def _adaptive_router(self):
# adaptive-router init is deferred to set_model_list, so the deployment
# must be part of the Router's initial model_list
return litellm.Router(
model_list=[
{
"model_name": "fast",
"litellm_params": {"model": "gpt-4o-mini", "api_key": "fake-key"},
},
{
"model_name": "adaptive-test-router",
"litellm_params": {
"model": "auto_router/adaptive_router",
"adaptive_router_config": {"available_models": ["fast"]},
},
"model_info": {"id": "adaptive-row"},
},
]
)
def test_adaptive_router_delete_retires_post_call_hook(self):
router = self._adaptive_router()
assert "adaptive-test-router" in router.adaptive_routers
assert self._count_adaptive_hooks() == 1
router.delete_deployment(id="adaptive-row")
assert "adaptive-test-router" not in router.adaptive_routers
assert self._count_adaptive_hooks() == 0, "deleted adaptive router's hook must be retired"
def test_adaptive_router_upsert_applies_new_config_without_leaking_hooks(self):
from litellm.types.router import Deployment, LiteLLM_Params
router = self._adaptive_router()
original = router.adaptive_routers["adaptive-test-router"]
assert self._count_adaptive_hooks() == 1
router.upsert_deployment(
Deployment(
model_name="adaptive-test-router",
litellm_params=LiteLLM_Params(
model="auto_router/adaptive_router",
adaptive_router_config={"available_models": ["fast"], "explore_rate": 0.5},
),
model_info={"id": "adaptive-row"},
)
)
assert "adaptive-test-router" in router.get_model_names()
assert "adaptive-test-router" in router.adaptive_routers, (
"upserted adaptive deployment must still have a registered router"
)
assert router.adaptive_routers["adaptive-test-router"] is not original, (
"upsert must rebuild the adaptive router with the new config"
)
assert self._count_adaptive_hooks() == 1, (
"exactly one hook must remain after upsert (old retired, new registered)"
)