diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c132b12ac6..dc22abf2cf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -835,7 +835,7 @@ use_background_health_checks = None use_queue = False health_check_interval = None health_check_details = None -health_check_results = {} +health_check_results: dict[str, Union[int, List[Dict[str, Any]]]] = {} queue: List = [] litellm_proxy_budget_name = "litellm-proxy-budget" litellm_proxy_admin_name = LITELLM_PROXY_ADMIN_NAME @@ -1231,13 +1231,15 @@ async def _run_background_health_check(): """ global health_check_results, llm_model_list, health_check_interval, health_check_details - # make 1 deep copy of llm_model_list -> use this for all background health checks - _llm_model_list = copy.deepcopy(llm_model_list) - - if _llm_model_list is None: + if health_check_interval is None or not isinstance( + health_check_interval, int + ) or health_check_interval <= 0: return while True: + # make 1 deep copy of llm_model_list on every health check iteration + _llm_model_list = copy.deepcopy(llm_model_list) or [] + healthy_endpoints, unhealthy_endpoints = await perform_health_check( model_list=_llm_model_list, details=health_check_details ) @@ -1248,10 +1250,7 @@ async def _run_background_health_check(): health_check_results["healthy_count"] = len(healthy_endpoints) health_check_results["unhealthy_count"] = len(unhealthy_endpoints) - if health_check_interval is not None and isinstance( - health_check_interval, float - ): - await asyncio.sleep(health_check_interval) + await asyncio.sleep(health_check_interval) class StreamingCallbackError(Exception): diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index dda39d2bd5..4b7d86b8ca 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2194,6 +2194,50 @@ async def test_get_ui_settings_spend_logs_threshold(): proxy_state.set_proxy_state_variable("spend_logs_row_count", 0) +@pytest.mark.asyncio +async def test_run_background_health_check_reflects_llm_model_list(monkeypatch): + """ + Test that _run_background_health_check reflects changes to llm_model_list in each health check iteration. + """ + import litellm.proxy.proxy_server as proxy_server + import copy + + test_model_list_1 = [{"model_name": "model-a"}] + test_model_list_2 = [{"model_name": "model-b"}] + called_model_lists = [] + + async def fake_perform_health_check(model_list, details): + called_model_lists.append(copy.deepcopy(model_list)) + return (["healthy"], ["unhealthy"]) + + monkeypatch.setattr(proxy_server, "health_check_interval", 1) + monkeypatch.setattr(proxy_server, "health_check_details", None) + monkeypatch.setattr(proxy_server, "llm_model_list", copy.deepcopy(test_model_list_1)) + monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check) + monkeypatch.setattr(proxy_server, "health_check_results", {}) + + async def fake_sleep(interval): + raise asyncio.CancelledError() + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + try: + await proxy_server._run_background_health_check() + except asyncio.CancelledError: + pass + + monkeypatch.setattr(proxy_server, "llm_model_list", copy.deepcopy(test_model_list_2)) + + try: + await proxy_server._run_background_health_check() + except asyncio.CancelledError: + pass + + assert len(called_model_lists) >= 2 + assert called_model_lists[0] == test_model_list_1 + assert called_model_lists[1] == test_model_list_2 + + def test_get_timeout_from_request(): from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup