Merge 7199af0449 into c2c2a623c0

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"]

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(

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
 # ============================================================================