diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 36fdea605c2..a4d99d9859e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1464,6 +1464,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() @@ -10635,14 +10641,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() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b078a65759e..7f5609b2e39 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -38,6 +38,7 @@ from typing import ( Literal, Optional, Protocol, + TypeAlias, TypeVar, Union, cast, @@ -268,6 +269,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: ... @@ -4284,6 +4294,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 @@ -6330,6 +6341,78 @@ 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() + try: + await self._view_setup_task + except asyncio.CancelledError: + pass + 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", self._view_setup_gate_table() + ) + return "table_missing" + await self.check_view_exists() + await self._set_spend_logs_row_count_in_proxy_state() + 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.", + self._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", self._view_setup_gate_table()) + ) + return rows[0]["present"] + + @staticmethod + def _view_setup_gate_table() -> str: + pg_schema: Final = os.getenv("DATABASE_SCHEMA", "public") + return f'"{pg_schema}"."{_VIEW_SETUP_GATE_TABLE}"' + async def _db_health_watchdog_loop(self) -> None: while True: try: diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 9ac97f57f47..233aa81c2af 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -359,6 +359,19 @@ 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: + """Re-probe a route that depends on startup work the proxy finishes after it + starts serving, such as the spend views it creates once migrations land.""" + 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( diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py index 8cb3e3927f0..67fd88bc84d 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py @@ -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]}" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e2e2045e826..7ab142c5fd0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1628,6 +1628,40 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): assert master_key == test_resolved_key +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): + """The view setup task keeps polling for the spend-log table while migrations + run, so a shutdown inside that window has to cancel it rather than leave it + to die with the event loop.""" + 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 @@ -13307,6 +13341,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 @@ -13368,13 +13403,39 @@ 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): + """View creation used to be two fire-and-forget ``asyncio.create_task`` calls + that raised and vanished when the migrations Job had not created + ``LiteLLM_SpendLogs`` yet (LIT-5211). Startup must hand the work to the client's + held task, which waits for the table, and must not call the two coroutines directly.""" + 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 diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index 18b02ac7772..bb34486771c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -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,217 @@ 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, monkeypatch: pytest.MonkeyPatch +) -> None: + """On a fresh database the migrations Job can still be running when the proxy + boots. The views reference ``LiteLLM_SpendLogs``, so creating them before the + table exists raised inside a fire-and-forget task and the views never appeared.""" + monkeypatch.delenv("DATABASE_SCHEMA", raising=False) + 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", "views", "row_count"], + "probe_args": (_PROBE_SQL, '"public"."LiteLLM_SpendLogs"'), + } + + +@pytest.mark.asyncio +async def test_view_setup_probes_the_configured_database_schema( + 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_tenant"."LiteLLM_SpendLogs"') + + +@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: + """``LiteLLM_SpendLogs`` lands early in the migration set while + ``LiteLLM_VerificationTokenView`` references columns the newest migrations add, + so the first attempt after the table appears can still fail.""" + 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", "views", "probe", "views", "row_count"], + } + + +@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", "views", "row_count"], + } + + +@pytest.mark.asyncio +async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.delenv("DATABASE_SCHEMA", raising=False) + _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": '"public"."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: + """A database role without CREATE on the schema fails every attempt even though + the table is there, so the timeout must blame that error, not missing migrations.""" + _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}