scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 133

Merge pull request #9569 from BerriAI/litellm_fix_db_unavailable

Browse source 
[Reliability Fix] - Allow Pods to startup + passing /health/readiness when `allow_requests_on_db_unavailable: True` and DB is down
This commit is contained in:
Ishaan Jaff 2025-03-26 21:18:52 -07:00 • committed by GitHub
commit 7097ce544e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 370 additions and 151 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
docs/my-website/docs/proxy/prod.md
View file

@ -114,6 +114,8 @@ When `allow_requests_on_db_unavailable` is set to `true`, LiteLLM will handle er
|---------------|-------------------|----------------|
| Prisma Errors | ✅ Request will be allowed | Covers issues like DB connection resets or rejections from the DB via Prisma, the ORM used by LiteLLM. |
| Httpx Errors | ✅ Request will be allowed | Occurs when the database is unreachable, allowing the request to proceed despite the DB outage. |
| Pod Startup Behavior | ✅ Pods start regardless | LiteLLM Pods will start even if the database is down or unreachable, ensuring higher uptime guarantees for deployments. |
| Health/Readiness Check | ✅ Always returns 200 OK | The /health/readiness endpoint returns a 200 OK status to ensure that pods remain operational even when the database is unavailable.
| LiteLLM Budget Errors or Model Errors | ❌ Request will be blocked | Triggered when the DB is reachable but the authentication token is invalid, lacks access, or exceeds budget limits. |

38
litellm/proxy/auth/auth_exception_handler.py
View file

@ -9,13 +9,9 @@ from fastapi import HTTPException, Request, status
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
ProxyErrorTypes,
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import _get_request_ip_address
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.types.services import ServiceTypes
if TYPE_CHECKING:
@ -58,8 +54,8 @@ class UserAPIKeyAuthExceptionHandler:
)
if (
UserAPIKeyAuthExceptionHandler.should_allow_request_on_db_unavailable()
and UserAPIKeyAuthExceptionHandler.is_database_connection_error(e)
PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
and PrismaDBExceptionHandler.is_database_connection_error(e)
):
# log this as a DB failure on prometheus
proxy_logging_obj.service_logging_obj.service_failure_hook(
@ -125,29 +121,3 @@ class UserAPIKeyAuthExceptionHandler:
param=getattr(e, "param", "None"),
code=status.HTTP_401_UNAUTHORIZED,
)
@staticmethod
def should_allow_request_on_db_unavailable() -> bool:
"""
Returns True if the request should be allowed to proceed despite the DB connection error
"""
from litellm.proxy.proxy_server import general_settings
if general_settings.get("allow_requests_on_db_unavailable", False) is True:
return True
return False
@staticmethod
def is_database_connection_error(e: Exception) -> bool:
"""
Returns True if the exception is from a database outage / connection error
"""
import prisma
if isinstance(e, DB_CONNECTION_ERROR_TYPES):
return True
if isinstance(e, prisma.errors.PrismaError):
return True
if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection:
return True
return False

61
litellm/proxy/db/exception_handler.py Normal file
View file

@ -0,0 +1,61 @@
from typing import Union
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
ProxyErrorTypes,
ProxyException,
)
from litellm.secret_managers.main import str_to_bool
class PrismaDBExceptionHandler:
"""
Class to handle DB Exceptions or Connection Errors
"""
@staticmethod
def should_allow_request_on_db_unavailable() -> bool:
"""
Returns True if the request should be allowed to proceed despite the DB connection error
"""
from litellm.proxy.proxy_server import general_settings
_allow_requests_on_db_unavailable: Union[bool, str] = general_settings.get(
"allow_requests_on_db_unavailable", False
)
if isinstance(_allow_requests_on_db_unavailable, bool):
return _allow_requests_on_db_unavailable
if str_to_bool(_allow_requests_on_db_unavailable) is True:
return True
return False
@staticmethod
def is_database_connection_error(e: Exception) -> bool:
"""
Returns True if the exception is from a database outage / connection error
"""
import prisma
if isinstance(e, DB_CONNECTION_ERROR_TYPES):
return True
if isinstance(e, prisma.errors.PrismaError):
return True
if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection:
return True
return False
@staticmethod
def handle_db_exception(e: Exception):
"""
Primary handler for `allow_requests_on_db_unavailable` flag. Decides whether to raise a DB Exception or not based on the flag.
- If exception is a DB Connection Error, and `allow_requests_on_db_unavailable` is True,
- Do not raise an exception, return None
- Else, raise the exception
"""
if (
PrismaDBExceptionHandler.is_database_connection_error(e)
and PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
):
return None
raise e

24
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -20,6 +20,7 @@ from litellm.proxy._types import (
WebhookEvent,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.health_check import (
_clean_endpoint_data,
_update_litellm_params_for_health_check,
@ -381,20 +382,23 @@ async def _db_health_readiness_check():
global db_health_cache
# Note - Intentionally don't try/except this so it raises an exception when it fails
try:
# if timedelta is less than 2 minutes return DB Status
time_diff = datetime.now() - db_health_cache["last_updated"]
if db_health_cache["status"] != "unknown" and time_diff < timedelta(minutes=2):
return db_health_cache
# if timedelta is less than 2 minutes return DB Status
time_diff = datetime.now() - db_health_cache["last_updated"]
if db_health_cache["status"] != "unknown" and time_diff < timedelta(minutes=2):
if prisma_client is None:
db_health_cache = {"status": "disconnected", "last_updated": datetime.now()}
return db_health_cache
await prisma_client.health_check()
db_health_cache = {"status": "connected", "last_updated": datetime.now()}
return db_health_cache
if prisma_client is None:
db_health_cache = {"status": "disconnected", "last_updated": datetime.now()}
except Exception as e:
PrismaDBExceptionHandler.handle_db_exception(e)
return db_health_cache
await prisma_client.health_check()
db_health_cache = {"status": "connected", "last_updated": datetime.now()}
return db_health_cache
@router.get(
"/settings",

75
litellm/proxy/proxy_server.py
View file

@ -176,6 +176,7 @@ from litellm.proxy.common_utils.proxy_state import ProxyState
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router
from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config
from litellm.proxy.guardrails.guardrail_endpoints import router as guardrails_router
@ -453,18 +454,6 @@ async def proxy_startup_event(app: FastAPI):
import json
init_verbose_loggers()
### LOAD MASTER KEY ###
# check if master key set in environment - load from there
master_key = get_secret("LITELLM_MASTER_KEY", None) # type: ignore
# check if DATABASE_URL in environment - load from there
if prisma_client is None:
_db_url: Optional[str] = get_secret("DATABASE_URL", None) # type: ignore
prisma_client = await ProxyStartupEvent._setup_prisma_client(
database_url=_db_url,
proxy_logging_obj=proxy_logging_obj,
user_api_key_cache=user_api_key_cache,
)
## CHECK PREMIUM USER
verbose_proxy_logger.debug(
"litellm.proxy.proxy_server.py::startup() - CHECKING PREMIUM USER - {}".format(
@ -521,6 +510,18 @@ async def proxy_startup_event(app: FastAPI):
if isinstance(worker_config, dict):
await initialize(**worker_config)
### LOAD MASTER KEY ###
# check if master key set in environment - load from there
master_key = get_secret("LITELLM_MASTER_KEY", None) # type: ignore
# check if DATABASE_URL in environment - load from there
if prisma_client is None:
_db_url: Optional[str] = get_secret("DATABASE_URL", None) # type: ignore
prisma_client = await ProxyStartupEvent._setup_prisma_client(
database_url=_db_url,
proxy_logging_obj=proxy_logging_obj,
user_api_key_cache=user_api_key_cache,
)
ProxyStartupEvent._initialize_startup_logging(
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
@ -3362,33 +3363,37 @@ class ProxyStartupEvent:
- Sets up prisma client
- Adds necessary views to proxy
"""
prisma_client: Optional[PrismaClient] = None
if database_url is not None:
try:
prisma_client = PrismaClient(
database_url=database_url, proxy_logging_obj=proxy_logging_obj
)
except Exception as e:
raise e
try:
prisma_client: Optional[PrismaClient] = None
if database_url is not None:
try:
prisma_client = PrismaClient(
database_url=database_url, proxy_logging_obj=proxy_logging_obj
)
except Exception as e:
raise e
await prisma_client.connect()
await prisma_client.connect()
## Add necessary views to proxy ##
asyncio.create_task(
prisma_client.check_view_exists()
) # check if all necessary views exist. Don't block execution
## Add necessary views to proxy ##
asyncio.create_task(
prisma_client.check_view_exists()
) # check if all necessary views exist. Don't block execution
asyncio.create_task(
prisma_client._set_spend_logs_row_count_in_proxy_state()
) # set the spend logs row count in proxy state. Don't block execution
asyncio.create_task(
prisma_client._set_spend_logs_row_count_in_proxy_state()
) # set the spend logs row count in proxy state. Don't block execution
# run a health check to ensure the DB is ready
if (
get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False)
is not True
):
await prisma_client.health_check()
return prisma_client
# run a health check to ensure the DB is ready
if (
get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False)
is not True
):
await prisma_client.health_check()
return prisma_client
except Exception as e:
PrismaDBExceptionHandler.handle_db_exception(e)
return None
@classmethod
def _init_dd_tracer(cls):

72
tests/litellm/proxy/auth/test_auth_exception_handler.py
View file

@ -29,78 +29,6 @@ from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler
# Test is_database_connection_error method
@pytest.mark.parametrize(
"prisma_error",
[
PrismaError(),
DataError(data={"user_facing_error": {"meta": {"table": "test_table"}}}),
UniqueViolationError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
ForeignKeyViolationError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
MissingRequiredValueError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
RawQueryError(data={"user_facing_error": {"meta": {"table": "test_table"}}}),
TableNotFoundError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
RecordNotFoundError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
HTTPClientClosedError(),
ClientNotConnectedError(),
],
)
def test_is_database_connection_error_prisma_errors(prisma_error):
"""
Test that all Prisma errors are considered database connection errors
"""
handler = UserAPIKeyAuthExceptionHandler()
assert handler.is_database_connection_error(prisma_error) == True
def test_is_database_connection_generic_errors():
"""
Test non-Prisma error cases for database connection checking
"""
handler = UserAPIKeyAuthExceptionHandler()
# Test with ProxyException (DB connection)
db_proxy_exception = ProxyException(
message="DB Connection Error",
type=ProxyErrorTypes.no_db_connection,
param="test-param",
)
assert handler.is_database_connection_error(db_proxy_exception) == True
# Test with non-DB error
regular_exception = Exception("Regular error")
assert handler.is_database_connection_error(regular_exception) == False
# Test should_allow_request_on_db_unavailable method
@patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
)
def test_should_allow_request_on_db_unavailable_true():
handler = UserAPIKeyAuthExceptionHandler()
assert handler.should_allow_request_on_db_unavailable() == True
@patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": False},
)
def test_should_allow_request_on_db_unavailable_false():
handler = UserAPIKeyAuthExceptionHandler()
assert handler.should_allow_request_on_db_unavailable() == False
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prisma_error",

148
tests/litellm/proxy/db/test_exception_handler.py Normal file
View file

@ -0,0 +1,148 @@
import asyncio
import json
import os
import sys
from unittest.mock import MagicMock, patch
import pytest
from fastapi import HTTPException, Request, status
from prisma import errors as prisma_errors
from prisma.errors import (
ClientNotConnectedError,
DataError,
ForeignKeyViolationError,
HTTPClientClosedError,
MissingRequiredValueError,
PrismaError,
RawQueryError,
RecordNotFoundError,
TableNotFoundError,
UniqueViolationError,
)
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
# Test is_database_connection_error method
@pytest.mark.parametrize(
"prisma_error",
[
PrismaError(),
DataError(data={"user_facing_error": {"meta": {"table": "test_table"}}}),
UniqueViolationError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
ForeignKeyViolationError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
MissingRequiredValueError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
RawQueryError(data={"user_facing_error": {"meta": {"table": "test_table"}}}),
TableNotFoundError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
RecordNotFoundError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
HTTPClientClosedError(),
ClientNotConnectedError(),
],
)
def test_is_database_connection_error_prisma_errors(prisma_error):
"""
Test that all Prisma errors are considered database connection errors
"""
assert PrismaDBExceptionHandler.is_database_connection_error(prisma_error) == True
def test_is_database_connection_generic_errors():
"""
Test non-Prisma error cases for database connection checking
"""
assert (
PrismaDBExceptionHandler.is_database_connection_error(
Exception("Regular error")
)
== False
)
# Test with ProxyException (DB connection)
db_proxy_exception = ProxyException(
message="DB Connection Error",
type=ProxyErrorTypes.no_db_connection,
param="test-param",
)
assert (
PrismaDBExceptionHandler.is_database_connection_error(db_proxy_exception)
== True
)
# Test with non-DB error
regular_exception = Exception("Regular error")
assert (
PrismaDBExceptionHandler.is_database_connection_error(regular_exception)
== False
)
# Test should_allow_request_on_db_unavailable method
@patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
)
def test_should_allow_request_on_db_unavailable_true():
assert PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() == True
@patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": False},
)
def test_should_allow_request_on_db_unavailable_false():
assert PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() == False
@patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
)
def test_handle_db_exception_with_connection_error():
"""
Test that DB connection errors are handled gracefully when allow_requests_on_db_unavailable is True
"""
db_error = ClientNotConnectedError()
result = PrismaDBExceptionHandler.handle_db_exception(db_error)
assert result is None
@patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": False},
)
def test_handle_db_exception_raises_error():
"""
Test that DB connection errors are raised when allow_requests_on_db_unavailable is False
"""
db_error = ClientNotConnectedError()
with pytest.raises(ClientNotConnectedError):
PrismaDBExceptionHandler.handle_db_exception(db_error)
def test_handle_db_exception_with_non_db_error():
"""
Test that non-DB errors are always raised regardless of allow_requests_on_db_unavailable setting
"""
regular_error = litellm.BudgetExceededError(
current_cost=10,
max_budget=10,
)
with pytest.raises(litellm.BudgetExceededError):
PrismaDBExceptionHandler.handle_db_exception(regular_error)

101
tests/litellm/proxy/health_endpoints/test_health_endpoints.py Normal file
View file

@ -0,0 +1,101 @@
import asyncio
import json
import os
import sys
from datetime import datetime, timedelta
from unittest.mock import MagicMock, patch
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
import pytest
from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, PrismaError
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.health_endpoints._health_endpoints import (
_db_health_readiness_check,
db_health_cache,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prisma_error",
[
PrismaError(),
ClientNotConnectedError(),
HTTPClientClosedError(),
],
)
async def test_db_health_readiness_check_with_prisma_error(prisma_error):
"""
Test that when prisma_client.health_check() raises a PrismaError and
allow_requests_on_db_unavailable is True, the function should not raise an error
and return the cached health status.
"""
# Mock the prisma client
mock_prisma_client = MagicMock()
mock_prisma_client.health_check.side_effect = prisma_error
# Reset the health cache to a known state
global db_health_cache
db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(minutes=5),
}
# Patch the imports and general_settings
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
):
# Call the function
result = await _db_health_readiness_check()
# Verify that the function called health_check
mock_prisma_client.health_check.assert_called_once()
# Verify that the function returned the cache
assert result is not None
assert result["status"] == "unknown" # Should retain the status from the cache
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prisma_error",
[
PrismaError(),
ClientNotConnectedError(),
HTTPClientClosedError(),
],
)
async def test_db_health_readiness_check_with_error_and_flag_off(prisma_error):
"""
Test that when prisma_client.health_check() raises a DB error but
allow_requests_on_db_unavailable is False, the exception should be raised.
"""
# Mock the prisma client
mock_prisma_client = MagicMock()
mock_prisma_client.health_check.side_effect = prisma_error
# Reset the health cache
global db_health_cache
db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(minutes=5),
}
# Patch the imports and general_settings where the flag is False
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": False},
):
# The function should raise the exception
with pytest.raises(Exception) as excinfo:
await _db_health_readiness_check()
# Verify that the raised exception is the same
assert excinfo.value == prisma_error