From 2f1393365fdaf547f758789a87451d289329cdf4 Mon Sep 17 00:00:00 2001 From: Mihidum Hettiyahandi <55163074+mihidumh@users.noreply.github.com> Date: Wed, 15 Jul 2026 09:29:43 +1000 Subject: [PATCH] fix(router): guard registry cleanup to strategy deployments; retire adaptive hooks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- litellm/router.py | 72 ++++++++++-- tests/test_litellm/test_router.py | 175 ++++++++++++++++++++++++++++++ 2 files changed, 235 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 855d6c38098..fc3f6deaa67 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cd2fdb9aab0..2380a8f5b20 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -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)" + )