Fix/shared health check polling (#26434)

Squash-merged by litellm-agent from noahnistler's PR.

litellm/proxy/health_check_utils/shared_health_check_manager.py

@@ -253,27 +253,63 @@ class SharedHealthCheckManager:
                 # 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
 
-            # Still no cache, fall back to local health check
+                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", []),
+                        {},
+                    )
+
+                # 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(

tests/test_litellm/proxy/test_shared_health_check.py

@@ -322,13 +322,13 @@ class TestSharedHealthCheckManager:
     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 @@ class TestSharedHealthCheckManager:
                 )
             )
 
-        # 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 @@ class TestSharedHealthCheckManager:
                 )
             )
 
-        # 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
         )