Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/health_check_utils/shared_health_check_manager.py: 0%
147 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import asyncio
2import json
3import time
4from collections.abc import Mapping, Sequence
5from typing import TYPE_CHECKING, Any, Final
7from litellm._logging import verbose_proxy_logger
8from litellm.caching.redis_cache import RedisCache
9from litellm.constants import (
10 DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
11 DEFAULT_SHARED_HEALTH_CHECK_TTL,
12)
13from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
14from litellm.proxy.health_check import perform_health_check
16if TYPE_CHECKING:
17 from litellm.router import Router
20class SharedHealthCheckManager:
21 """
22 Manager for coordinating health checks across multiple pods using Redis.
24 This class implements a shared health check state mechanism that:
25 - Prevents duplicate health checks across pods
26 - Caches health check results with configurable TTL
27 - Uses Redis locks to ensure only one pod runs health checks at a time
28 - Allows other pods to read cached results instead of running redundant checks
29 """
31 def __init__(
32 self,
33 redis_cache: RedisCache | None = None,
34 health_check_ttl: int = DEFAULT_SHARED_HEALTH_CHECK_TTL,
35 lock_ttl: int = DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
36 ):
37 self.redis_cache = redis_cache
38 self.health_check_ttl = health_check_ttl
39 self.lock_ttl = lock_ttl
40 self.pod_id = f"pod_{int(time.time() * 1000)}"
42 @staticmethod
43 def get_health_check_lock_key() -> str:
44 """Get the Redis key for health check lock."""
45 return "health_check_lock"
47 @staticmethod
48 def get_health_check_cache_key() -> str:
49 """Get the Redis key for health check results cache."""
50 return "health_check_results"
52 @staticmethod
53 def get_model_health_check_lock_key(model_name: str) -> str:
54 """Get the Redis key for model-specific health check lock."""
55 return f"health_check_lock:{model_name}"
57 @staticmethod
58 def get_model_health_check_cache_key(model_name: str) -> str:
59 """Get the Redis key for model-specific health check results cache."""
60 return f"health_check_results:{model_name}"
62 async def acquire_health_check_lock(self) -> bool:
63 """
64 Attempt to acquire the global health check lock.
66 Returns:
67 bool: True if lock was acquired, False otherwise
68 """
69 if self.redis_cache is None:
70 verbose_proxy_logger.debug("redis_cache is None, skipping lock acquisition")
71 return False
73 try:
74 lock_key: Final = self.get_health_check_lock_key()
75 acquired: Final = await self.redis_cache.async_set_cache(
76 lock_key,
77 self.pod_id,
78 nx=True, # Only set if key doesn't exist
79 ttl=self.lock_ttl,
80 )
82 if acquired:
83 verbose_proxy_logger.info("Pod %s acquired health check lock", self.pod_id)
84 else:
85 verbose_proxy_logger.debug("Pod %s failed to acquire health check lock", self.pod_id)
87 return bool(acquired)
88 except Exception as e:
89 verbose_proxy_logger.error("Error acquiring health check lock: %s", str(e))
90 return False
92 async def release_health_check_lock(self) -> None:
93 """Release the global health check lock."""
94 if self.redis_cache is None:
95 return
97 try:
98 lock_key: Final = self.get_health_check_lock_key()
99 # Only release if we own the lock
100 current_owner: Final = await self.redis_cache.async_get_cache(lock_key)
101 if current_owner == self.pod_id:
102 await self.redis_cache.async_delete_cache(lock_key)
103 verbose_proxy_logger.info("Pod %s released health check lock", self.pod_id)
104 except Exception as e:
105 verbose_proxy_logger.error("Error releasing health check lock: %s", str(e))
107 async def get_cached_health_check_results(self) -> dict[str, Any] | None:
108 """
109 Get cached health check results from Redis.
111 Returns:
112 Optional[Dict]: Cached health check results or None if not found/expired
113 """
114 if self.redis_cache is None:
115 return None
117 try:
118 cache_key: Final = self.get_health_check_cache_key()
119 cached_data: Final = await self.redis_cache.async_get_cache(cache_key)
121 if cached_data is None:
122 return None
124 # Parse the cached data
125 if isinstance(cached_data, str):
126 cached_results = json.loads(cached_data)
127 else:
128 cached_results = cached_data
130 # Check if the cache is still valid
131 cache_timestamp: Final = cached_results.get("timestamp", 0)
132 current_time: Final = time.time()
134 if current_time - cache_timestamp > self.health_check_ttl:
135 verbose_proxy_logger.debug("Cached health check results expired")
136 return None
138 verbose_proxy_logger.debug("Using cached health check results")
139 return cached_results
141 except Exception as e:
142 verbose_proxy_logger.error("Error getting cached health check results: %s", str(e))
143 return None
145 async def cache_health_check_results(
146 self,
147 healthy_endpoints: Sequence[Mapping[str, object]],
148 unhealthy_endpoints: Sequence[Mapping[str, object]],
149 ) -> None:
150 """
151 Cache health check results in Redis.
153 Args:
154 healthy_endpoints: List of healthy endpoints
155 unhealthy_endpoints: List of unhealthy endpoints
156 """
157 if self.redis_cache is None:
158 return
160 try:
161 cache_data: Final = {
162 "healthy_endpoints": healthy_endpoints,
163 "unhealthy_endpoints": unhealthy_endpoints,
164 "healthy_count": len(healthy_endpoints),
165 "unhealthy_count": len(unhealthy_endpoints),
166 "timestamp": time.time(),
167 "checked_by": self.pod_id,
168 }
170 cache_key: Final = self.get_health_check_cache_key()
171 await self.redis_cache.async_set_cache(
172 cache_key,
173 safe_dumps(cache_data),
174 ttl=self.health_check_ttl,
175 )
177 verbose_proxy_logger.info(
178 "Cached health check results for %d healthy and %d unhealthy endpoints",
179 len(healthy_endpoints),
180 len(unhealthy_endpoints),
181 )
183 except Exception as e:
184 verbose_proxy_logger.error("Error caching health check results: %s", str(e))
186 async def perform_shared_health_check(
187 self,
188 model_list: list[dict[str, Any]],
189 details: bool = True,
190 max_concurrency: int | None = None,
191 health_check_skip_disabled_background_models: bool = False,
192 router: "Router | None" = None,
193 ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]:
194 """
195 Perform health check with shared state coordination.
197 This method:
198 1. First checks if there are recent cached results
199 2. If no recent cache, tries to acquire lock to run health check
200 3. If lock acquired, runs health check and caches results
201 4. If lock not acquired, waits briefly and tries to get cached results again
202 5. Falls back to running health check locally if no cache available
204 Args:
205 model_list: List of models to check
206 details: Whether to include detailed information
207 max_concurrency: Optional limit on concurrent health check requests
208 health_check_skip_disabled_background_models: Remove models with disable_background_health_check: true
210 Returns:
211 Tuple of (healthy_endpoints, unhealthy_endpoints)
212 """
213 # First, try to get cached results
214 cached_results = await self.get_cached_health_check_results()
215 if cached_results is not None:
216 return (
217 cached_results.get("healthy_endpoints", []),
218 cached_results.get("unhealthy_endpoints", []),
219 {},
220 )
222 # No recent cache, try to acquire lock
223 lock_acquired: Final = await self.acquire_health_check_lock()
225 if lock_acquired:
226 try:
227 # We have the lock, run health check
228 verbose_proxy_logger.info(
229 "Pod %s running health check for %d models",
230 self.pod_id,
231 len(model_list),
232 )
234 (
235 healthy_endpoints,
236 unhealthy_endpoints,
237 exceptions_by_model_id,
238 ) = await perform_health_check(
239 model_list=model_list,
240 details=details,
241 max_concurrency=max_concurrency,
242 health_check_skip_disabled_background_models=health_check_skip_disabled_background_models,
243 router=router,
244 )
246 # Cache the results
247 await self.cache_health_check_results(healthy_endpoints, unhealthy_endpoints)
249 return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id
251 finally:
252 # Always release the lock
253 await self.release_health_check_lock()
254 else:
255 # If Redis is not configured, skip polling — there is no cache
256 # to wait for.
257 if self.redis_cache is None:
258 return await perform_health_check(
259 model_list=model_list,
260 details=details,
261 max_concurrency=max_concurrency,
262 health_check_skip_disabled_background_models=health_check_skip_disabled_background_models,
263 router=router,
264 )
266 # Lock not acquired — poll for cached results until the lock
267 # holder finishes or the lock expires, rather than falling back
268 # to a redundant local health check after only 2 seconds.
269 verbose_proxy_logger.debug("Pod %s waiting for other pod to complete health check", self.pod_id)
271 poll_interval: Final = 5 # seconds between cache checks
272 max_wait: Final = self.lock_ttl # wait at most as long as the lock can live
273 elapsed = 0
275 while elapsed < max_wait:
276 await asyncio.sleep(poll_interval)
277 elapsed += poll_interval
279 cached_results = await self.get_cached_health_check_results()
280 if cached_results is not None:
281 verbose_proxy_logger.info(
282 "Pod %s using cached health check results after waiting %ds",
283 self.pod_id,
284 elapsed,
285 )
286 return (
287 cached_results.get("healthy_endpoints", []),
288 cached_results.get("unhealthy_endpoints", []),
289 {},
290 )
292 # Check if the lock is still held — if it was released without
293 # caching (e.g. the holder crashed), stop waiting early.
294 try:
295 lock_key = self.get_health_check_lock_key()
296 current_owner = await self.redis_cache.async_get_cache(lock_key)
297 if current_owner is None:
298 verbose_proxy_logger.debug(
299 "Pod %s detected lock released without cache, stopping wait",
300 self.pod_id,
301 )
302 break
303 except Exception:
304 # Redis hiccup — continue polling rather than crashing out
305 pass
307 # Exhausted wait — fall back to local health check
308 verbose_proxy_logger.warning(
309 "Pod %s falling back to local health check after waiting %ds (no cache available)",
310 self.pod_id,
311 elapsed,
312 )
314 return await perform_health_check(
315 model_list=model_list,
316 details=details,
317 max_concurrency=max_concurrency,
318 health_check_skip_disabled_background_models=health_check_skip_disabled_background_models,
319 router=router,
320 )
322 async def is_health_check_in_progress(self) -> bool:
323 """
324 Check if a health check is currently in progress by another pod.
326 Returns:
327 bool: True if health check is in progress, False otherwise
328 """
329 if self.redis_cache is None:
330 return False
332 try:
333 lock_key: Final = self.get_health_check_lock_key()
334 current_owner: Final = await self.redis_cache.async_get_cache(lock_key)
335 return current_owner is not None and current_owner != self.pod_id
336 except Exception as e:
337 verbose_proxy_logger.error("Error checking health check lock status: %s", str(e))
338 return False
340 async def get_health_check_status(self) -> dict[str, object]:
341 """
342 Get the current status of health check coordination.
344 Returns:
345 Dict containing status information
346 """
347 status: Final[dict[str, object]] = {
348 "pod_id": self.pod_id,
349 "redis_available": self.redis_cache is not None,
350 "lock_ttl": self.lock_ttl,
351 "cache_ttl": self.health_check_ttl,
352 }
354 if self.redis_cache is not None:
355 try:
356 # Check if there's a current lock
357 lock_key: Final = self.get_health_check_lock_key()
358 current_owner: Final = await self.redis_cache.async_get_cache(lock_key)
359 status["lock_owner"] = current_owner
360 status["lock_in_progress"] = current_owner is not None
362 # Check cache status
363 cached_results: Final = await self.get_cached_health_check_results()
364 status["cache_available"] = cached_results is not None
365 if cached_results:
366 status["cache_age_seconds"] = time.time() - cached_results.get("timestamp", 0)
367 status["last_checked_by"] = cached_results.get("checked_by")
369 except Exception as e:
370 status["error"] = str(e)
372 return status