diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 7979b7d09d..7642ad74b2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1264,5 +1264,5 @@ model LiteLLM_AdaptiveRouterSession { last_activity_at DateTime @default(now()) @updatedAt @@id([session_id, router_name, model_name]) - @@index([last_activity_at]) + @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1ec6f3a2ad..cddc173949 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -952,11 +952,16 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 _run_background_health_check() ) # start the background health check coroutine. - # Start adaptive-router queue flusher and load persisted state if any AdaptiveRouter is configured. + # Start adaptive-router queue flusher unconditionally — adaptive routers + # may be added later via `/config/reload`, and the flusher is a no-op when + # `llm_router.adaptive_routers` is empty. Per-router DB state is loaded + # lazily by the flusher on first tick (see `_state_loaded` flag) so + # hot-reloaded routers also get their persisted priors. if llm_router is not None and getattr(llm_router, "adaptive_routers", None): for _ar in llm_router.adaptive_routers.values(): await _ar.load_state_from_db(prisma_client) - asyncio.create_task(_adaptive_router_flusher_loop()) + _ar._state_loaded = True + asyncio.create_task(_adaptive_router_flusher_loop()) ## [Optional] Initialize dd tracer ProxyStartupEvent._init_dd_tracer() @@ -2450,6 +2455,13 @@ async def _adaptive_router_flusher_loop(): if not adaptive_routers or prisma_client is None: continue for ar in adaptive_routers.values(): + # Lazy state load: covers adaptive routers registered via + # `/config/reload` after proxy boot. + if not getattr(ar, "_state_loaded", False): + try: + await ar.load_state_from_db(prisma_client) + finally: + ar._state_loaded = True await ar.queue.flush_state_to_db(prisma_client) await ar.queue.flush_session_to_db(prisma_client) except asyncio.CancelledError: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 7979b7d09d..7642ad74b2 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1264,5 +1264,5 @@ model LiteLLM_AdaptiveRouterSession { last_activity_at DateTime @default(now()) @updatedAt @@id([session_id, router_name, model_name]) - @@index([last_activity_at]) + @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } diff --git a/litellm/router.py b/litellm/router.py index 6c7e73e680..07053db7d3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6943,6 +6943,26 @@ class Router: """Locate every adaptive-router deployment in the finalized model_list and build an AdaptiveRouter for each. Safe no-op when none are configured. Idempotent: skips any deployment whose model_name is already initialized.""" + # Drop any adaptive-router hooks left over from a previous Router + # instance (e.g. after `/config/reload` replaced `llm_router`). Without + # this, stale AdaptiveRouterPostCallHook callbacks from the old Router + # remain wired up in `litellm.callbacks` and double-fire signal + # recording for every request. + 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, + ): + litellm.logging_callback_manager.remove_callbacks_by_type( + _cb_list, AdaptiveRouterPostCallHook + ) + for entry in self.model_list or []: lp = ( entry.get("litellm_params") @@ -7052,6 +7072,7 @@ class Router: deployment.model_name, len(config.available_models), ) + def _is_quality_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """ Check if the deployment is a quality-router deployment. diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index d6ffd61b7b..8ab4a72d51 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -92,6 +92,9 @@ class AdaptiveRouter: # Evicted opportunistically in `get_or_create_session_state`. self._session_states_expiry: Dict[Tuple[str, str], float] = {} self._skipped_updates_total: int = 0 + # Set to True once the proxy flusher has loaded persisted priors from + # Postgres. Checked to support lazy-load on hot-reloaded routers. + self._state_loaded: bool = False self._lock = asyncio.Lock() self._init_cold_start_cells() diff --git a/schema.prisma b/schema.prisma index 7979b7d09d..7642ad74b2 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1264,5 +1264,5 @@ model LiteLLM_AdaptiveRouterSession { last_activity_at DateTime @default(now()) @updatedAt @@id([session_id, router_name, model_name]) - @@index([last_activity_at]) + @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py b/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py index 73cb66616e..604155e122 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py @@ -418,6 +418,59 @@ def test_finalize_adaptive_router_if_configured_initializes_and_is_idempotent(): assert r.adaptive_routers["my-router"] is original +def test_finalize_prunes_stale_adaptive_router_hooks_from_callbacks(): + """Replacing the Router (hot-reload path) must not leave stale + AdaptiveRouterPostCallHook instances in `litellm.callbacks` — otherwise + every request double-fires signal recording.""" + import litellm + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + model_list = [ + { + "model_name": "fast", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + }, + { + "model_name": "my-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["fast"]}, + }, + }, + ] + + # Snapshot any pre-existing AdaptiveRouterPostCallHook entries so we can + # restore them — other tests may have registered hooks we shouldn't drop. + pre_hooks = [ + cb for cb in litellm.callbacks if isinstance(cb, AdaptiveRouterPostCallHook) + ] + for cb in pre_hooks: + litellm.callbacks.remove(cb) + + try: + Router(model_list=model_list) + Router(model_list=model_list) # simulate hot-reload + + adaptive_hooks = [ + cb + for cb in litellm.callbacks + if isinstance(cb, AdaptiveRouterPostCallHook) + ] + assert len(adaptive_hooks) == 1, ( + f"expected exactly one AdaptiveRouterPostCallHook after hot-reload, " + f"got {len(adaptive_hooks)}" + ) + finally: + # Best-effort cleanup: remove whatever this test added, then restore. + for cb in list(litellm.callbacks): + if isinstance(cb, AdaptiveRouterPostCallHook): + litellm.callbacks.remove(cb) + for cb in pre_hooks: + litellm.callbacks.append(cb) + + def test_finalize_adaptive_router_if_configured_noop_when_none_configured(): """With no adaptive deployments in model_list, the finalizer leaves `adaptive_routers` empty."""