phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

allow proxy to startup on DB unavailable

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-26 19:50:57 -07:00
parent 497570b2a6
commit 88ef97b9d1
2 changed files with 50 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

6
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,7 +382,7 @@ 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):
@ -394,6 +395,9 @@ async def _db_health_readiness_check():
await prisma_client.health_check()
db_health_cache = {"status": "connected", "last_updated": datetime.now()}
return db_health_cache
except Exception as e:
PrismaDBExceptionHandler.handle_db_exception(e)
return db_health_cache
@router.get(

22
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
@ -456,15 +457,6 @@ async def proxy_startup_event(app: FastAPI):
### 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(
@ -527,6 +519,15 @@ async def proxy_startup_event(app: FastAPI):
redis_usage_cache=redis_usage_cache,
)
# 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,
)
## JWT AUTH ##
ProxyStartupEvent._initialize_jwt_auth(
general_settings=general_settings,
@ -3362,6 +3363,7 @@ class ProxyStartupEvent:
- Sets up prisma client
- Adds necessary views to proxy
"""
try:
prisma_client: Optional[PrismaClient] = None
if database_url is not None:
try:
@ -3389,6 +3391,8 @@ class ProxyStartupEvent:
):
await prisma_client.health_check()
return prisma_client
except Exception as e:
PrismaDBExceptionHandler.handle_db_exception(e)
@classmethod
def _init_dd_tracer(cls):