mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(proxy): warn at startup when max_budget is set but no database is connected (#36041)
* warn at startup when a proxy-wide budget is set but no DB is connected litellm.max_budget is only enforced via DB-loaded global spend, so a DB-less proxy silently ignores it. Log a one-time startup warning. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): inject max_budget into DB-less budget warning Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover DB-less budget warning startup call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): pin DB-less budget warning call site Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): stabilize budget warning call-site pin Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: tin <tin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
f6587faef5
commit
b7749f67f1
2 changed files with 66 additions and 0 deletions
|
|
@ -1111,6 +1111,10 @@ async def proxy_startup_event(app: FastAPI):
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
)
|
||||
ProxyStartupEvent._warn_budget_without_db(
|
||||
max_budget=litellm.max_budget,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
### START BATCH WRITING DB + CHECKING NEW MODELS###
|
||||
if prisma_client is not None:
|
||||
|
|
@ -7825,6 +7829,19 @@ def giveup(e):
|
|||
|
||||
|
||||
class ProxyStartupEvent:
|
||||
@staticmethod
|
||||
def _warn_budget_without_db(max_budget: float | None, prisma_client: PrismaClient | None) -> None:
|
||||
if prisma_client is not None or not max_budget or max_budget <= 0:
|
||||
return
|
||||
|
||||
verbose_proxy_logger.warning(
|
||||
"A proxy-wide budget (litellm.max_budget=%s) is configured but no database is connected, "
|
||||
"so the budget will NOT be enforced and requests will never be blocked. Set DATABASE_URL or "
|
||||
"general_settings.database_url and restart. Redis and fail_closed_budget_enforcement do not "
|
||||
"cover the proxy-wide budget because there is no global spend counter; Redis alone is not a substitute.",
|
||||
max_budget,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _initialize_startup_logging(
|
||||
cls,
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import List, Optional, Union
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -31,6 +32,7 @@ from typing_extensions import TypedDict
|
|||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.proxy_server import (
|
||||
ProxyStartupEvent,
|
||||
_initialize_shared_aiohttp_session,
|
||||
_resolve_pydantic_type,
|
||||
_resolve_typed_dict_type,
|
||||
|
|
@ -728,3 +730,50 @@ def test_otel_global_provider_published_after_callback_init():
|
|||
"preset logger will not exist yet and a second generic logger will own "
|
||||
"the global provider, orphaning gen-ai spans"
|
||||
)
|
||||
|
||||
|
||||
def test_startup_warns_for_global_budget_without_database(caplog):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=None)
|
||||
|
||||
assert "litellm.max_budget=100.0" in caplog.text
|
||||
assert "will NOT be enforced" in caplog.text
|
||||
assert "requests will never be blocked" in caplog.text
|
||||
|
||||
|
||||
def test_startup_does_not_warn_for_global_budget_with_database(caplog):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=MagicMock())
|
||||
|
||||
assert "litellm.max_budget" not in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("max_budget", [0, None])
|
||||
def test_startup_does_not_warn_without_global_budget(caplog, max_budget):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
ProxyStartupEvent._warn_budget_without_db(max_budget=max_budget, prisma_client=None)
|
||||
|
||||
assert "litellm.max_budget" not in caplog.text
|
||||
|
||||
|
||||
def test_proxy_startup_event_warns_for_global_budget_without_database():
|
||||
"""Pin the lifespan call that prevents silent DB-less budgets.
|
||||
|
||||
The call must follow Prisma setup so DB-backed deployments do not false-positive.
|
||||
Direct ``_warn_budget_without_db`` tests cover the warning behavior itself.
|
||||
"""
|
||||
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
|
||||
source = inspect.getsource(wrapped)
|
||||
budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:")
|
||||
warn_pos = source.find("_warn_budget_without_db(")
|
||||
next_startup_section_pos = source.find(
|
||||
"await ProxyStartupEvent.initialize_scheduled_background_jobs(",
|
||||
budget_check_pos,
|
||||
)
|
||||
|
||||
assert budget_check_pos != -1, "global budget startup block not found"
|
||||
assert warn_pos != -1, "DB-less budget warning call not found"
|
||||
assert next_startup_section_pos != -1, "startup section after budget block not found"
|
||||
assert budget_check_pos < warn_pos < next_startup_section_pos, (
|
||||
"DB-less budget warning must run after Prisma setup and the DB-backed budget block"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue