diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index 5b8370fece8..ad58bc7e286 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -253,27 +253,63 @@ async def perform_shared_health_check( # Always release the lock await self.release_health_check_lock() else: - # Lock not acquired, wait briefly and try to get cached results + # If Redis is not configured, skip polling — there is no cache + # to wait for. + if self.redis_cache is None: + return await perform_health_check( + model_list=model_list, + details=details, + max_concurrency=max_concurrency, + ) + + # Lock not acquired — poll for cached results until the lock + # holder finishes or the lock expires, rather than falling back + # to a redundant local health check after only 2 seconds. verbose_proxy_logger.debug( "Pod %s waiting for other pod to complete health check", self.pod_id ) - # Wait a bit for the other pod to complete - await asyncio.sleep(2) + poll_interval = 5 # seconds between cache checks + max_wait = self.lock_ttl # wait at most as long as the lock can live + elapsed = 0 - # Try to get cached results again - cached_results = await self.get_cached_health_check_results() - if cached_results is not None: - return ( - cached_results.get("healthy_endpoints", []), - cached_results.get("unhealthy_endpoints", []), - {}, - ) + while elapsed < max_wait: + await asyncio.sleep(poll_interval) + elapsed += poll_interval + + cached_results = await self.get_cached_health_check_results() + if cached_results is not None: + verbose_proxy_logger.info( + "Pod %s using cached health check results after waiting %ds", + self.pod_id, + elapsed, + ) + return ( + cached_results.get("healthy_endpoints", []), + cached_results.get("unhealthy_endpoints", []), + {}, + ) - # Still no cache, fall back to local health check + # Check if the lock is still held — if it was released without + # caching (e.g. the holder crashed), stop waiting early. + try: + lock_key = self.get_health_check_lock_key() + current_owner = await self.redis_cache.async_get_cache(lock_key) + if current_owner is None: + verbose_proxy_logger.debug( + "Pod %s detected lock released without cache, stopping wait", + self.pod_id, + ) + break + except Exception: + # Redis hiccup — continue polling rather than crashing out + pass + + # Exhausted wait — fall back to local health check verbose_proxy_logger.warning( - "Pod %s falling back to local health check (no cache available)", + "Pod %s falling back to local health check after waiting %ds (no cache available)", self.pod_id, + elapsed, ) return await perform_health_check( diff --git a/tests/test_litellm/proxy/test_shared_health_check.py b/tests/test_litellm/proxy/test_shared_health_check.py index 04099f2634b..1530d336085 100644 --- a/tests/test_litellm/proxy/test_shared_health_check.py +++ b/tests/test_litellm/proxy/test_shared_health_check.py @@ -322,13 +322,13 @@ async def test_perform_shared_health_check_with_lock_acquisition( async def test_perform_shared_health_check_lock_failed_then_cache( self, shared_health_manager, mock_redis_cache ): - """Test performing shared health check when lock fails but cache becomes available""" + """Test performing shared health check when lock fails but cache becomes available during polling""" # First call: no cache, lock fails - # Second call: cache available + # Polling finds cache on first iteration mock_redis_cache.async_get_cache.side_effect = [ - None, # No cache initially + None, # No cache initially (get_cached_health_check_results) json.dumps( - { # Cache available after waiting + { # Cache available on first poll iteration "healthy_endpoints": [{"model": "cached-model"}], "unhealthy_endpoints": [], "healthy_count": 1, @@ -350,18 +350,68 @@ async def test_perform_shared_health_check_lock_failed_then_cache( ) ) - # Should wait and then get cached results - mock_sleep.assert_called_once_with(2) + # Should poll once (5s interval) and find cached results + mock_sleep.assert_called_once_with(5) assert healthy == [{"model": "cached-model"}] assert unhealthy == [] @pytest.mark.asyncio - async def test_perform_shared_health_check_fallback( + async def test_perform_shared_health_check_fallback(self, mock_redis_cache): + """Test performing shared health check with fallback to local health check""" + # Use short lock_ttl so the polling loop only runs 2 iterations + manager = SharedHealthCheckManager( + redis_cache=mock_redis_cache, + health_check_ttl=300, + lock_ttl=10, + ) + + # No cache ever, lock always held by another pod + mock_redis_cache.async_get_cache.side_effect = [ + None, # Initial cache check + None, # Iteration 1: cache check + "other_pod", # Iteration 1: lock check (still held) + None, # Iteration 2: cache check + "other_pod", # Iteration 2: lock check (still held) + ] + mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails + + model_list = [ + {"model_name": "test-model", "litellm_params": {"model": "test-model"}} + ] + expected_healthy = [{"model": "test-model", "status": "healthy"}] + expected_unhealthy = [] + + with ( + patch("asyncio.sleep") as mock_sleep, + patch( + "litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check" + ) as mock_perform, + ): + mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) + + healthy, unhealthy, _ = await manager.perform_shared_health_check( + model_list, details=True + ) + + # Should poll twice (5s * 2 = 10s >= lock_ttl) then fall back + assert mock_sleep.call_count == 2 + mock_sleep.assert_called_with(5) + mock_perform.assert_called_once_with( + model_list=model_list, details=True, max_concurrency=None + ) + assert healthy == expected_healthy + assert unhealthy == expected_unhealthy + + @pytest.mark.asyncio + async def test_perform_shared_health_check_early_exit_orphaned_lock( self, shared_health_manager, mock_redis_cache ): - """Test performing shared health check with fallback to local health check""" - # No cache, lock fails, no cache after waiting - mock_redis_cache.async_get_cache.return_value = None + """Test that polling exits early when the lock disappears without a cache write (crash recovery)""" + mock_redis_cache.async_get_cache.side_effect = [ + None, # Initial cache check + None, # Iteration 1: cache check (still no cache) + None, # Iteration 1: lock check -> lock gone (holder crashed) + ] mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails model_list = [ @@ -384,8 +434,77 @@ async def test_perform_shared_health_check_fallback( ) ) - # Should fall back to local health check - mock_sleep.assert_called_once_with(2) + # Should detect orphaned lock after 1 iteration and fall back immediately + mock_sleep.assert_called_once_with(5) + mock_perform.assert_called_once_with( + model_list=model_list, details=True, max_concurrency=None + ) + assert healthy == expected_healthy + assert unhealthy == expected_unhealthy + + @pytest.mark.asyncio + async def test_perform_shared_health_check_redis_error_during_polling( + self, shared_health_manager, mock_redis_cache + ): + """Test that a transient Redis error during lock polling doesn't crash the loop""" + cached_data = json.dumps( + { + "healthy_endpoints": [{"model": "cached-model"}], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + "timestamp": time.time() - 100, + } + ) + mock_redis_cache.async_get_cache.side_effect = [ + None, # Initial cache check + None, # Iteration 1: cache check + Exception("Redis connection lost"), # Iteration 1: lock check errors + cached_data, # Iteration 2: cache check -> found! + ] + mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails + + model_list = [ + {"model_name": "test-model", "litellm_params": {"model": "test-model"}} + ] + + with patch("asyncio.sleep") as mock_sleep: + healthy, unhealthy, _ = ( + await shared_health_manager.perform_shared_health_check( + model_list, details=True + ) + ) + + # Should survive the Redis error and find cache on iteration 2 + assert mock_sleep.call_count == 2 + assert healthy == [{"model": "cached-model"}] + assert unhealthy == [] + + @pytest.mark.asyncio + async def test_perform_shared_health_check_no_redis_skips_polling(self): + """Test that polling is skipped entirely when redis_cache is None""" + manager = SharedHealthCheckManager(redis_cache=None) + + model_list = [ + {"model_name": "test-model", "litellm_params": {"model": "test-model"}} + ] + expected_healthy = [{"model": "test-model", "status": "healthy"}] + expected_unhealthy = [] + + with ( + patch("asyncio.sleep") as mock_sleep, + patch( + "litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check" + ) as mock_perform, + ): + mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) + + healthy, unhealthy, _ = await manager.perform_shared_health_check( + model_list, details=True + ) + + # Should NOT sleep at all — falls back to local health check immediately + mock_sleep.assert_not_called() mock_perform.assert_called_once_with( model_list=model_list, details=True, max_concurrency=None )