mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(proxy): prevent HealthCheckTable unbounded growth
LiteLLM_HealthCheckTable grows without bound when external monitoring tools (e.g. Uptime Kuma) ping health endpoints continuously. With 500 k+ rows the dashboard query loads every row into Python memory, causing 8 GB+ container memory usage and 400% CPU spikes. Two changes fix this: 1. get_all_latest_health_checks: add a checked_at >= (now - TTL) WHERE filter so only recent rows are fetched. Previously the method did a full-table find_many with no bounds and deduplicated in Python. With the filter the DB does the work and Python only sees O(models) rows. 2. cleanup_old_health_checks: new method that deletes rows older than TTL via a single delete_many call. save_health_check_result calls it at most once per hour (tracked via _health_check_last_cleanup_ts) to keep the table bounded without adding per-insert overhead. TTL defaults to 7 days and is configurable via HEALTH_CHECK_TTL_DAYS. Fixes #25623
This commit is contained in:
Claude
parent
d319cd8cc6
commit
0009995729
2 changed files with 324 additions and 115 deletions
|
|
@ -1965,9 +1965,9 @@ class ProxyLogging:
|
|||
normalized_call_type = CallTypes.aembedding.value
|
||||
if normalized_call_type is not None:
|
||||
litellm_logging_obj.call_type = normalized_call_type
|
||||
litellm_logging_obj.model_call_details[
|
||||
"call_type"
|
||||
] = normalized_call_type
|
||||
litellm_logging_obj.model_call_details["call_type"] = (
|
||||
normalized_call_type
|
||||
)
|
||||
# Pass-through endpoints are logged via the callback loop's
|
||||
# async_post_call_failure_hook — skip pre_call and failure handlers.
|
||||
if litellm_logging_obj.call_type == CallTypes.pass_through.value:
|
||||
|
|
@ -2065,10 +2065,12 @@ class ProxyLogging:
|
|||
raise
|
||||
else:
|
||||
try:
|
||||
guardrail_response = await callback.async_post_call_success_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
response=response,
|
||||
guardrail_response = (
|
||||
await callback.async_post_call_success_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
response=response,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
_enrich_http_exception_with_guardrail_context(e, callback)
|
||||
|
|
@ -2281,13 +2283,15 @@ class ProxyLogging:
|
|||
"async_post_call_streaming_iterator_hook"
|
||||
in type(callback).__dict__
|
||||
):
|
||||
current_response = self._wrap_streaming_iterator_with_enrichment(
|
||||
_callback,
|
||||
_callback.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=current_response,
|
||||
request_data=request_data,
|
||||
),
|
||||
current_response = (
|
||||
self._wrap_streaming_iterator_with_enrichment(
|
||||
_callback,
|
||||
_callback.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=current_response,
|
||||
request_data=request_data,
|
||||
),
|
||||
)
|
||||
)
|
||||
elif "apply_guardrail" in type(callback).__dict__:
|
||||
request_data["guardrail_to_apply"] = callback
|
||||
|
|
@ -2300,13 +2304,15 @@ class ProxyLogging:
|
|||
),
|
||||
)
|
||||
else:
|
||||
current_response = self._wrap_streaming_iterator_with_enrichment(
|
||||
_callback,
|
||||
_callback.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=current_response,
|
||||
request_data=request_data,
|
||||
),
|
||||
current_response = (
|
||||
self._wrap_streaming_iterator_with_enrichment(
|
||||
_callback,
|
||||
_callback.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=current_response,
|
||||
request_data=request_data,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Actually iterate through the chained async generator and yield chunks
|
||||
|
|
@ -2512,6 +2518,7 @@ class PrismaClient:
|
|||
self._watching_engine: bool = False
|
||||
self._engine_confirmed_dead: bool = False
|
||||
self._engine_wait_thread: Optional[threading.Thread] = None
|
||||
self._health_check_last_cleanup_ts: float = 0.0
|
||||
verbose_proxy_logger.debug("Success - Created Prisma Client")
|
||||
|
||||
def get_request_status(
|
||||
|
|
@ -4487,7 +4494,17 @@ class PrismaClient:
|
|||
)
|
||||
|
||||
verbose_proxy_logger.debug(f"Saving health check data: {health_check_data}")
|
||||
return await self.db.litellm_healthchecktable.create(data=health_check_data)
|
||||
result = await self.db.litellm_healthchecktable.create(
|
||||
data=health_check_data
|
||||
)
|
||||
|
||||
# Prune old rows at most once per hour to prevent unbounded table growth
|
||||
now = time.time()
|
||||
if now - self._health_check_last_cleanup_ts >= 3600:
|
||||
self._health_check_last_cleanup_ts = now
|
||||
await self.cleanup_old_health_checks()
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -4525,25 +4542,29 @@ class PrismaClient:
|
|||
|
||||
async def get_all_latest_health_checks(self):
|
||||
"""
|
||||
Get the latest health check for each model
|
||||
Get the latest health check for each model.
|
||||
|
||||
Only considers rows within the retention window (HEALTH_CHECK_TTL_DAYS) to
|
||||
avoid full-table scans on deployments with large health check histories.
|
||||
"""
|
||||
try:
|
||||
# Get all unique model names first
|
||||
health_check_ttl_days = int(os.getenv("HEALTH_CHECK_TTL_DAYS", "7"))
|
||||
cutoff = datetime.utcnow() - timedelta(days=health_check_ttl_days)
|
||||
all_checks = await self.db.litellm_healthchecktable.find_many(
|
||||
order={"checked_at": "desc"}
|
||||
where={"checked_at": {"gte": cutoff}},
|
||||
order={"checked_at": "desc"},
|
||||
)
|
||||
|
||||
# Group by model_name and get the latest for each
|
||||
# Group by model and return only the latest entry per model
|
||||
latest_checks = {}
|
||||
for check in all_checks:
|
||||
# Create a unique key: prefer model_id if available, otherwise use model_name
|
||||
# This ensures we get the latest check for each unique model
|
||||
# Prefer model_id as the dedup key; fall back to model_name
|
||||
if check.model_id:
|
||||
key = (check.model_id, check.model_name)
|
||||
else:
|
||||
key = (None, check.model_name)
|
||||
|
||||
# Only add if we haven't seen this key yet (since checks are ordered by checked_at desc)
|
||||
# First entry wins because rows are ordered by checked_at desc
|
||||
if key not in latest_checks:
|
||||
latest_checks[key] = check
|
||||
|
||||
|
|
@ -4552,6 +4573,25 @@ class PrismaClient:
|
|||
verbose_proxy_logger.error(f"Error getting all latest health checks: {e}")
|
||||
return []
|
||||
|
||||
async def cleanup_old_health_checks(self) -> None:
|
||||
"""
|
||||
Delete LiteLLM_HealthCheckTable rows older than HEALTH_CHECK_TTL_DAYS (default 7).
|
||||
|
||||
Called from save_health_check_result at most once per hour to prevent
|
||||
unbounded table growth without adding per-insert overhead.
|
||||
"""
|
||||
try:
|
||||
health_check_ttl_days = int(os.getenv("HEALTH_CHECK_TTL_DAYS", "7"))
|
||||
cutoff = datetime.utcnow() - timedelta(days=health_check_ttl_days)
|
||||
result = await self.db.litellm_healthchecktable.delete_many(
|
||||
where={"checked_at": {"lt": cutoff}}
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
f"Health check cleanup: deleted {result} rows older than {health_check_ttl_days} days"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error cleaning up old health checks: {e}")
|
||||
|
||||
|
||||
### HELPER FUNCTIONS ###
|
||||
|
||||
|
|
@ -4947,25 +4987,27 @@ async def update_daily_tag_spend(
|
|||
):
|
||||
"""
|
||||
Separate scheduler job to commit daily tag spend updates.
|
||||
|
||||
|
||||
Runs at a longer interval (2.3x default) than the main update_spend job
|
||||
to reduce query contention for DailyTagSpend table.
|
||||
|
||||
|
||||
This is called by a dedicated scheduler job and does NOT process:
|
||||
- Regular spend updates (user, key, team, org)
|
||||
- End-user spend
|
||||
- Agent spend
|
||||
- Spend logs
|
||||
|
||||
|
||||
Only processes tag spend transactions from the daily_tag_spend_update_queue.
|
||||
|
||||
|
||||
Args:
|
||||
prisma_client: PrismaClient instance
|
||||
proxy_logging_obj: ProxyLogging instance for error handling
|
||||
"""
|
||||
n_retry_times = 3
|
||||
try:
|
||||
if proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis():
|
||||
if (
|
||||
proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis()
|
||||
):
|
||||
await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis(
|
||||
prisma_client=prisma_client,
|
||||
n_retry_times=n_retry_times,
|
||||
|
|
|
|||
|
|
@ -24,33 +24,53 @@ from litellm.proxy.utils import PrismaClient
|
|||
def mock_prisma():
|
||||
"""Simplified mock PrismaClient with bound methods"""
|
||||
client = MagicMock()
|
||||
client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "test-id"})
|
||||
client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[{"id": "1", "model_name": "test"}])
|
||||
|
||||
client.db.litellm_healthchecktable.create = AsyncMock(
|
||||
return_value={"id": "test-id"}
|
||||
)
|
||||
client.db.litellm_healthchecktable.find_many = AsyncMock(
|
||||
return_value=[{"id": "1", "model_name": "test"}]
|
||||
)
|
||||
|
||||
# Bind actual methods
|
||||
import types
|
||||
for method in ['save_health_check_result', '_validate_response_time', '_clean_details',
|
||||
'get_health_check_history', 'get_all_latest_health_checks']:
|
||||
|
||||
for method in [
|
||||
"save_health_check_result",
|
||||
"_validate_response_time",
|
||||
"_clean_details",
|
||||
"get_health_check_history",
|
||||
"get_all_latest_health_checks",
|
||||
]:
|
||||
setattr(client, method, types.MethodType(getattr(PrismaClient, method), client))
|
||||
|
||||
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status,healthy,unhealthy,should_succeed", [
|
||||
("healthy", 1, 0, True),
|
||||
("unhealthy", 0, 1, True),
|
||||
("healthy", 1, 0, False), # Database error case
|
||||
])
|
||||
async def test_save_health_check_result(mock_prisma, status, healthy, unhealthy, should_succeed):
|
||||
@pytest.mark.parametrize(
|
||||
"status,healthy,unhealthy,should_succeed",
|
||||
[
|
||||
("healthy", 1, 0, True),
|
||||
("unhealthy", 0, 1, True),
|
||||
("healthy", 1, 0, False), # Database error case
|
||||
],
|
||||
)
|
||||
async def test_save_health_check_result(
|
||||
mock_prisma, status, healthy, unhealthy, should_succeed
|
||||
):
|
||||
"""Test health check result saving with various scenarios"""
|
||||
if not should_succeed:
|
||||
mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception("DB Error")
|
||||
|
||||
mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception(
|
||||
"DB Error"
|
||||
)
|
||||
|
||||
result = await mock_prisma.save_health_check_result(
|
||||
model_name="test-model", status=status, healthy_count=healthy, unhealthy_count=unhealthy
|
||||
model_name="test-model",
|
||||
status=status,
|
||||
healthy_count=healthy,
|
||||
unhealthy_count=unhealthy,
|
||||
)
|
||||
|
||||
|
||||
if should_succeed:
|
||||
mock_prisma.db.litellm_healthchecktable.create.assert_called_once()
|
||||
else:
|
||||
|
|
@ -66,24 +86,31 @@ async def test_get_health_check_history(mock_prisma):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("healthy_count,unhealthy_count,expected_status", [
|
||||
(1, 0, "healthy"),
|
||||
(0, 1, "unhealthy"),
|
||||
(2, 1, "healthy"),
|
||||
])
|
||||
@pytest.mark.parametrize(
|
||||
"healthy_count,unhealthy_count,expected_status",
|
||||
[
|
||||
(1, 0, "healthy"),
|
||||
(0, 1, "unhealthy"),
|
||||
(2, 1, "healthy"),
|
||||
],
|
||||
)
|
||||
async def test_save_health_check_to_db(healthy_count, unhealthy_count, expected_status):
|
||||
"""Test _save_health_check_to_db function with different endpoint counts"""
|
||||
mock_client = MagicMock()
|
||||
mock_client.save_health_check_result = AsyncMock()
|
||||
|
||||
|
||||
healthy_endpoints = [{"model": "test"}] * healthy_count
|
||||
unhealthy_endpoints = [{"error": "test error"}] * unhealthy_count
|
||||
|
||||
|
||||
await _save_health_check_to_db(
|
||||
mock_client, "test-model", healthy_endpoints, unhealthy_endpoints,
|
||||
1234567890.0, "test-user"
|
||||
mock_client,
|
||||
"test-model",
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
1234567890.0,
|
||||
"test-user",
|
||||
)
|
||||
|
||||
|
||||
call_args = mock_client.save_health_check_result.call_args[1]
|
||||
assert call_args["status"] == expected_status
|
||||
assert call_args["healthy_count"] == healthy_count
|
||||
|
|
@ -99,6 +126,7 @@ async def test_save_health_check_to_db_no_client():
|
|||
|
||||
# Tests for background health check functions
|
||||
|
||||
|
||||
def test_build_model_param_to_info_mapping():
|
||||
"""Test building model parameter to info mapping"""
|
||||
model_list = [
|
||||
|
|
@ -118,9 +146,9 @@ def test_build_model_param_to_info_mapping():
|
|||
"litellm_params": {"model": "gpt-3.5-turbo"}, # Same model param
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
result = _build_model_param_to_info_mapping(model_list)
|
||||
|
||||
|
||||
assert "gpt-3.5-turbo" in result
|
||||
assert "gpt-4" in result
|
||||
assert len(result["gpt-3.5-turbo"]) == 2 # Two models share same param
|
||||
|
|
@ -139,7 +167,7 @@ def test_build_model_param_to_info_mapping_no_model_name():
|
|||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
result = _build_model_param_to_info_mapping(model_list)
|
||||
assert len(result) == 0
|
||||
|
||||
|
|
@ -154,25 +182,25 @@ def test_aggregate_health_check_results():
|
|||
{"model_name": "gpt-4", "model_id": "model-456"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
healthy_endpoints = [
|
||||
{"model": "gpt-3.5-turbo"},
|
||||
]
|
||||
unhealthy_endpoints = [
|
||||
{"model": "gpt-4", "error": "Rate limit exceeded"},
|
||||
]
|
||||
|
||||
|
||||
result = _aggregate_health_check_results(
|
||||
model_param_to_info, healthy_endpoints, unhealthy_endpoints
|
||||
)
|
||||
|
||||
|
||||
# Check gpt-3.5-turbo is healthy
|
||||
gpt35_key = ("model-123", "gpt-3.5-turbo")
|
||||
assert gpt35_key in result
|
||||
assert result[gpt35_key]["healthy_count"] == 1
|
||||
assert result[gpt35_key]["unhealthy_count"] == 0
|
||||
assert result[gpt35_key]["error_message"] is None
|
||||
|
||||
|
||||
# Check gpt-4 is unhealthy
|
||||
gpt4_key = ("model-456", "gpt-4")
|
||||
assert gpt4_key in result
|
||||
|
|
@ -188,17 +216,17 @@ def test_aggregate_health_check_results_multiple_endpoints():
|
|||
{"model_name": "gpt-3.5-turbo", "model_id": "model-123"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
healthy_endpoints = [
|
||||
{"model": "gpt-3.5-turbo"},
|
||||
{"model": "gpt-3.5-turbo"},
|
||||
]
|
||||
unhealthy_endpoints = []
|
||||
|
||||
|
||||
result = _aggregate_health_check_results(
|
||||
model_param_to_info, healthy_endpoints, unhealthy_endpoints
|
||||
)
|
||||
|
||||
|
||||
key = ("model-123", "gpt-3.5-turbo")
|
||||
assert result[key]["healthy_count"] == 2
|
||||
assert result[key]["unhealthy_count"] == 0
|
||||
|
|
@ -209,7 +237,7 @@ async def test_save_health_check_results_if_changed_status_changed():
|
|||
"""Test saving when status changes"""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.save_health_check_result = AsyncMock()
|
||||
|
||||
|
||||
model_results = {
|
||||
("model-123", "gpt-3.5-turbo"): {
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
|
|
@ -219,7 +247,7 @@ async def test_save_health_check_results_if_changed_status_changed():
|
|||
"error_message": None,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Latest check shows unhealthy, new result is healthy (status changed)
|
||||
latest_checks_map = {
|
||||
"model-123": MagicMock(
|
||||
|
|
@ -227,12 +255,16 @@ async def test_save_health_check_results_if_changed_status_changed():
|
|||
checked_at=datetime.now(timezone.utc) - timedelta(minutes=5),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
start_time = 1234567890.0
|
||||
await _save_health_check_results_if_changed(
|
||||
mock_prisma, model_results, latest_checks_map, start_time, "background_health_check"
|
||||
mock_prisma,
|
||||
model_results,
|
||||
latest_checks_map,
|
||||
start_time,
|
||||
"background_health_check",
|
||||
)
|
||||
|
||||
|
||||
# Should save because status changed
|
||||
mock_prisma.save_health_check_result.assert_called_once()
|
||||
call_kwargs = mock_prisma.save_health_check_result.call_args[1]
|
||||
|
|
@ -246,7 +278,7 @@ async def test_save_health_check_results_if_changed_status_unchanged_recent():
|
|||
"""Test skipping save when status unchanged and checked recently"""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.save_health_check_result = AsyncMock()
|
||||
|
||||
|
||||
model_results = {
|
||||
("model-123", "gpt-3.5-turbo"): {
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
|
|
@ -256,7 +288,7 @@ async def test_save_health_check_results_if_changed_status_unchanged_recent():
|
|||
"error_message": None,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Latest check shows healthy, new result is healthy (status unchanged)
|
||||
# And checked recently (within 1 hour)
|
||||
latest_checks_map = {
|
||||
|
|
@ -265,12 +297,16 @@ async def test_save_health_check_results_if_changed_status_unchanged_recent():
|
|||
checked_at=datetime.now(timezone.utc) - timedelta(minutes=30),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
start_time = 1234567890.0
|
||||
await _save_health_check_results_if_changed(
|
||||
mock_prisma, model_results, latest_checks_map, start_time, "background_health_check"
|
||||
mock_prisma,
|
||||
model_results,
|
||||
latest_checks_map,
|
||||
start_time,
|
||||
"background_health_check",
|
||||
)
|
||||
|
||||
|
||||
# Should NOT save because status unchanged and checked recently
|
||||
mock_prisma.save_health_check_result.assert_not_called()
|
||||
|
||||
|
|
@ -280,7 +316,7 @@ async def test_save_health_check_results_if_changed_status_unchanged_old():
|
|||
"""Test saving when status unchanged but last check is old (>1 hour)"""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.save_health_check_result = AsyncMock()
|
||||
|
||||
|
||||
model_results = {
|
||||
("model-123", "gpt-3.5-turbo"): {
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
|
|
@ -290,7 +326,7 @@ async def test_save_health_check_results_if_changed_status_unchanged_old():
|
|||
"error_message": None,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Latest check shows healthy, new result is healthy (status unchanged)
|
||||
# But checked >1 hour ago
|
||||
latest_checks_map = {
|
||||
|
|
@ -299,12 +335,16 @@ async def test_save_health_check_results_if_changed_status_unchanged_old():
|
|||
checked_at=datetime.now(timezone.utc) - timedelta(hours=2),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
start_time = 1234567890.0
|
||||
await _save_health_check_results_if_changed(
|
||||
mock_prisma, model_results, latest_checks_map, start_time, "background_health_check"
|
||||
mock_prisma,
|
||||
model_results,
|
||||
latest_checks_map,
|
||||
start_time,
|
||||
"background_health_check",
|
||||
)
|
||||
|
||||
|
||||
# Should save because last check is old (>1 hour)
|
||||
mock_prisma.save_health_check_result.assert_called_once()
|
||||
|
||||
|
|
@ -314,7 +354,7 @@ async def test_save_health_check_results_if_changed_no_previous_check():
|
|||
"""Test saving when there's no previous check"""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.save_health_check_result = AsyncMock()
|
||||
|
||||
|
||||
model_results = {
|
||||
("model-123", "gpt-3.5-turbo"): {
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
|
|
@ -324,15 +364,19 @@ async def test_save_health_check_results_if_changed_no_previous_check():
|
|||
"error_message": None,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# No previous check
|
||||
latest_checks_map = {}
|
||||
|
||||
|
||||
start_time = 1234567890.0
|
||||
await _save_health_check_results_if_changed(
|
||||
mock_prisma, model_results, latest_checks_map, start_time, "background_health_check"
|
||||
mock_prisma,
|
||||
model_results,
|
||||
latest_checks_map,
|
||||
start_time,
|
||||
"background_health_check",
|
||||
)
|
||||
|
||||
|
||||
# Should save because no previous check
|
||||
mock_prisma.save_health_check_result.assert_called_once()
|
||||
|
||||
|
|
@ -343,7 +387,7 @@ async def test_save_background_health_checks_to_db():
|
|||
mock_prisma = MagicMock()
|
||||
mock_prisma.save_health_check_result = AsyncMock()
|
||||
mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[])
|
||||
|
||||
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
|
|
@ -351,20 +395,25 @@ async def test_save_background_health_checks_to_db():
|
|||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
healthy_endpoints = [{"model": "gpt-3.5-turbo"}]
|
||||
unhealthy_endpoints = []
|
||||
|
||||
|
||||
start_time = 1234567890.0
|
||||
|
||||
|
||||
await _save_background_health_checks_to_db(
|
||||
mock_prisma, model_list, healthy_endpoints, unhealthy_endpoints, start_time, "background_health_check"
|
||||
mock_prisma,
|
||||
model_list,
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
start_time,
|
||||
"background_health_check",
|
||||
)
|
||||
|
||||
|
||||
# Should call get_all_latest_health_checks and save_health_check_result
|
||||
mock_prisma.get_all_latest_health_checks.assert_called_once()
|
||||
mock_prisma.save_health_check_result.assert_called_once()
|
||||
|
||||
|
||||
call_kwargs = mock_prisma.save_health_check_result.call_args[1]
|
||||
assert call_kwargs["model_name"] == "gpt-3.5-turbo"
|
||||
assert call_kwargs["model_id"] == "model-123"
|
||||
|
|
@ -385,8 +434,10 @@ async def test_save_background_health_checks_to_db_no_prisma():
|
|||
async def test_save_background_health_checks_to_db_exception_handling():
|
||||
"""Test exception handling in background health check save"""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.get_all_latest_health_checks = AsyncMock(side_effect=Exception("DB Error"))
|
||||
|
||||
mock_prisma.get_all_latest_health_checks = AsyncMock(
|
||||
side_effect=Exception("DB Error")
|
||||
)
|
||||
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
|
|
@ -394,12 +445,12 @@ async def test_save_background_health_checks_to_db_exception_handling():
|
|||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
# Should not raise exception, should handle gracefully
|
||||
await _save_background_health_checks_to_db(
|
||||
mock_prisma, model_list, [], [], 0.0, "background_health_check"
|
||||
)
|
||||
|
||||
|
||||
# Function should complete without raising
|
||||
|
||||
|
||||
|
|
@ -411,32 +462,34 @@ async def test_get_all_latest_health_checks_with_model_id(mock_prisma):
|
|||
mock_check1.model_id = "model-123"
|
||||
mock_check1.model_name = "gpt-3.5-turbo"
|
||||
mock_check1.checked_at = datetime.now(timezone.utc) - timedelta(minutes=10)
|
||||
|
||||
|
||||
mock_check2 = MagicMock()
|
||||
mock_check2.model_id = "model-456"
|
||||
mock_check2.model_name = "gpt-3.5-turbo"
|
||||
mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=5)
|
||||
|
||||
|
||||
mock_check3 = MagicMock()
|
||||
mock_check3.model_id = "model-123"
|
||||
mock_check3.model_name = "gpt-3.5-turbo"
|
||||
mock_check3.checked_at = datetime.now(timezone.utc) - timedelta(minutes=1) # Latest for model-123
|
||||
|
||||
mock_check3.checked_at = datetime.now(timezone.utc) - timedelta(
|
||||
minutes=1
|
||||
) # Latest for model-123
|
||||
|
||||
# Order by checked_at desc
|
||||
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(
|
||||
return_value=[mock_check3, mock_check2, mock_check1]
|
||||
)
|
||||
|
||||
|
||||
result = await mock_prisma.get_all_latest_health_checks()
|
||||
|
||||
|
||||
# Should return 2 unique models (by model_id)
|
||||
assert len(result) == 2
|
||||
|
||||
|
||||
# Should have latest check for each model_id
|
||||
model_ids = {check.model_id for check in result}
|
||||
assert "model-123" in model_ids
|
||||
assert "model-456" in model_ids
|
||||
|
||||
|
||||
# model-123 should have the latest check (1 minute ago)
|
||||
model123_check = next(c for c in result if c.model_id == "model-123")
|
||||
assert model123_check.checked_at == mock_check3.checked_at
|
||||
|
|
@ -449,18 +502,18 @@ async def test_get_all_latest_health_checks_without_model_id(mock_prisma):
|
|||
mock_check1.model_id = None
|
||||
mock_check1.model_name = "gpt-3.5-turbo"
|
||||
mock_check1.checked_at = datetime.now(timezone.utc) - timedelta(minutes=10)
|
||||
|
||||
|
||||
mock_check2 = MagicMock()
|
||||
mock_check2.model_id = None
|
||||
mock_check2.model_name = "gpt-3.5-turbo"
|
||||
mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=1) # Latest
|
||||
|
||||
|
||||
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(
|
||||
return_value=[mock_check2, mock_check1]
|
||||
)
|
||||
|
||||
|
||||
result = await mock_prisma.get_all_latest_health_checks()
|
||||
|
||||
|
||||
# Should return 1 unique model (by model_name)
|
||||
assert len(result) == 1
|
||||
assert result[0].model_name == "gpt-3.5-turbo"
|
||||
|
|
@ -480,7 +533,14 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c
|
|||
healthy = [{"model": "gpt-4"}]
|
||||
unhealthy = []
|
||||
|
||||
async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None, max_concurrency=None):
|
||||
async def mock_perform_health_check(
|
||||
model_list,
|
||||
model=None,
|
||||
cli_model=None,
|
||||
details=True,
|
||||
model_id=None,
|
||||
max_concurrency=None,
|
||||
):
|
||||
return healthy, unhealthy, {}
|
||||
|
||||
with patch(
|
||||
|
|
@ -505,5 +565,112 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c
|
|||
assert result["unhealthy_count"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_latest_health_checks_applies_ttl_filter(mock_prisma):
|
||||
"""get_all_latest_health_checks should pass a checked_at >= cutoff filter so
|
||||
it does not load the entire table into memory."""
|
||||
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
await mock_prisma.get_all_latest_health_checks()
|
||||
|
||||
call_kwargs = mock_prisma.db.litellm_healthchecktable.find_many.call_args[1]
|
||||
assert "where" in call_kwargs
|
||||
assert "checked_at" in call_kwargs["where"]
|
||||
assert "gte" in call_kwargs["where"]["checked_at"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_latest_health_checks_ttl_env_override(mock_prisma):
|
||||
"""HEALTH_CHECK_TTL_DAYS env var should control the cutoff window."""
|
||||
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
with patch.dict(os.environ, {"HEALTH_CHECK_TTL_DAYS": "3"}):
|
||||
await mock_prisma.get_all_latest_health_checks()
|
||||
|
||||
call_kwargs = mock_prisma.db.litellm_healthchecktable.find_many.call_args[1]
|
||||
cutoff = call_kwargs["where"]["checked_at"]["gte"]
|
||||
# The cutoff should be roughly 3 days ago (within a small tolerance)
|
||||
expected = datetime.utcnow() - timedelta(days=3)
|
||||
assert abs((expected - cutoff).total_seconds()) < 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_old_health_checks_calls_delete_many(mock_prisma):
|
||||
"""cleanup_old_health_checks should delete rows older than TTL via delete_many."""
|
||||
import types
|
||||
|
||||
mock_prisma.db.litellm_healthchecktable.delete_many = AsyncMock(return_value=42)
|
||||
mock_prisma.cleanup_old_health_checks = types.MethodType(
|
||||
PrismaClient.cleanup_old_health_checks, mock_prisma
|
||||
)
|
||||
|
||||
await mock_prisma.cleanup_old_health_checks()
|
||||
|
||||
call_kwargs = mock_prisma.db.litellm_healthchecktable.delete_many.call_args[1]
|
||||
assert "where" in call_kwargs
|
||||
assert "checked_at" in call_kwargs["where"]
|
||||
assert "lt" in call_kwargs["where"]["checked_at"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_health_check_result_triggers_cleanup_after_one_hour():
|
||||
"""cleanup_old_health_checks should be called when last cleanup was > 1 hour ago."""
|
||||
import types
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "x"})
|
||||
mock_client.db.litellm_healthchecktable.delete_many = AsyncMock(return_value=0)
|
||||
mock_client._health_check_last_cleanup_ts = 0.0 # never cleaned up
|
||||
# Bind helper methods properly so `self` is resolved correctly
|
||||
mock_client._validate_response_time = types.MethodType(
|
||||
PrismaClient._validate_response_time, mock_client
|
||||
)
|
||||
mock_client._clean_details = types.MethodType(
|
||||
PrismaClient._clean_details, mock_client
|
||||
)
|
||||
mock_client.save_health_check_result = types.MethodType(
|
||||
PrismaClient.save_health_check_result, mock_client
|
||||
)
|
||||
mock_client.cleanup_old_health_checks = types.MethodType(
|
||||
PrismaClient.cleanup_old_health_checks, mock_client
|
||||
)
|
||||
|
||||
await mock_client.save_health_check_result(
|
||||
model_name="gpt-4", status="healthy", healthy_count=1, unhealthy_count=0
|
||||
)
|
||||
|
||||
mock_client.db.litellm_healthchecktable.delete_many.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_health_check_result_skips_cleanup_within_one_hour():
|
||||
"""cleanup_old_health_checks should NOT be called if last cleanup was < 1 hour ago."""
|
||||
import types
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "x"})
|
||||
mock_client.db.litellm_healthchecktable.delete_many = AsyncMock(return_value=0)
|
||||
mock_client._health_check_last_cleanup_ts = time.time() # just ran
|
||||
# Bind helper methods properly so `self` is resolved correctly
|
||||
mock_client._validate_response_time = types.MethodType(
|
||||
PrismaClient._validate_response_time, mock_client
|
||||
)
|
||||
mock_client._clean_details = types.MethodType(
|
||||
PrismaClient._clean_details, mock_client
|
||||
)
|
||||
mock_client.save_health_check_result = types.MethodType(
|
||||
PrismaClient.save_health_check_result, mock_client
|
||||
)
|
||||
mock_client.cleanup_old_health_checks = types.MethodType(
|
||||
PrismaClient.cleanup_old_health_checks, mock_client
|
||||
)
|
||||
|
||||
await mock_client.save_health_check_result(
|
||||
model_name="gpt-4", status="healthy", healthy_count=1, unhealthy_count=0
|
||||
)
|
||||
|
||||
mock_client.db.litellm_healthchecktable.delete_many.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
pytest.main([__file__])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue