diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..310b32e2c8a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1342,8 +1342,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await ProxyStartupEvent._sync_ui_settings_to_general_settings() # Start background health checks AFTER models are loaded and index is built - if use_background_health_checks: - asyncio.create_task(_run_background_health_check()) # start the background health check coroutine. + await _reconcile_background_health_check_task() # Start adaptive-router queue flusher unconditionally — adaptive routers # may be added later via `/config/reload`, and the flusher is a no-op when @@ -2376,6 +2375,7 @@ health_check_concurrency = None health_check_details = None health_check_results: dict[str, int | list[dict[str, Any]]] = {} background_health_check_loop_active = False +background_health_check_task: asyncio.Task[None] | None = None background_health_check_cycle_seq = 0 queue: Final[list] = [] litellm_proxy_budget_name: Final = LITELLM_PROXY_BUDGET_NAME @@ -4233,6 +4233,27 @@ async def _run_background_health_check(): await asyncio.sleep(health_check_interval) +async def _reconcile_background_health_check_task() -> None: + global background_health_check_task, background_health_check_loop_active # noqa: PLW0603 # reload reconciliation updates module task state + + if use_background_health_checks: + if background_health_check_task is None or background_health_check_task.done(): + background_health_check_task = asyncio.create_task(_run_background_health_check()) + return + + if background_health_check_task is not None: + if not background_health_check_task.done(): + background_health_check_task.cancel() + try: + await background_health_check_task + except asyncio.CancelledError: + pass + elif not background_health_check_task.cancelled(): + background_health_check_task.exception() + background_health_check_task = None + background_health_check_loop_active = False + + class StreamingCallbackError(Exception): pass @@ -7044,14 +7065,50 @@ class ProxyConfig: except ValueError: verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") - async def _update_general_settings(self, db_general_settings: Json | None): + async def _update_general_settings(self, db_general_settings: Json | None): # noqa: C901 # DB reload branches combine independent settings """ Pull from DB, read general settings value """ global general_settings, store_model_in_db + global use_background_health_checks, use_shared_health_check # noqa: PLW0603 # DB reload runtime state + global health_check_interval, health_check_concurrency, health_check_details # noqa: PLW0603 # DB reload runtime state if db_general_settings is None: return _general_settings: Final = dict(db_general_settings) + health_check_settings: Final = frozenset("background_health_checks use_shared_health_check health_check_interval health_check_concurrency health_check_details enable_health_check_routing health_check_staleness_threshold health_check_ignore_transient_errors".split()) # noqa: E501 # immutable DB setting key list # fmt: skip + + def _db_setting_is_active(key: str) -> bool: + return key in _general_settings and key not in self._yaml_general_settings_keys + + for key in health_check_settings: + if key not in _general_settings or key in self._yaml_general_settings_keys: + continue + general_settings[key] = _general_settings[key] + + if _db_setting_is_active("background_health_checks"): + use_background_health_checks = _general_settings["background_health_checks"] + if _db_setting_is_active("use_shared_health_check"): + use_shared_health_check = _general_settings["use_shared_health_check"] + if _db_setting_is_active("health_check_interval"): + health_check_interval = _general_settings["health_check_interval"] + if _db_setting_is_active("health_check_concurrency"): + health_check_concurrency = _general_settings["health_check_concurrency"] + if _db_setting_is_active("health_check_details"): + health_check_details = _general_settings["health_check_details"] + + if llm_router is not None: + if _db_setting_is_active("enable_health_check_routing"): + llm_router.enable_health_check_routing = _general_settings["enable_health_check_routing"] + if _db_setting_is_active("health_check_staleness_threshold"): + llm_router.health_state_cache.staleness_threshold = float( + _general_settings["health_check_staleness_threshold"] + ) + if _db_setting_is_active("health_check_ignore_transient_errors"): + llm_router.health_check_ignore_transient_errors = _general_settings[ + "health_check_ignore_transient_errors" + ] + await _reconcile_background_health_check_task() + ## MAX PARALLEL REQUESTS ## if "max_parallel_requests" in _general_settings: general_settings["max_parallel_requests"] = _general_settings["max_parallel_requests"] diff --git a/tests/test_litellm/proxy/proxy_server/test_background_health.py b/tests/test_litellm/proxy/proxy_server/test_background_health.py index d99aa637955..7bce5aef179 100644 --- a/tests/test_litellm/proxy/proxy_server/test_background_health.py +++ b/tests/test_litellm/proxy/proxy_server/test_background_health.py @@ -25,6 +25,7 @@ from litellm.proxy.proxy_server import ( _adaptive_router_flusher_loop, _get_endpoint_exception_status, _get_process_rss_mb, + _reconcile_background_health_check_task, _run_background_health_check, _run_direct_health_check_with_instrumentation, _rss_mb_for_log, @@ -615,6 +616,85 @@ async def test_adaptive_router_flusher_loop_times_out_when_sleep_real(monkeypatc # _run_background_health_check # --------------------------------------------------------------------------- +@pytest.mark.asyncio +async def test_reconcile_background_health_check_task_starts_enabled_loop(monkeypatch): + fake_task = MagicMock() + + def _create_task(coroutine): + coroutine.close() + return fake_task + + monkeypatch.setattr(proxy_server, "use_background_health_checks", True) + monkeypatch.setattr(proxy_server, "background_health_check_task", None) + monkeypatch.setattr(proxy_server.asyncio, "create_task", _create_task) + + await _reconcile_background_health_check_task() + + assert proxy_server.background_health_check_task is fake_task + + +@pytest.mark.asyncio +async def test_reconcile_background_health_check_task_stops_disabled_loop(monkeypatch): + fake_task = AsyncMock() + + monkeypatch.setattr(proxy_server, "use_background_health_checks", False) + monkeypatch.setattr(proxy_server, "background_health_check_task", fake_task) + monkeypatch.setattr(proxy_server, "background_health_check_loop_active", True) + + await _reconcile_background_health_check_task() + + fake_task.cancel.assert_called_once_with() + fake_task.assert_awaited_once_with() + assert proxy_server.background_health_check_task is None + assert proxy_server.background_health_check_loop_active is False + + +@pytest.mark.asyncio +async def test_reconcile_background_health_check_task_restarts_completed_loop(monkeypatch): + old_task = MagicMock() + old_task.done.return_value = True + new_task = MagicMock() + + def _create_task(coroutine): + coroutine.close() + return new_task + + monkeypatch.setattr(proxy_server, "use_background_health_checks", True) + monkeypatch.setattr(proxy_server, "background_health_check_task", old_task) + monkeypatch.setattr(proxy_server.asyncio, "create_task", _create_task) + + await _reconcile_background_health_check_task() + + assert proxy_server.background_health_check_task is new_task + + +@pytest.mark.asyncio +async def test_reconcile_background_health_check_task_handles_cancellation(monkeypatch): + task = AsyncMock(side_effect=asyncio.CancelledError) + + monkeypatch.setattr(proxy_server, "use_background_health_checks", False) + monkeypatch.setattr(proxy_server, "background_health_check_task", task) + + await _reconcile_background_health_check_task() + + task.cancel.assert_called_once_with() + task.assert_awaited_once_with() + assert proxy_server.background_health_check_task is None + + +@pytest.mark.asyncio +async def test_reconcile_background_health_check_task_collects_completed_error(monkeypatch): + task = MagicMock() + task.done.return_value = True + task.cancelled.return_value = False + + monkeypatch.setattr(proxy_server, "use_background_health_checks", False) + monkeypatch.setattr(proxy_server, "background_health_check_task", task) + + await _reconcile_background_health_check_task() + + task.exception.assert_called_once_with() + assert proxy_server.background_health_check_task is None @pytest.mark.asyncio async def test_run_background_health_check_returns_immediately_when_interval_invalid( diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 03a24ec5e98..72945048a0a 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7121,6 +7121,111 @@ async def test_async_data_generator_cleanup_on_midstream_error(): # ============================================================================ +@pytest.mark.asyncio +async def test_update_general_settings_propagates_health_check_settings(monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + monkeypatch.setattr(proxy_config, "_yaml_general_settings_keys", set()) + monkeypatch.setattr(ps, "general_settings", {}) + monkeypatch.setattr(ps, "use_background_health_checks", False) + monkeypatch.setattr(ps, "use_shared_health_check", False) + monkeypatch.setattr(ps, "health_check_interval", None) + monkeypatch.setattr(ps, "health_check_concurrency", None) + monkeypatch.setattr(ps, "health_check_details", True) + monkeypatch.setattr(ps, "llm_router", None) + monkeypatch.setattr(ps, "_reconcile_background_health_check_task", AsyncMock()) + + await proxy_config._update_general_settings(None) + await proxy_config._update_general_settings( + db_general_settings={ + "background_health_checks": True, + "use_shared_health_check": True, + "health_check_interval": 300, + "health_check_concurrency": 4, + "health_check_details": False, + "enable_health_check_routing": True, + } + ) + + assert ps.use_background_health_checks is True + assert ps.use_shared_health_check is True + assert ps.health_check_interval == 300 + assert ps.health_check_concurrency == 4 + assert ps.health_check_details is False + assert ps.general_settings["enable_health_check_routing"] is True + + +@pytest.mark.asyncio +async def test_update_general_settings_updates_health_state_cache_threshold(monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + monkeypatch.setattr(proxy_config, "_yaml_general_settings_keys", set()) + fake_router = types.SimpleNamespace( + health_state_cache=types.SimpleNamespace(staleness_threshold=600.0), + enable_health_check_routing=False, + health_check_ignore_transient_errors=False, + ) + monkeypatch.setattr(ps, "general_settings", {}) + monkeypatch.setattr(ps, "llm_router", fake_router) + monkeypatch.setattr(ps, "_reconcile_background_health_check_task", AsyncMock()) + + await proxy_config._update_general_settings(None) + await proxy_config._update_general_settings(db_general_settings={"health_check_staleness_threshold": 900}) + + assert fake_router.health_state_cache.staleness_threshold == 900.0 + + +@pytest.mark.asyncio +async def test_update_general_settings_updates_router_health_options(monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + monkeypatch.setattr(proxy_config, "_yaml_general_settings_keys", set()) + fake_router = types.SimpleNamespace( + health_state_cache=types.SimpleNamespace(staleness_threshold=600.0), + enable_health_check_routing=False, + health_check_ignore_transient_errors=False, + ) + monkeypatch.setattr(ps, "general_settings", {}) + monkeypatch.setattr(ps, "llm_router", fake_router) + monkeypatch.setattr(ps, "_reconcile_background_health_check_task", AsyncMock()) + + await proxy_config._update_general_settings(None) + await proxy_config._update_general_settings( + db_general_settings={ + "enable_health_check_routing": True, + "health_check_staleness_threshold": 900, + "health_check_ignore_transient_errors": True, + } + ) + + assert fake_router.enable_health_check_routing is True + assert fake_router.health_state_cache.staleness_threshold == 900.0 + assert fake_router.health_check_ignore_transient_errors is True + + +@pytest.mark.asyncio +async def test_reconcile_background_health_check_task_starts_enabled_task(monkeypatch): + from litellm.proxy import proxy_server as ps + + async def _noop_background_health_check(): + return None + + monkeypatch.setattr(ps, "use_background_health_checks", True) + monkeypatch.setattr(ps, "background_health_check_task", None) + monkeypatch.setattr(ps, "_run_background_health_check", _noop_background_health_check) + + await ps._reconcile_background_health_check_task() + + assert ps.background_health_check_task is not None + await ps.background_health_check_task + + # store_model_in_db DB Config Override Tests # ============================================================================