Merge pull request #41974 from BerriAI/litellm_fix_startup_view_creation_race

fix(proxy): wait for the spend-log table before creating startup views
This commit is contained in:
Mateo Wang 2026-09-19 20:33:49 -07:00 • committed by GitHub
commit 9d7f77988a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 411 additions and 38 deletions

View file

@ -1510,6 +1510,12 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
except Exception as e:
verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e)
if prisma_client is not None and hasattr(prisma_client, "stop_view_setup_task"):
try:
await prisma_client.stop_view_setup_task()
except Exception as e:
verbose_proxy_logger.error("Error stopping the spend view setup task: %s", e)
await _drain_spend_event_producer_on_shutdown()
await flush_spend_counters_on_shutdown()
@ -10808,14 +10814,7 @@ class ProxyStartupEvent:
if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"):
await prisma_client.db.start_token_refresh_task()
## 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
prisma_client.start_view_setup_task()
if hasattr(prisma_client, "start_db_health_watchdog_task"):
await prisma_client.start_db_health_watchdog_task()

View file

@ -38,6 +38,7 @@ from typing import (
Literal,
Optional,
Protocol,
TypeAlias,
TypeVar,
Union,
cast,
@ -269,6 +270,15 @@ class _RelTuplesRow(TypedDict):
reltuples: ReadOnly[int]
_VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0
_VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0
_VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"'
_VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...])
_ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"]
_ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception
class _EndUserBatchTable(Protocol):
def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
@ -4297,6 +4307,7 @@ class PrismaClient:
self.db = writer_wrapper # Client to connect to Prisma db
self._db_reconnect_lock = asyncio.Lock()
self._db_health_watchdog_task: asyncio.Task | None = None
self._view_setup_task: asyncio.Task[_ViewSetupOutcome] | None = None
self._db_last_reconnect_attempt_ts: float = 0.0
self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15")))
self._db_read_only_recreate_ts: float = 0.0
@ -6343,6 +6354,71 @@ class PrismaClient:
self._db_health_watchdog_task = None
verbose_proxy_logger.info("Stopped Prisma DB health watchdog")
def start_view_setup_task(self) -> None:
if self._view_setup_task is not None:
return
self._view_setup_task = asyncio.create_task(self._run_view_setup())
async def stop_view_setup_task(self) -> None:
if self._view_setup_task is None:
return
self._view_setup_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self._view_setup_task
self._view_setup_task = None
async def _run_view_setup(
self,
poll_interval_seconds: float = _VIEW_SETUP_POLL_INTERVAL_SECONDS,
deadline_seconds: float = _VIEW_SETUP_DEADLINE_SECONDS,
) -> _ViewSetupOutcome:
deadline: Final = time.monotonic() + deadline_seconds
while True:
if (attempt := await self._attempt_view_setup()) == "ready":
return "ready"
if time.monotonic() >= deadline:
self._log_view_setup_timeout(attempt, deadline_seconds)
return "timed_out"
await asyncio.sleep(poll_interval_seconds)
async def _attempt_view_setup(self) -> _ViewSetupAttempt:
try:
if not await self._view_setup_gate_table_present():
verbose_proxy_logger.debug(
"Waiting for table %s before creating the spend views", _VIEW_SETUP_GATE_TABLE
)
return "table_missing"
await self._set_spend_logs_row_count_in_proxy_state()
await self.check_view_exists()
return "ready"
except Exception as e:
verbose_proxy_logger.warning("Spend view setup attempt failed, retrying until the schema settles: %s", e)
return e
def _log_view_setup_timeout(
self, last_attempt: Literal["table_missing"] | Exception, deadline_seconds: float
) -> None:
if isinstance(last_attempt, Exception):
verbose_proxy_logger.error(
"Gave up creating the spend views after %ss; the last attempt failed with: %s. "
"Fix that error and restart the proxy.",
deadline_seconds,
last_attempt,
)
return
verbose_proxy_logger.error(
"Gave up creating the spend views: table %s did not appear within %ss. "
"Run the database migrations against this database and restart the proxy.",
_VIEW_SETUP_GATE_TABLE,
deadline_seconds,
)
async def _view_setup_gate_table_present(self) -> bool:
rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python(
await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE)
)
return rows[0]["present"]
async def _db_health_watchdog_loop(self) -> None:
while True:
try:

View file

@ -359,6 +359,17 @@ class SpendClient:
def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult:
return self.proxy.transport.probe(path, params=params)
def probe_until_healthy(self, path: str, *, params: DateRangeParams) -> ProbeResult:
outcome: Final = await_converged(
lambda: self.probe(path, params=params),
converged=lambda result: result.healthy,
timeout=self.proxy.poll_timeout,
interval=self.proxy.poll_interval,
now=time.monotonic,
sleep=time.sleep,
)
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
def create_user(self, *, email: str, role: UserRole, user_id: str) -> str:
return unwrap(
self.proxy.transport.post(

View file

@ -17,9 +17,11 @@ fast: no batch-write wait, no provider calls.
"""
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
from e2e_http import ProbeResult
from models import DateRangeParams
from spend_e2e_client import SpendClient
@ -72,15 +74,10 @@ SPEND_ROUTES = (
_SPEND_PREFIXES = ("/spend", "/global/spend", "/global/activity")
_MISSING_VIEW_SKIP = pytest.mark.skip(
reason=(
"LIT-5211: on a fresh database the proxy's startup view creation can lose the race "
"against schema migrations, leaving MonthlyGlobalSpend/DailyTagSpend/Last30d* views "
"missing and these routes 500ing until the views exist"
)
)
_VIEW_BACKED_ROUTES = frozenset(
# Served from the MonthlyGlobalSpend / DailyTagSpend / Last30d* views, which the
# proxy creates in the background once the schema migrations have landed, so on a
# fresh database they can 500 for a while after the proxy starts serving.
_VIEW_BACKED_ROUTES: Final = frozenset(
(
"/global/spend",
"/global/spend/keys",
@ -98,15 +95,15 @@ def _date_range() -> DateRangeParams:
return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat())
@pytest.mark.parametrize(
"route",
tuple(
pytest.param(route, marks=_MISSING_VIEW_SKIP) if route in _VIEW_BACKED_ROUTES else route
for route in SPEND_ROUTES
),
)
def _probe(client: SpendClient, route: str) -> ProbeResult:
if route in _VIEW_BACKED_ROUTES:
return client.probe_until_healthy(route, params=_date_range())
return client.probe(route, params=_date_range())
@pytest.mark.parametrize("route", SPEND_ROUTES)
def test_spend_route_responsive(client: SpendClient, route: str) -> None:
result = client.probe(route, params=_date_range())
result = _probe(client, route)
print(f"{route} -> {result.status_code}\n{result.body[:600]}")
assert result.healthy, f"{route} -> {result.status_code}\n{result.body[:600]}"

View file

@ -2382,7 +2382,7 @@ async def test_proxy_model_group_info_rerank(prisma_client): # noqa: F811 # py
@pytest.mark.asyncio
async def test_proxy_server_prisma_setup():
from litellm.proxy.proxy_server import ProxyStartupEvent, proxy_state
from litellm.proxy.proxy_server import ProxyStartupEvent
from litellm.proxy.utils import ProxyLogging
from litellm.caching import DualCache
@ -2393,35 +2393,28 @@ async def test_proxy_server_prisma_setup():
) as mock_prisma_client:
mock_client = mock_prisma_client.return_value # This is the mocked instance
mock_client.connect = AsyncMock() # Mock the connect method
mock_client.check_view_exists = AsyncMock() # Mock the check_view_exists method
mock_client.start_view_setup_task = MagicMock()
mock_client.health_check = AsyncMock() # Mock the health_check method
mock_client._set_spend_logs_row_count_in_proxy_state = (
AsyncMock()
) # Mock the _set_spend_logs_row_count_in_proxy_state method
mock_client.start_db_health_watchdog_task = AsyncMock()
# Mock the db attribute with start_token_refresh_task for RDS IAM token refresh
mock_db = MagicMock()
mock_db.start_token_refresh_task = AsyncMock()
mock_client.db = mock_db
await ProxyStartupEvent._setup_prisma_client(
prisma_client = await ProxyStartupEvent._setup_prisma_client(
database_url=os.getenv("DATABASE_URL"),
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
user_api_key_cache=user_api_key_cache,
)
# Verify our mocked methods were called
assert prisma_client is mock_client
mock_client.connect.assert_called_once()
mock_client.check_view_exists.assert_called_once()
mock_client.start_view_setup_task.assert_called_once()
# Note: This is REALLY IMPORTANT to check that the health check is called
# This is how we ensure the DB is ready before proceeding
mock_client.health_check.assert_called_once()
# check that the spend logs row count is set in proxy state
mock_client._set_spend_logs_row_count_in_proxy_state.assert_called_once()
assert proxy_state.get_proxy_state_variable("spend_logs_row_count") is not None
@pytest.mark.asyncio
async def test_proxy_server_prisma_setup_invalid_db(monkeypatch):

View file

@ -1810,6 +1810,37 @@ async def test_proxy_startup_boots_an_unsafe_master_key_under_the_override(monke
assert announced == []
class _ShutdownAwarePrisma(MockPrisma):
def __init__(self):
super().__init__()
self.stop_view_setup_task = AsyncMock()
@pytest.mark.asyncio
async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path):
import yaml
from fastapi import FastAPI
from litellm.proxy.proxy_server import proxy_startup_event
fake_prisma = _ShutdownAwarePrisma()
config_path = tmp_path / "config.yaml"
with open(config_path, "w") as f:
yaml.dump({"general_settings": {"master_key": "sk-12345"}}, f)
monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path))
monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma)
monkeypatch.setattr(proxy_server_module, "store_model_in_db", False)
async with proxy_startup_event(FastAPI()):
stopped_while_serving = fake_prisma.stop_view_setup_task.await_count
actual = {
"stopped_while_serving": stopped_while_serving,
"stopped_after_shutdown": fake_prisma.stop_view_setup_task.await_count,
}
assert actual == {"stopped_while_serving": 0, "stopped_after_shutdown": 1}
def test_team_info_masking():
"""
Test that sensitive team information is properly masked
@ -13608,6 +13639,7 @@ def _mock_startup_prisma_client(health_check_error=None, connect_error=None):
client.db.start_token_refresh_task = AsyncMock()
client.check_view_exists = AsyncMock()
client._set_spend_logs_row_count_in_proxy_state = AsyncMock()
client.start_view_setup_task = MagicMock()
client.start_db_health_watchdog_task = AsyncMock()
client.health_check = AsyncMock(side_effect=health_check_error)
return client
@ -13669,13 +13701,35 @@ async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_ch
mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out"))
call_order = MagicMock()
call_order.attach_mock(mock_client.start_view_setup_task, "view_setup")
call_order.attach_mock(mock_client.start_db_health_watchdog_task, "watchdog")
call_order.attach_mock(mock_client.health_check, "health_check")
await _run_setup_prisma_client(mock_client)
assert mock_client.start_db_health_watchdog_task.await_count == 1
assert [call[0] for call in call_order.mock_calls] == ["watchdog", "health_check"]
assert [call[0] for call in call_order.mock_calls] == ["view_setup", "watchdog", "health_check"]
@pytest.mark.asyncio
async def test_setup_prisma_client_hands_view_creation_to_the_held_task(monkeypatch):
monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True")
mock_client = _mock_startup_prisma_client()
result = await _run_setup_prisma_client(mock_client)
actual = {
"result": result,
"view_setup_started": mock_client.start_view_setup_task.call_count,
"direct_view_creation": mock_client.check_view_exists.await_count,
"direct_row_count": mock_client._set_spend_logs_row_count_in_proxy_state.await_count,
}
assert actual == {
"result": mock_client,
"view_setup_started": 1,
"direct_view_creation": 0,
"direct_row_count": 0,
}
@pytest.mark.asyncio

View file

@ -5,11 +5,15 @@ Symbols pinned here:
- ``PrismaClient.writer_db``
- ``PrismaClient.connect``
- ``PrismaClient.disconnect``
- ``PrismaClient.start_view_setup_task``
- ``PrismaClient.stop_view_setup_task``
- ``PrismaClient._run_view_setup``
"""
from __future__ import annotations
import asyncio
import logging
from typing import Any
from unittest.mock import AsyncMock, MagicMock
@ -17,6 +21,27 @@ import pytest
from litellm.proxy.utils import PrismaClient
_PROBE_SQL = "SELECT to_regclass($1) IS NOT NULL AS present"
def _absent() -> list[dict[str, bool]]:
return [{"present": False}]
def _present() -> list[dict[str, bool]]:
return [{"present": True}]
def _wire_view_setup(prisma_client: PrismaClient, probe: AsyncMock) -> MagicMock:
prisma_client.db.query_raw = probe
prisma_client.check_view_exists = AsyncMock()
prisma_client._set_spend_logs_row_count_in_proxy_state = AsyncMock()
call_order = MagicMock()
call_order.attach_mock(probe, "probe")
call_order.attach_mock(prisma_client.check_view_exists, "views")
call_order.attach_mock(prisma_client._set_spend_logs_row_count_in_proxy_state, "row_count")
return call_order
@pytest.mark.asyncio
async def test_prismaclient_init_wires_default_config(
@ -205,3 +230,221 @@ async def test_disconnect_raises_when_underlying_fails(
prisma_client.db.disconnect = AsyncMock(side_effect=RuntimeError("disconnect boom"))
with pytest.raises(RuntimeError, match="disconnect boom"):
await prisma_client.disconnect()
@pytest.mark.asyncio
async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views(prisma_client: PrismaClient) -> None:
probe = AsyncMock(side_effect=[_absent(), _absent(), _present()])
call_order = _wire_view_setup(prisma_client, probe)
outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5)
actual = {
"outcome": outcome,
"calls": [call[0] for call in call_order.mock_calls],
"probe_args": probe.await_args.args,
}
assert actual == {
"outcome": "ready",
"calls": ["probe", "probe", "probe", "row_count", "views"],
"probe_args": (_PROBE_SQL, '"LiteLLM_SpendLogs"'),
}
@pytest.mark.asyncio
async def test_view_setup_probe_resolves_through_the_connection_search_path(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("DATABASE_SCHEMA", "litellm_tenant")
probe = AsyncMock(return_value=_present())
_wire_view_setup(prisma_client, probe)
await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5)
assert probe.await_args.args == (_PROBE_SQL, '"LiteLLM_SpendLogs"')
@pytest.mark.asyncio
async def test_view_setup_sets_the_row_count_even_when_view_creation_keeps_failing(
prisma_client: PrismaClient,
) -> None:
_wire_view_setup(prisma_client, AsyncMock(return_value=_present()))
prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public")
outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.02)
actual = {
"outcome": outcome,
"row_count_set": prisma_client._set_spend_logs_row_count_in_proxy_state.await_count >= 1,
}
assert actual == {"outcome": "timed_out", "row_count_set": True}
@pytest.mark.asyncio
async def test_view_setup_gives_up_when_the_table_never_appears(prisma_client: PrismaClient) -> None:
probe = AsyncMock(return_value=_absent())
_wire_view_setup(prisma_client, probe)
outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.02)
actual = {
"outcome": outcome,
"kept_polling": probe.await_count > 1,
"views_attempted": prisma_client.check_view_exists.await_count,
"row_count_attempted": prisma_client._set_spend_logs_row_count_in_proxy_state.await_count,
}
assert actual == {
"outcome": "timed_out",
"kept_polling": True,
"views_attempted": 0,
"row_count_attempted": 0,
}
@pytest.mark.asyncio
async def test_view_setup_retries_when_view_creation_fails_mid_migration(prisma_client: PrismaClient) -> None:
probe = AsyncMock(return_value=_present())
call_order = _wire_view_setup(prisma_client, probe)
prisma_client.check_view_exists.side_effect = [RuntimeError('column "tpd_limit" does not exist'), None]
outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5)
actual = {
"outcome": outcome,
"calls": [call[0] for call in call_order.mock_calls],
}
assert actual == {
"outcome": "ready",
"calls": ["probe", "row_count", "views", "probe", "row_count", "views"],
}
@pytest.mark.asyncio
async def test_view_setup_retries_when_the_table_probe_itself_fails(prisma_client: PrismaClient) -> None:
probe = AsyncMock(side_effect=[RuntimeError("connection reset"), _present()])
call_order = _wire_view_setup(prisma_client, probe)
outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5)
actual = {
"outcome": outcome,
"calls": [call[0] for call in call_order.mock_calls],
}
assert actual == {
"outcome": "ready",
"calls": ["probe", "probe", "row_count", "views"],
}
@pytest.mark.asyncio
async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout(
prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture
) -> None:
_wire_view_setup(prisma_client, AsyncMock(return_value=_absent()))
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01)
errors = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR]
actual = {
"outcome": outcome,
"error_count": len(errors),
"names_table": '"LiteLLM_SpendLogs"' in errors[0],
"tells_operator_to_migrate": "migrations" in errors[0] and "restart" in errors[0],
}
assert actual == {
"outcome": "timed_out",
"error_count": 1,
"names_table": True,
"tells_operator_to_migrate": True,
}
@pytest.mark.asyncio
async def test_run_view_setup_reports_the_last_error_when_views_keep_failing_on_a_present_table(
prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture
) -> None:
_wire_view_setup(prisma_client, AsyncMock(return_value=_present()))
prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public")
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01)
errors = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR]
actual = {
"outcome": outcome,
"error_count": len(errors),
"names_the_error": "permission denied for schema public" in errors[0],
"blames_missing_migrations": "did not appear" in errors[0],
"tells_operator_to_restart": "restart" in errors[0],
}
assert actual == {
"outcome": "timed_out",
"error_count": 1,
"names_the_error": True,
"blames_missing_migrations": False,
"tells_operator_to_restart": True,
}
@pytest.mark.asyncio
async def test_run_view_setup_stays_quiet_when_views_are_ready(
prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture
) -> None:
_wire_view_setup(prisma_client, AsyncMock(return_value=_present()))
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01)
actual = {
"outcome": outcome,
"errors": [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR],
}
assert actual == {"outcome": "ready", "errors": []}
@pytest.mark.asyncio
async def test_stop_view_setup_task_cancels_a_task_parked_between_polls(prisma_client: PrismaClient) -> None:
probe = AsyncMock(return_value=_absent())
_wire_view_setup(prisma_client, probe)
prisma_client.start_view_setup_task()
task = prisma_client._view_setup_task
await asyncio.sleep(0)
await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1)
actual = {
"probed_before_parking": probe.await_count,
"task_cancelled": task is not None and task.cancelled(),
"reference_cleared": prisma_client._view_setup_task,
"views_attempted": prisma_client.check_view_exists.await_count,
}
assert actual == {
"probed_before_parking": 1,
"task_cancelled": True,
"reference_cleared": None,
"views_attempted": 0,
}
@pytest.mark.asyncio
async def test_stop_view_setup_task_is_a_noop_without_a_task(prisma_client: PrismaClient) -> None:
await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1)
assert prisma_client._view_setup_task is None
@pytest.mark.asyncio
async def test_start_view_setup_task_twice_keeps_the_first_task(prisma_client: PrismaClient) -> None:
_wire_view_setup(prisma_client, AsyncMock(return_value=_absent()))
prisma_client.start_view_setup_task()
first = prisma_client._view_setup_task
prisma_client.start_view_setup_task()
second = prisma_client._view_setup_task
await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1)
actual = {
"first_is_task": isinstance(first, asyncio.Task),
"second_is_first": second is first,
}
assert actual == {"first_is_task": True, "second_is_first": True}