mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
refactor(redis): extract shared health store
This commit is contained in:
Devin AI
parent
0b0fcd06de
commit
69a516138e
3 changed files with 527 additions and 673 deletions
|
|
@ -1,12 +1,9 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
|
@ -14,25 +11,23 @@ from litellm.constants import (
|
|||
DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
|
||||
DEFAULT_SHARED_HEALTH_CHECK_TTL,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy.health_check import perform_health_check
|
||||
from litellm.redis.stores.shared_health import (
|
||||
CachedHealthCheckResults,
|
||||
SharedHealthStore,
|
||||
TransportSharedHealthStore,
|
||||
lock_key,
|
||||
model_lock_key,
|
||||
model_results_key,
|
||||
results_key,
|
||||
)
|
||||
from litellm.redis.transport.contract import RedisFailure
|
||||
from litellm.router_utils.health_state_cache import HEALTH_CHECKS_TARGET
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _CachedHealthCheckResults(TypedDict):
|
||||
healthy_endpoints: NotRequired[ReadOnly[list[Mapping[str, object]]]]
|
||||
unhealthy_endpoints: NotRequired[ReadOnly[list[Mapping[str, object]]]]
|
||||
timestamp: NotRequired[ReadOnly[int | float]]
|
||||
checked_by: NotRequired[ReadOnly[object]]
|
||||
|
||||
|
||||
_CACHED_HEALTH_CHECK_RESULTS: Final = TypeAdapter(_CachedHealthCheckResults)
|
||||
|
||||
|
||||
class SharedHealthCheckManager:
|
||||
"""
|
||||
Manager for coordinating health checks across multiple pods using Redis.
|
||||
|
|
@ -51,6 +46,9 @@ class SharedHealthCheckManager:
|
|||
lock_ttl: int = DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
|
||||
):
|
||||
self.redis_cache = redis_cache
|
||||
self._store: SharedHealthStore | None = (
|
||||
TransportSharedHealthStore(redis_cache.transport()) if redis_cache is not None else None
|
||||
)
|
||||
self.health_check_ttl = health_check_ttl
|
||||
self.lock_ttl = lock_ttl
|
||||
self.pod_id = f"pod_{int(time.time() * 1000)}"
|
||||
|
|
@ -58,22 +56,22 @@ class SharedHealthCheckManager:
|
|||
@staticmethod
|
||||
def get_health_check_lock_key() -> str:
|
||||
"""Get the Redis key for health check lock."""
|
||||
return "health_check_lock"
|
||||
return lock_key()
|
||||
|
||||
@staticmethod
|
||||
def get_health_check_cache_key() -> str:
|
||||
"""Get the Redis key for health check results cache."""
|
||||
return "health_check_results"
|
||||
return results_key()
|
||||
|
||||
@staticmethod
|
||||
def get_model_health_check_lock_key(model_name: str) -> str:
|
||||
"""Get the Redis key for model-specific health check lock."""
|
||||
return f"health_check_lock:{model_name}"
|
||||
return model_lock_key(model_name)
|
||||
|
||||
@staticmethod
|
||||
def get_model_health_check_cache_key(model_name: str) -> str:
|
||||
"""Get the Redis key for model-specific health check results cache."""
|
||||
return f"health_check_results:{model_name}"
|
||||
return model_results_key(model_name)
|
||||
|
||||
@with_service_target(HEALTH_CHECKS_TARGET)
|
||||
async def acquire_health_check_lock(self) -> bool:
|
||||
|
|
@ -83,18 +81,14 @@ class SharedHealthCheckManager:
|
|||
Returns:
|
||||
bool: True if lock was acquired, False otherwise
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
if self._store is None:
|
||||
verbose_proxy_logger.debug("redis_cache is None, skipping lock acquisition")
|
||||
return False
|
||||
|
||||
try:
|
||||
lock_key: Final = self.get_health_check_lock_key()
|
||||
acquired: Final = await self.redis_cache.async_set_cache(
|
||||
lock_key,
|
||||
self.pod_id,
|
||||
nx=True, # Only set if key doesn't exist
|
||||
ttl=self.lock_ttl,
|
||||
)
|
||||
acquired: Final = await self._store.acquire(self.pod_id, ttl=timedelta(seconds=self.lock_ttl))
|
||||
if isinstance(acquired, RedisFailure):
|
||||
return False
|
||||
|
||||
if acquired:
|
||||
verbose_proxy_logger.info("Pod %s acquired health check lock", self.pod_id)
|
||||
|
|
@ -102,49 +96,38 @@ class SharedHealthCheckManager:
|
|||
verbose_proxy_logger.debug("Pod %s failed to acquire health check lock", self.pod_id)
|
||||
|
||||
return bool(acquired)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error acquiring health check lock: %s", str(e))
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@with_service_target(HEALTH_CHECKS_TARGET)
|
||||
async def release_health_check_lock(self) -> None:
|
||||
"""Release the global health check lock."""
|
||||
if self.redis_cache is None:
|
||||
if self._store is None:
|
||||
return
|
||||
|
||||
try:
|
||||
lock_key: Final = self.get_health_check_lock_key()
|
||||
# Only release if we own the lock
|
||||
current_owner: Final = await self.redis_cache.async_get_cache(lock_key)
|
||||
if current_owner == self.pod_id:
|
||||
await self.redis_cache.async_delete_cache(lock_key)
|
||||
released: Final = await self._store.release(self.pod_id)
|
||||
if released is True:
|
||||
verbose_proxy_logger.info("Pod %s released health check lock", self.pod_id)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error releasing health check lock: %s", str(e))
|
||||
except Exception:
|
||||
return
|
||||
|
||||
@with_service_target(HEALTH_CHECKS_TARGET)
|
||||
async def get_cached_health_check_results(self) -> _CachedHealthCheckResults | None:
|
||||
async def get_cached_health_check_results(self) -> CachedHealthCheckResults | None:
|
||||
"""
|
||||
Get cached health check results from Redis.
|
||||
|
||||
Returns:
|
||||
Optional[Dict]: Cached health check results or None if not found/expired
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
if self._store is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
cache_key: Final = self.get_health_check_cache_key()
|
||||
cached_data: Final = await self.redis_cache.async_get_cache(cache_key)
|
||||
|
||||
if cached_data is None:
|
||||
cached_results: Final = await self._store.get()
|
||||
if cached_results is None or isinstance(cached_results, RedisFailure):
|
||||
return None
|
||||
|
||||
# Parse the cached data
|
||||
cached_results: Final = _CACHED_HEALTH_CHECK_RESULTS.validate_python(
|
||||
json.loads(cached_data) if isinstance(cached_data, str) else cached_data
|
||||
)
|
||||
|
||||
# Check if the cache is still valid
|
||||
cache_timestamp: Final = cached_results.get("timestamp", 0)
|
||||
current_time: Final = time.time()
|
||||
|
|
@ -156,8 +139,7 @@ class SharedHealthCheckManager:
|
|||
verbose_proxy_logger.debug("Using cached health check results")
|
||||
return cached_results
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error getting cached health check results: %s", str(e))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@with_service_target(HEALTH_CHECKS_TARGET)
|
||||
|
|
@ -173,7 +155,7 @@ class SharedHealthCheckManager:
|
|||
healthy_endpoints: List of healthy endpoints
|
||||
unhealthy_endpoints: List of unhealthy endpoints
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
if self._store is None:
|
||||
return
|
||||
|
||||
try:
|
||||
|
|
@ -186,12 +168,9 @@ class SharedHealthCheckManager:
|
|||
"checked_by": self.pod_id,
|
||||
}
|
||||
|
||||
cache_key: Final = self.get_health_check_cache_key()
|
||||
await self.redis_cache.async_set_cache(
|
||||
cache_key,
|
||||
safe_dumps(cache_data),
|
||||
ttl=self.health_check_ttl,
|
||||
)
|
||||
written: Final = await self._store.put(cache_data, ttl=timedelta(seconds=self.health_check_ttl))
|
||||
if isinstance(written, RedisFailure):
|
||||
return
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"Cached health check results for %d healthy and %d unhealthy endpoints",
|
||||
|
|
@ -199,8 +178,8 @@ class SharedHealthCheckManager:
|
|||
len(unhealthy_endpoints),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error caching health check results: %s", str(e))
|
||||
except Exception:
|
||||
return
|
||||
|
||||
@with_service_target(HEALTH_CHECKS_TARGET)
|
||||
async def perform_shared_health_check(
|
||||
|
|
@ -274,7 +253,7 @@ class SharedHealthCheckManager:
|
|||
else:
|
||||
# If Redis is not configured, skip polling — there is no cache
|
||||
# to wait for.
|
||||
if self.redis_cache is None:
|
||||
if self._store is None:
|
||||
return await perform_health_check(
|
||||
model_list=model_list,
|
||||
details=details,
|
||||
|
|
@ -312,9 +291,12 @@ class SharedHealthCheckManager:
|
|||
# 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:
|
||||
if (
|
||||
isinstance(current_owner := await self._store.owner(), RedisFailure)
|
||||
and current_owner.kind == "breaker_open"
|
||||
):
|
||||
continue
|
||||
if current_owner is None or isinstance(current_owner, RedisFailure):
|
||||
verbose_proxy_logger.debug(
|
||||
"Pod %s detected lock released without cache, stopping wait",
|
||||
self.pod_id,
|
||||
|
|
@ -347,15 +329,15 @@ class SharedHealthCheckManager:
|
|||
Returns:
|
||||
bool: True if health check is in progress, False otherwise
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
if self._store is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
lock_key: Final = self.get_health_check_lock_key()
|
||||
current_owner: Final = await self.redis_cache.async_get_cache(lock_key)
|
||||
return current_owner is not None and current_owner != self.pod_id
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error checking health check lock status: %s", str(e))
|
||||
current_owner: Final = await self._store.owner()
|
||||
if isinstance(current_owner, RedisFailure):
|
||||
return False
|
||||
return current_owner is not None and current_owner.value != self.pod_id
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@with_service_target(HEALTH_CHECKS_TARGET)
|
||||
|
|
@ -373,11 +355,12 @@ class SharedHealthCheckManager:
|
|||
"cache_ttl": self.health_check_ttl,
|
||||
}
|
||||
|
||||
if self.redis_cache is not None:
|
||||
if self._store is not None:
|
||||
try:
|
||||
# Check if there's a current lock
|
||||
lock_key: Final = self.get_health_check_lock_key()
|
||||
current_owner: Final = await self.redis_cache.async_get_cache(lock_key)
|
||||
owner: Final = await self._store.owner()
|
||||
if isinstance(owner, RedisFailure) and owner.kind == "breaker_open":
|
||||
return {**status, "error": owner.message}
|
||||
current_owner: Final = None if owner is None or isinstance(owner, RedisFailure) else owner.value
|
||||
status["lock_owner"] = current_owner
|
||||
status["lock_in_progress"] = current_owner is not None
|
||||
|
||||
|
|
|
|||
116
litellm/redis/stores/shared_health.py
Normal file
116
litellm/redis/stores/shared_health.py
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
import ast
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.redis.transport.contract import RedisFailure, RedisTransport
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class CachedHealthCheckResults(TypedDict):
|
||||
healthy_endpoints: NotRequired[ReadOnly[list[Mapping[str, object]]]]
|
||||
unhealthy_endpoints: NotRequired[ReadOnly[list[Mapping[str, object]]]]
|
||||
timestamp: NotRequired[ReadOnly[int | float]]
|
||||
checked_by: NotRequired[ReadOnly[object]]
|
||||
|
||||
|
||||
_RESULTS: Final = TypeAdapter(CachedHealthCheckResults)
|
||||
_VALUE: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SharedHealthLockOwner:
|
||||
value: object
|
||||
|
||||
|
||||
def lock_key() -> str:
|
||||
return "health_check_lock"
|
||||
|
||||
|
||||
def results_key() -> str:
|
||||
return "health_check_results"
|
||||
|
||||
|
||||
def model_lock_key(model_name: str) -> str:
|
||||
return f"{lock_key()}:{model_name}"
|
||||
|
||||
|
||||
def model_results_key(model_name: str) -> str:
|
||||
return f"{results_key()}:{model_name}"
|
||||
|
||||
|
||||
def _decode_cache_value(raw: bytes) -> object:
|
||||
decoded: Final = raw.decode("utf-8")
|
||||
try:
|
||||
return _VALUE.validate_python(json.loads(decoded))
|
||||
except ValueError:
|
||||
return _VALUE.validate_python(ast.literal_eval(decoded))
|
||||
|
||||
|
||||
class SharedHealthStore(Protocol):
|
||||
async def acquire(self, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ...
|
||||
|
||||
async def owner(self) -> SharedHealthLockOwner | None | RedisFailure: ...
|
||||
|
||||
async def release(self, owner: str) -> bool | RedisFailure: ...
|
||||
|
||||
async def put(self, snapshot: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: ...
|
||||
|
||||
async def get(self) -> CachedHealthCheckResults | None | RedisFailure: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TransportSharedHealthStore:
|
||||
transport: RedisTransport
|
||||
|
||||
async def acquire(self, owner: str, *, ttl: timedelta) -> bool | RedisFailure:
|
||||
if ttl <= timedelta(0):
|
||||
return RedisFailure(kind="server", message="invalid expire time in 'set' command")
|
||||
return await self.transport.set(lock_key(), json.dumps(owner), ttl=ttl, nx=True)
|
||||
|
||||
async def owner(self) -> SharedHealthLockOwner | None | RedisFailure:
|
||||
raw: Final = await self.transport.get(lock_key())
|
||||
if raw is None or isinstance(raw, RedisFailure):
|
||||
return raw
|
||||
try:
|
||||
value: Final = _decode_cache_value(raw)
|
||||
return None if value is None else SharedHealthLockOwner(value)
|
||||
except Exception as exc:
|
||||
return RedisFailure(kind="decode", message=f"shared health owner is invalid: {exc}")
|
||||
|
||||
async def release(self, owner: str) -> bool | RedisFailure:
|
||||
held: Final = await self.owner()
|
||||
if isinstance(held, RedisFailure):
|
||||
return held
|
||||
if held is None or held.value != owner:
|
||||
return False
|
||||
deleted: Final = await self.transport.delete(lock_key())
|
||||
return deleted if isinstance(deleted, RedisFailure) else bool(deleted)
|
||||
|
||||
async def put(self, snapshot: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure:
|
||||
if ttl <= timedelta(0):
|
||||
return RedisFailure(kind="server", message="invalid expire time in 'set' command")
|
||||
try:
|
||||
encoded: Final = json.dumps(safe_dumps(snapshot))
|
||||
except Exception as exc:
|
||||
return RedisFailure(kind="decode", message=f"shared health snapshot could not be encoded: {exc}")
|
||||
written: Final = await self.transport.set(results_key(), encoded, ttl=ttl)
|
||||
return written if isinstance(written, RedisFailure) else None
|
||||
|
||||
async def get(self) -> CachedHealthCheckResults | None | RedisFailure:
|
||||
raw: Final = await self.transport.get(results_key())
|
||||
if raw is None or isinstance(raw, RedisFailure):
|
||||
return raw
|
||||
try:
|
||||
value: Final = _decode_cache_value(raw)
|
||||
if value is None:
|
||||
return None
|
||||
return _RESULTS.validate_python(json.loads(value) if isinstance(value, str) else value)
|
||||
except Exception as exc:
|
||||
return RedisFailure(kind="decode", message=f"shared health snapshot is invalid: {exc}")
|
||||
|
|
@ -1,605 +1,360 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import fakeredis
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from pydantic import TypeAdapter
|
||||
from redis.exceptions import TimeoutError as RedisTimeoutError
|
||||
|
||||
from litellm.proxy.health_check_utils.shared_health_check_manager import (
|
||||
SharedHealthCheckManager,
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager
|
||||
from tests._support.fake_redis import fake_redis_cache
|
||||
|
||||
NOW: Final = 1_700_000_000.25
|
||||
LOCK: Final = "shared:health_check_lock"
|
||||
RESULTS: Final = "shared:health_check_results"
|
||||
_VALUE: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HealthRedis:
|
||||
server: fakeredis.FakeServer
|
||||
raw: fakeredis.FakeAsyncRedis
|
||||
cache: RedisCache
|
||||
manager: SharedHealthCheckManager
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(loop_scope="function")
|
||||
async def health(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[HealthRedis]:
|
||||
monkeypatch.setattr(time, "time", lambda: NOW)
|
||||
server: Final = fakeredis.FakeServer()
|
||||
raw: Final = fakeredis.FakeAsyncRedis(server=server)
|
||||
cache: Final = fake_redis_cache(server, namespace="shared")
|
||||
manager: Final = SharedHealthCheckManager(cache, health_check_ttl=300, lock_ttl=60)
|
||||
manager.pod_id = 'pod_雪"'
|
||||
try:
|
||||
yield HealthRedis(server, raw, cache, manager)
|
||||
finally:
|
||||
await raw.aclose()
|
||||
await cache.init_async_client().aclose()
|
||||
|
||||
|
||||
def test_public_key_helpers_keep_global_and_model_keys() -> None:
|
||||
assert SharedHealthCheckManager.get_health_check_lock_key() == "health_check_lock"
|
||||
assert SharedHealthCheckManager.get_health_check_cache_key() == "health_check_results"
|
||||
assert SharedHealthCheckManager.get_model_health_check_lock_key("a:雪") == "health_check_lock:a:雪"
|
||||
assert SharedHealthCheckManager.get_model_health_check_cache_key("a:雪") == "health_check_results:a:雪"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_old_lock_writer_new_status_and_owner_checked_release(health: HealthRedis) -> None:
|
||||
await health.cache.async_set_cache("health_check_lock", "other_雪", ttl=60, nx=True)
|
||||
assert await health.manager.acquire_health_check_lock() is False
|
||||
assert await health.manager.is_health_check_in_progress() is True
|
||||
assert await health.manager.get_health_check_status() == {
|
||||
"pod_id": health.manager.pod_id,
|
||||
"redis_available": True,
|
||||
"lock_ttl": 60,
|
||||
"cache_ttl": 300,
|
||||
"lock_owner": "other_雪",
|
||||
"lock_in_progress": True,
|
||||
"cache_available": False,
|
||||
}
|
||||
await health.manager.release_health_check_lock()
|
||||
assert await health.raw.get(LOCK) == b'"other_\\u96ea"'
|
||||
health.manager.pod_id = "other_雪"
|
||||
assert await health.manager.is_health_check_in_progress() is False
|
||||
await health.manager.release_health_check_lock()
|
||||
assert await health.raw.get(LOCK) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_lock_writer_old_reader_exact_bytes_namespace_and_expiry(health: HealthRedis) -> None:
|
||||
assert await health.manager.acquire_health_check_lock() is True
|
||||
assert await health.raw.get(LOCK) == b'"pod_\\u96ea\\""'
|
||||
assert await health.cache.async_get_cache("health_check_lock") == health.manager.pod_id
|
||||
assert await health.raw.pttl(LOCK) == 60_000
|
||||
assert await health.raw.keys("*") == [LOCK.encode()]
|
||||
assert await health.manager.acquire_health_check_lock() is False
|
||||
await health.manager.release_health_check_lock()
|
||||
assert await health.raw.get(LOCK) is None
|
||||
assert await health.manager.is_health_check_in_progress() is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_snapshot_new_writer_old_reader_exact_double_encoding_and_ttl(health: HealthRedis) -> None:
|
||||
healthy: Final = [{"model": "雪", "latency": 1.25, "count": 2, "text": "a\x00b"}]
|
||||
unhealthy: Final = [{"model": "bad", "exception_status": 503}]
|
||||
expected: Final = {
|
||||
"healthy_endpoints": [{"model": "雪", "latency": 1.25, "count": 2, "text": "ab"}],
|
||||
"unhealthy_endpoints": unhealthy,
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 1,
|
||||
"timestamp": NOW,
|
||||
"checked_by": health.manager.pod_id,
|
||||
}
|
||||
await health.manager.cache_health_check_results(healthy, unhealthy)
|
||||
await health.cache.async_set_cache("old_writer", safe_dumps(expected), ttl=300)
|
||||
assert await health.raw.get(RESULTS) == await health.raw.get("shared:old_writer")
|
||||
assert await health.raw.get(RESULTS) == json.dumps(json.dumps(expected)).encode()
|
||||
assert await health.raw.pttl(RESULTS) == 300_000
|
||||
assert await health.raw.get("shared:shared:health_check_results") is None
|
||||
assert _VALUE.validate_json(await health.cache.async_get_cache("health_check_results")) == expected
|
||||
assert await health.manager.get_cached_health_check_results() == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("encoding", ["double", "single", "literal"])
|
||||
async def test_old_snapshot_formats_preserve_extra_fields_and_cached_flow(health: HealthRedis, encoding: str) -> None:
|
||||
snapshot: Final = {
|
||||
"healthy_endpoints": [{"model": "雪", "latency": 1.25}],
|
||||
"unhealthy_endpoints": [],
|
||||
"timestamp": NOW - 300,
|
||||
"checked_by": {"legacy": "pod"},
|
||||
"extra": [0, False, None],
|
||||
}
|
||||
if encoding == "double":
|
||||
await health.cache.async_set_cache("health_check_results", safe_dumps(snapshot), ttl=300)
|
||||
elif encoding == "single":
|
||||
await health.cache.async_set_cache("health_check_results", snapshot, ttl=300)
|
||||
else:
|
||||
await health.raw.set(RESULTS, repr(snapshot), ex=300)
|
||||
assert await health.manager.get_cached_health_check_results() == snapshot
|
||||
assert await health.manager.perform_shared_health_check([]) == (snapshot["healthy_endpoints"], [], {})
|
||||
assert await health.manager.get_health_check_status() == {
|
||||
"pod_id": health.manager.pod_id,
|
||||
"redis_available": True,
|
||||
"lock_ttl": 60,
|
||||
"cache_ttl": 300,
|
||||
"lock_owner": None,
|
||||
"lock_in_progress": False,
|
||||
"cache_available": True,
|
||||
"cache_age_seconds": 300,
|
||||
"last_checked_by": {"legacy": "pod"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"snapshot",
|
||||
[
|
||||
{},
|
||||
{"timestamp": NOW - 301},
|
||||
{"timestamp": str(NOW)},
|
||||
{"timestamp": True},
|
||||
{"timestamp": None},
|
||||
{"timestamp": NOW, "healthy_endpoints": None},
|
||||
{"timestamp": NOW, "unhealthy_endpoints": [1]},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestSharedHealthCheckManager:
|
||||
"""Test cases for SharedHealthCheckManager"""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis_cache(self):
|
||||
"""Mock Redis cache for testing"""
|
||||
cache = AsyncMock()
|
||||
cache.async_set_cache = AsyncMock()
|
||||
cache.async_get_cache = AsyncMock()
|
||||
cache.async_delete_cache = AsyncMock()
|
||||
return cache
|
||||
|
||||
@pytest.fixture
|
||||
def shared_health_manager(self, mock_redis_cache):
|
||||
"""Create SharedHealthCheckManager instance for testing"""
|
||||
return SharedHealthCheckManager(
|
||||
redis_cache=mock_redis_cache,
|
||||
health_check_ttl=300,
|
||||
lock_ttl=60,
|
||||
)
|
||||
|
||||
def test_initialization(self, mock_redis_cache):
|
||||
"""Test SharedHealthCheckManager initialization"""
|
||||
manager = SharedHealthCheckManager(
|
||||
redis_cache=mock_redis_cache,
|
||||
health_check_ttl=300,
|
||||
lock_ttl=60,
|
||||
)
|
||||
|
||||
assert manager.redis_cache == mock_redis_cache
|
||||
assert manager.health_check_ttl == 300
|
||||
assert manager.lock_ttl == 60
|
||||
assert manager.pod_id.startswith("pod_")
|
||||
|
||||
def test_initialization_without_redis(self):
|
||||
"""Test SharedHealthCheckManager initialization without Redis"""
|
||||
manager = SharedHealthCheckManager(redis_cache=None)
|
||||
|
||||
assert manager.redis_cache is None
|
||||
assert manager.health_check_ttl == 300 # Default value
|
||||
assert manager.lock_ttl == 60 # Default value
|
||||
|
||||
def test_get_health_check_lock_key(self):
|
||||
"""Test getting health check lock key"""
|
||||
key = SharedHealthCheckManager.get_health_check_lock_key()
|
||||
assert key == "health_check_lock"
|
||||
|
||||
def test_get_health_check_cache_key(self):
|
||||
"""Test getting health check cache key"""
|
||||
key = SharedHealthCheckManager.get_health_check_cache_key()
|
||||
assert key == "health_check_results"
|
||||
|
||||
def test_get_model_health_check_lock_key(self):
|
||||
"""Test getting model-specific health check lock key"""
|
||||
key = SharedHealthCheckManager.get_model_health_check_lock_key("test-model")
|
||||
assert key == "health_check_lock:test-model"
|
||||
|
||||
def test_get_model_health_check_cache_key(self):
|
||||
"""Test getting model-specific health check cache key"""
|
||||
key = SharedHealthCheckManager.get_model_health_check_cache_key("test-model")
|
||||
assert key == "health_check_results:test-model"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_health_check_lock_success(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test successful lock acquisition"""
|
||||
mock_redis_cache.async_set_cache.return_value = True
|
||||
|
||||
result = await shared_health_manager.acquire_health_check_lock()
|
||||
|
||||
assert result is True
|
||||
mock_redis_cache.async_set_cache.assert_called_once_with(
|
||||
"health_check_lock",
|
||||
shared_health_manager.pod_id,
|
||||
nx=True,
|
||||
ttl=60,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_health_check_lock_failure(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test failed lock acquisition"""
|
||||
mock_redis_cache.async_set_cache.return_value = False
|
||||
|
||||
result = await shared_health_manager.acquire_health_check_lock()
|
||||
|
||||
assert result is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_health_check_lock_no_redis(self):
|
||||
"""Test lock acquisition without Redis"""
|
||||
manager = SharedHealthCheckManager(redis_cache=None)
|
||||
|
||||
result = await manager.acquire_health_check_lock()
|
||||
|
||||
assert result is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_health_check_lock_exception(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test lock acquisition with exception"""
|
||||
mock_redis_cache.async_set_cache.side_effect = Exception("Redis error")
|
||||
|
||||
result = await shared_health_manager.acquire_health_check_lock()
|
||||
|
||||
assert result is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_health_check_lock_success(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test successful lock release"""
|
||||
mock_redis_cache.async_get_cache.return_value = shared_health_manager.pod_id
|
||||
|
||||
await shared_health_manager.release_health_check_lock()
|
||||
|
||||
mock_redis_cache.async_get_cache.assert_called_once_with("health_check_lock")
|
||||
mock_redis_cache.async_delete_cache.assert_called_once_with("health_check_lock")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_health_check_lock_wrong_owner(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test lock release when not the owner"""
|
||||
mock_redis_cache.async_get_cache.return_value = "other_pod_id"
|
||||
|
||||
await shared_health_manager.release_health_check_lock()
|
||||
|
||||
mock_redis_cache.async_get_cache.assert_called_once_with("health_check_lock")
|
||||
mock_redis_cache.async_delete_cache.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_health_check_lock_no_redis(self):
|
||||
"""Test lock release without Redis"""
|
||||
manager = SharedHealthCheckManager(redis_cache=None)
|
||||
|
||||
# Should not raise exception
|
||||
await manager.release_health_check_lock()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_health_check_results_success(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test getting cached health check results successfully"""
|
||||
current_time = time.time()
|
||||
cached_data = {
|
||||
"healthy_endpoints": [{"model": "test-model"}],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"timestamp": current_time - 100, # 100 seconds ago
|
||||
"checked_by": "test_pod",
|
||||
}
|
||||
mock_redis_cache.async_get_cache.return_value = json.dumps(cached_data)
|
||||
|
||||
result = await shared_health_manager.get_cached_health_check_results()
|
||||
|
||||
assert result is not None
|
||||
assert result["healthy_count"] == 1
|
||||
assert result["unhealthy_count"] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_health_check_results_expired(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test getting expired cached health check results"""
|
||||
current_time = time.time()
|
||||
cached_data = {
|
||||
"healthy_endpoints": [{"model": "test-model"}],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"timestamp": current_time - 400, # 400 seconds ago (expired)
|
||||
"checked_by": "test_pod",
|
||||
}
|
||||
mock_redis_cache.async_get_cache.return_value = json.dumps(cached_data)
|
||||
|
||||
result = await shared_health_manager.get_cached_health_check_results()
|
||||
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_health_check_results_no_cache(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test getting cached results when no cache exists"""
|
||||
mock_redis_cache.async_get_cache.return_value = None
|
||||
|
||||
result = await shared_health_manager.get_cached_health_check_results()
|
||||
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_health_check_results_no_redis(self):
|
||||
"""Test getting cached results without Redis"""
|
||||
manager = SharedHealthCheckManager(redis_cache=None)
|
||||
|
||||
result = await manager.get_cached_health_check_results()
|
||||
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_health_check_results_success(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test caching health check results successfully"""
|
||||
healthy_endpoints = [{"model": "test-model-1"}]
|
||||
unhealthy_endpoints = [{"model": "test-model-2"}]
|
||||
|
||||
await shared_health_manager.cache_health_check_results(
|
||||
healthy_endpoints, unhealthy_endpoints
|
||||
)
|
||||
|
||||
mock_redis_cache.async_set_cache.assert_called_once()
|
||||
call_args = mock_redis_cache.async_set_cache.call_args
|
||||
assert call_args[0][0] == "health_check_results" # key
|
||||
assert call_args[1]["ttl"] == 300 # ttl
|
||||
|
||||
# Verify cached data structure
|
||||
cached_data = json.loads(call_args[0][1])
|
||||
assert cached_data["healthy_endpoints"] == healthy_endpoints
|
||||
assert cached_data["unhealthy_endpoints"] == unhealthy_endpoints
|
||||
assert cached_data["healthy_count"] == 1
|
||||
assert cached_data["unhealthy_count"] == 1
|
||||
assert "timestamp" in cached_data
|
||||
assert cached_data["checked_by"] == shared_health_manager.pod_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_health_check_results_no_redis(self):
|
||||
"""Test caching results without Redis"""
|
||||
manager = SharedHealthCheckManager(redis_cache=None)
|
||||
|
||||
# Should not raise exception
|
||||
await manager.cache_health_check_results([], [])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_perform_shared_health_check_with_cache(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test performing shared health check when cache is available"""
|
||||
# Mock cached results
|
||||
cached_data = {
|
||||
"healthy_endpoints": [{"model": "cached-model"}],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"timestamp": time.time() - 100,
|
||||
}
|
||||
mock_redis_cache.async_get_cache.return_value = json.dumps(cached_data)
|
||||
|
||||
model_list = [
|
||||
{"model_name": "test-model", "litellm_params": {"model": "test-model"}}
|
||||
]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check"
|
||||
) as mock_perform:
|
||||
healthy, unhealthy, _ = (
|
||||
await shared_health_manager.perform_shared_health_check(
|
||||
model_list, details=True
|
||||
)
|
||||
)
|
||||
|
||||
# Should return cached results, not call perform_health_check
|
||||
assert healthy == [{"model": "cached-model"}]
|
||||
assert unhealthy == []
|
||||
mock_perform.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_perform_shared_health_check_with_lock_acquisition(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test performing shared health check when acquiring lock"""
|
||||
# No cached results
|
||||
mock_redis_cache.async_get_cache.return_value = None
|
||||
# Lock acquisition succeeds
|
||||
mock_redis_cache.async_set_cache.return_value = True
|
||||
|
||||
model_list = [
|
||||
{"model_name": "test-model", "litellm_params": {"model": "test-model"}}
|
||||
]
|
||||
expected_healthy = [{"model": "test-model", "status": "healthy"}]
|
||||
expected_unhealthy = []
|
||||
|
||||
with 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 shared_health_manager.perform_shared_health_check(
|
||||
model_list, details=True
|
||||
)
|
||||
)
|
||||
|
||||
# Should call perform_health_check and cache results
|
||||
mock_perform.assert_called_once_with(
|
||||
model_list=model_list,
|
||||
details=True,
|
||||
max_concurrency=None,
|
||||
health_check_skip_disabled_background_models=False,
|
||||
router=None,
|
||||
)
|
||||
assert healthy == expected_healthy
|
||||
assert unhealthy == expected_unhealthy
|
||||
|
||||
# Should cache the results
|
||||
assert mock_redis_cache.async_set_cache.call_count >= 2 # Lock + cache
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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 during polling"""
|
||||
# First call: no cache, lock fails
|
||||
# Polling finds cache on first iteration
|
||||
mock_redis_cache.async_get_cache.side_effect = [
|
||||
None, # No cache initially (get_cached_health_check_results)
|
||||
json.dumps(
|
||||
{ # Cache available on first poll iteration
|
||||
"healthy_endpoints": [{"model": "cached-model"}],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"timestamp": time.time() - 100,
|
||||
}
|
||||
),
|
||||
]
|
||||
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: # Mock sleep to avoid actual delay
|
||||
healthy, unhealthy, _ = (
|
||||
await shared_health_manager.perform_shared_health_check(
|
||||
model_list, details=True
|
||||
)
|
||||
)
|
||||
|
||||
# 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(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,
|
||||
health_check_skip_disabled_background_models=False,
|
||||
router=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 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 = [
|
||||
{"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 shared_health_manager.perform_shared_health_check(
|
||||
model_list, details=True
|
||||
)
|
||||
)
|
||||
|
||||
# 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,
|
||||
health_check_skip_disabled_background_models=False,
|
||||
router=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,
|
||||
health_check_skip_disabled_background_models=False,
|
||||
router=None,
|
||||
)
|
||||
assert healthy == expected_healthy
|
||||
assert unhealthy == expected_unhealthy
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_health_check_in_progress_true(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test checking if health check is in progress when it is"""
|
||||
mock_redis_cache.async_get_cache.return_value = "other_pod_id"
|
||||
|
||||
result = await shared_health_manager.is_health_check_in_progress()
|
||||
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_health_check_in_progress_false(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test checking if health check is in progress when it's not"""
|
||||
mock_redis_cache.async_get_cache.return_value = None
|
||||
|
||||
result = await shared_health_manager.is_health_check_in_progress()
|
||||
|
||||
assert result is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_health_check_in_progress_own_lock(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test checking if health check is in progress when we own the lock"""
|
||||
mock_redis_cache.async_get_cache.return_value = shared_health_manager.pod_id
|
||||
|
||||
result = await shared_health_manager.is_health_check_in_progress()
|
||||
|
||||
assert result is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_health_check_status(
|
||||
self, shared_health_manager, mock_redis_cache
|
||||
):
|
||||
"""Test getting health check status"""
|
||||
current_time = time.time()
|
||||
cached_data = {
|
||||
"healthy_endpoints": [{"model": "test-model"}],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"timestamp": current_time - 100,
|
||||
"checked_by": "test_pod",
|
||||
}
|
||||
|
||||
mock_redis_cache.async_get_cache.side_effect = [
|
||||
"other_pod_id", # Lock owner
|
||||
json.dumps(cached_data), # Cached results
|
||||
]
|
||||
|
||||
status = await shared_health_manager.get_health_check_status()
|
||||
|
||||
assert status["pod_id"] == shared_health_manager.pod_id
|
||||
assert status["redis_available"] is True
|
||||
assert status["lock_ttl"] == 60
|
||||
assert status["cache_ttl"] == 300
|
||||
assert status["lock_owner"] == "other_pod_id"
|
||||
assert status["lock_in_progress"] is True
|
||||
assert status["cache_available"] is True
|
||||
assert status["last_checked_by"] == "test_pod"
|
||||
assert "cache_age_seconds" in status
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_health_check_status_no_redis(self):
|
||||
"""Test getting health check status without Redis"""
|
||||
manager = SharedHealthCheckManager(redis_cache=None)
|
||||
|
||||
status = await manager.get_health_check_status()
|
||||
|
||||
assert status["pod_id"] == manager.pod_id
|
||||
assert status["redis_available"] is False
|
||||
assert status["lock_ttl"] == 60
|
||||
assert status["cache_ttl"] == 300
|
||||
async def test_stale_or_schema_invalid_snapshot_is_a_miss(health: HealthRedis, snapshot: Mapping[str, object]) -> None:
|
||||
await health.cache.async_set_cache("health_check_results", safe_dumps(snapshot), ttl=300)
|
||||
assert await health.manager.get_cached_health_check_results() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("raw", [b"broken", b'"broken"', b"[]", b"null", b"\xff", b'"null"'])
|
||||
async def test_malformed_snapshot_and_null_are_misses(health: HealthRedis, raw: bytes) -> None:
|
||||
await health.raw.set(RESULTS, raw)
|
||||
assert await health.manager.get_cached_health_check_results() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_endpoints_and_future_timestamp_remain_valid(health: HealthRedis) -> None:
|
||||
snapshot: Final = {"timestamp": NOW + 1}
|
||||
await health.cache.async_set_cache("health_check_results", safe_dumps(snapshot), ttl=300)
|
||||
assert await health.manager.get_cached_health_check_results() == snapshot
|
||||
assert await health.manager.perform_shared_health_check([]) == ([], [], {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lock_holder_executes_health_check_caches_and_releases(health: HealthRedis) -> None:
|
||||
assert await health.manager.perform_shared_health_check([]) == ([], [], {})
|
||||
assert await health.manager.get_cached_health_check_results() == {
|
||||
"healthy_endpoints": [],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 0,
|
||||
"unhealthy_count": 0,
|
||||
"timestamp": NOW,
|
||||
"checked_by": health.manager.pod_id,
|
||||
}
|
||||
assert await health.raw.get(LOCK) is None
|
||||
assert await health.raw.pttl(RESULTS) == 300_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_waiter_reads_snapshot_published_during_poll(health: HealthRedis) -> None:
|
||||
await health.cache.async_set_cache("health_check_lock", "other", ttl=60)
|
||||
|
||||
async def publish(delay: float) -> None:
|
||||
assert delay == 5
|
||||
await health.manager.cache_health_check_results([{"model": "published"}], [])
|
||||
|
||||
with patch("asyncio.sleep", new=publish):
|
||||
result: Final = await health.manager.perform_shared_health_check([])
|
||||
assert result == ([{"model": "published"}], [], {})
|
||||
assert await health.raw.get(LOCK) == b'"other"'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("change", ["orphaned", "outage", "held"])
|
||||
async def test_wait_policy_falls_back_locally_without_writing_snapshot(health: HealthRedis, change: str) -> None:
|
||||
health.manager.lock_ttl = 10
|
||||
await health.cache.async_set_cache("health_check_lock", "other", ttl=60)
|
||||
|
||||
async def during_wait(delay: float) -> None:
|
||||
assert delay == 5
|
||||
if change == "orphaned":
|
||||
await health.raw.delete(LOCK)
|
||||
elif change == "outage":
|
||||
health.server.connected = False
|
||||
|
||||
with patch("asyncio.sleep", wraps=during_wait) as sleep:
|
||||
assert await health.manager.perform_shared_health_check([]) == ([], [], {})
|
||||
assert sleep.await_count == (2 if change == "held" else 1)
|
||||
health.server.connected = True
|
||||
assert await health.raw.get(RESULTS) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_outage_is_false_or_miss_then_breaker_refuses_after_recovery(health: HealthRedis) -> None:
|
||||
health.server.connected = False
|
||||
assert await health.manager.acquire_health_check_lock() is False
|
||||
assert await health.manager.get_cached_health_check_results() is None
|
||||
assert await health.manager.is_health_check_in_progress() is False
|
||||
assert await health.manager.release_health_check_lock() is None
|
||||
assert await health.manager.cache_health_check_results([], []) is None
|
||||
health.server.connected = True
|
||||
assert await health.manager.acquire_health_check_lock() is False
|
||||
assert await health.raw.get(LOCK) is None
|
||||
status: Final = await health.manager.get_health_check_status()
|
||||
assert "error" in status
|
||||
assert "lock_owner" not in status
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("ttl", [0, -1])
|
||||
async def test_nonpositive_ttl_does_not_write_a_short_lived_lock_or_snapshot(health: HealthRedis, ttl: int) -> None:
|
||||
health.manager.lock_ttl = ttl
|
||||
health.manager.health_check_ttl = ttl
|
||||
assert await health.manager.acquire_health_check_lock() is False
|
||||
await health.manager.cache_health_check_results([], [])
|
||||
assert await health.raw.get(LOCK) is None
|
||||
assert await health.raw.get(RESULTS) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lost_write_replies_remain_unknown_applied(health: HealthRedis) -> None:
|
||||
async def applied_set(
|
||||
name: str, value: str, *, px: int | None = None, ex: int | None = None, nx: bool = False
|
||||
) -> None:
|
||||
await health.raw.set(name, value, px=px, ex=ex, nx=nx)
|
||||
raise RedisTimeoutError("reply lost after set")
|
||||
|
||||
with patch.object(health.cache.init_async_client(), "set", new=applied_set):
|
||||
assert await health.manager.acquire_health_check_lock() is False
|
||||
await health.manager.cache_health_check_results([], [])
|
||||
assert await health.raw.get(LOCK) == json.dumps(health.manager.pod_id).encode()
|
||||
assert await health.raw.pttl(LOCK) == 60_000
|
||||
assert await health.manager.get_cached_health_check_results() == {
|
||||
"healthy_endpoints": [],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 0,
|
||||
"unhealthy_count": 0,
|
||||
"timestamp": NOW,
|
||||
"checked_by": health.manager.pod_id,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("applied", [False, True])
|
||||
async def test_lost_delete_reply_is_best_effort(health: HealthRedis, applied: bool) -> None:
|
||||
assert await health.manager.acquire_health_check_lock() is True
|
||||
|
||||
async def delete_timeout(name: str) -> None:
|
||||
if applied:
|
||||
await health.raw.delete(name)
|
||||
raise RedisTimeoutError("reply lost during delete")
|
||||
|
||||
with patch.object(health.cache.init_async_client(), "delete", new=delete_timeout):
|
||||
assert await health.manager.release_health_check_lock() is None
|
||||
assert await health.raw.get(LOCK) == (None if applied else json.dumps(health.manager.pod_id).encode())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_keeps_existing_get_then_delete_race(health: HealthRedis) -> None:
|
||||
assert await health.manager.acquire_health_check_lock() is True
|
||||
|
||||
async def replace_after_read(name: str) -> bytes | None:
|
||||
previous: Final = await health.raw.get(name)
|
||||
await health.raw.set(name, json.dumps("replacement"), ex=60)
|
||||
return previous
|
||||
|
||||
with patch.object(health.cache.init_async_client(), "get", new=replace_after_read):
|
||||
await health.manager.release_health_check_lock()
|
||||
assert await health.raw.get(LOCK) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_redis_runs_locally_without_polling() -> None:
|
||||
manager: Final = SharedHealthCheckManager()
|
||||
|
||||
async def no_wait(delay: float) -> None:
|
||||
pytest.fail(f"unexpected local polling delay {delay}")
|
||||
|
||||
with patch("asyncio.sleep", new=no_wait):
|
||||
assert await manager.perform_shared_health_check([]) == ([], [], {})
|
||||
assert await manager.acquire_health_check_lock() is False
|
||||
assert await manager.is_health_check_in_progress() is False
|
||||
assert await manager.get_cached_health_check_results() is None
|
||||
assert await manager.release_health_check_lock() is None
|
||||
assert await manager.cache_health_check_results([], []) is None
|
||||
assert await manager.get_health_check_status() == {
|
||||
"pod_id": manager.pod_id,
|
||||
"redis_available": False,
|
||||
"lock_ttl": 60,
|
||||
"cache_ttl": 300,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"raw, expected",
|
||||
[
|
||||
(b"42", 42),
|
||||
(b"[]", []),
|
||||
(b"false", False),
|
||||
(b'""', ""),
|
||||
(b"'legacy'", "legacy"),
|
||||
(b"null", None),
|
||||
(b"broken", None),
|
||||
(b"\xff", None),
|
||||
],
|
||||
)
|
||||
async def test_legacy_owner_decoding_keeps_lock_presence(health: HealthRedis, raw: bytes, expected: object) -> None:
|
||||
await health.raw.set(LOCK, raw)
|
||||
assert await health.manager.is_health_check_in_progress() is (expected is not None)
|
||||
status: Final = await health.manager.get_health_check_status()
|
||||
assert status["lock_owner"] == expected
|
||||
assert status["lock_in_progress"] is (expected is not None)
|
||||
await health.manager.release_health_check_lock()
|
||||
assert await health.raw.get(LOCK) == raw
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keys_expire_at_caller_selected_seconds(health: HealthRedis, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
assert await health.manager.acquire_health_check_lock() is True
|
||||
await health.manager.cache_health_check_results([], [])
|
||||
monkeypatch.setattr(time, "time", lambda: NOW + 61)
|
||||
assert await health.raw.pttl(LOCK) == -2
|
||||
assert await health.raw.pttl(RESULTS) == 239_000
|
||||
assert await health.manager.is_health_check_in_progress() is False
|
||||
monkeypatch.setattr(time, "time", lambda: NOW + 301)
|
||||
assert await health.manager.get_cached_health_check_results() is None
|
||||
assert await health.raw.pttl(RESULTS) == -2
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue