fix(proxy): wait for the spend-log table before creating startup views

litellm/proxy/proxy_server.py

@@ -1510,6 +1510,12 @@ async def _run_agent_grant_id_migration() -> 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 @@ async def _setup_prisma_client(
             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()

litellm/proxy/utils.py

@@ -38,6 +38,7 @@
     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 @@ def __init__(
             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 @@ async def stop_db_health_watchdog_task(self) -> None:
         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:

tests/e2e/quota_management/spend_tracking/spend_e2e_client.py

@@ -359,6 +359,17 @@ def spend_logs_page(
     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(

tests/e2e/quota_management/spend_tracking/test_spend_routes.py

@@ -17,9 +17,11 @@
 """
 
 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_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]}"
 

tests/proxy_unit_tests/test_proxy_server.py

@@ -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):

tests/test_litellm/proxy/test_proxy_server.py

@@ -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