diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 7ae60121c6f..703abc64d8c 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -73,6 +73,17 @@ class RoutingPrismaWrapper: `connect()` or `recreate_prisma_client()` clears the flag. This keeps the proxy serving traffic during transient reader outages instead of failing startup or returning errors for read-heavy endpoints. + + Writer degradation: a writer-side `connect()` failure while the reader + connects is likewise non-fatal — the wrapper sets + `_writer_unavailable=True`, logs a warning, and keeps serving reads from + the reader (key lookups, DB-stored model loads) so a proxy that starts + during a primary outage still serves inference from the replica. Writes + fail at call time until the writer recovers; the PrismaClient DB health + watchdog polls `writer_unavailable` and drives the writer reconnect, + which clears the flag via `recreate_prisma_client`. Only when BOTH sides + fail to connect does `connect()` raise (full DB outage — the existing + `allow_requests_on_db_unavailable` startup handling applies). """ def __init__(self, writer: PrismaWrapper, reader: PrismaWrapper): @@ -81,6 +92,7 @@ class RoutingPrismaWrapper: # When True, reads fall back to the writer. Flipped on by reader # connect/recreate failures and flipped off on the next reader recovery. self._reader_unavailable: bool = False + self._writer_unavailable: bool = False @property def writer(self) -> PrismaWrapper: @@ -94,17 +106,46 @@ class RoutingPrismaWrapper: def reader_unavailable(self) -> bool: return self._reader_unavailable + @property + def writer_unavailable(self) -> bool: + return self._writer_unavailable + + def mark_writer_recovered(self) -> None: + """Clear the degraded-writer flag after an external health probe proved + the writer reachable. Needed when recovery happens without + `recreate_prisma_client` (e.g. an IAM token refresh already recreated + the writer engine), which is otherwise the only runtime path that + clears the flag — without this, the watchdog would keep firing + reconnect attempts against an already-healthy writer.""" + self._writer_unavailable = False + def _should_use_reader(self) -> bool: return not self._reader_unavailable - async def connect(self, *args: Any, **kwargs: Any) -> None: - await self._writer.connect(*args, **kwargs) - verbose_proxy_logger.info("[writer] DB connected") + @staticmethod + async def _try_connect(client: PrismaWrapper, *args: Any, **kwargs: Any) -> Exception | None: + if client.is_connected() is True: + return None try: - await self._reader.connect(*args, **kwargs) + await client.connect(*args, **kwargs) + return None + except Exception as e: + return e + + async def connect(self, *args: Any, **kwargs: Any) -> None: + writer_error = await self._try_connect(self._writer, *args, **kwargs) + if writer_error is None: + self._writer_unavailable = False + verbose_proxy_logger.info("[writer] DB connected") + reader_error = await self._try_connect(self._reader, *args, **kwargs) + if reader_error is None: self._reader_unavailable = False verbose_proxy_logger.info("[reader] DB connected") - except Exception as e: + if writer_error is None and reader_error is None: + return + if writer_error is not None and reader_error is not None: + raise writer_error + if reader_error is not None: # Degrade gracefully: the proxy keeps serving traffic with reads # routed to the writer until the reader endpoint is reachable. # Aborting startup here would tie proxy availability to an @@ -113,8 +154,15 @@ class RoutingPrismaWrapper: verbose_proxy_logger.warning( "Failed to connect to read replica DB: %s. " "Falling back to the writer for reads until the reader is reachable.", - e, + reader_error, ) + return + self._writer_unavailable = True + verbose_proxy_logger.warning( + "Failed to connect to primary (writer) DB: %s. " + "Serving reads from the read replica; writes will fail until the writer recovers.", + writer_error, + ) async def disconnect(self, *args: Any, **kwargs: Any) -> None: first_error: BaseException | None = None @@ -172,6 +220,7 @@ class RoutingPrismaWrapper: ) if not writer_recreated: return False + self._writer_unavailable = False try: await self._recreate_reader(http_client=http_client) self._reader_unavailable = False diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4433a35f5d0..2880eef6908 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4527,6 +4527,8 @@ class PrismaClient: "Writer healthy on probe; skipping recreate (engine " "likely already replaced by a token refresh)." ) + if isinstance(self.db, RoutingPrismaWrapper): + self.db.mark_writer_recovered() await self._start_engine_watcher() return except Exception as probe_err: @@ -4719,6 +4721,11 @@ class PrismaClient: self.db.query_raw("SELECT 1"), timeout=self._db_health_watchdog_probe_timeout_seconds, ) + if isinstance(self.db, RoutingPrismaWrapper) and self.db.writer_unavailable: + await self.attempt_db_reconnect( + reason="db_health_watchdog_writer_unavailable", + timeout_seconds=self._db_watchdog_reconnect_timeout_seconds, + ) except asyncio.CancelledError: break except Exception as e: diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index 265940e51ed..35ef0a965f3 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -513,3 +513,94 @@ async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect( # The flag must STILL be True so the next attempt re-enters the heavy # branch instead of silently demoting to the lightweight path. assert client._engine_confirmed_dead is True + + +@pytest.mark.asyncio +async def test_db_health_watchdog_should_reconnect_degraded_writer( + mock_proxy_logging, +): + """LIT-3792: when the proxy booted during a primary outage (reads served + by the replica, writer never connected), a healthy reader probe must not + mask the degraded writer — the watchdog drives the writer reconnect.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + writer = MagicMock() + reader = MagicMock() + reader.query_raw = AsyncMock(return_value=[{"result": 1}]) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = True + client.db = routing + client.attempt_db_reconnect = AsyncMock(return_value=True) + client._db_health_watchdog_interval_seconds = 1 + client._db_watchdog_reconnect_timeout_seconds = 7.0 + client._db_health_watchdog_probe_timeout_seconds = 0.2 + + with patch( + "litellm.proxy.utils.asyncio.sleep", + AsyncMock(side_effect=[None, asyncio.CancelledError()]), + ): + await client._db_health_watchdog_loop() + + client.attempt_db_reconnect.assert_awaited_once_with( + reason="db_health_watchdog_writer_unavailable", + timeout_seconds=7.0, + ) + + +@pytest.mark.asyncio +async def test_db_health_watchdog_should_not_reconnect_healthy_writer( + mock_proxy_logging, +): + """A healthy probe with no degraded writer must not trigger reconnects.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + writer = MagicMock() + reader = MagicMock() + reader.query_raw = AsyncMock(return_value=[{"result": 1}]) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + client.db = routing + client.attempt_db_reconnect = AsyncMock(return_value=True) + client._db_health_watchdog_interval_seconds = 1 + client._db_health_watchdog_probe_timeout_seconds = 0.2 + + with patch( + "litellm.proxy.utils.asyncio.sleep", + AsyncMock(side_effect=[None, asyncio.CancelledError()]), + ): + await client._db_health_watchdog_loop() + + client.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_direct_reconnect_probe_success_clears_writer_unavailable( + mock_proxy_logging, +): + """If the writer probe inside _do_direct_reconnect succeeds (engine already + reconnected by another path, e.g. an IAM token refresh), the early return + skips recreate_prisma_client — the degraded-writer flag must still be + cleared there or the watchdog fires reconnect attempts forever.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + writer = MagicMock() + writer.query_raw = AsyncMock(return_value=[{"result": 1}]) + reader = MagicMock() + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = True + client.db = routing + client._start_engine_watcher = AsyncMock() + + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + await client._run_reconnect_cycle(timeout_seconds=5.0) + + writer.query_raw.assert_awaited_once_with("SELECT 1") + assert routing.writer_unavailable is False diff --git a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py index efc3a6cf5b7..92043f44ca9 100644 --- a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py @@ -885,3 +885,111 @@ def test_prisma_client_init_falls_back_to_writer_when_reader_iam_token_fails( "Failed to initialize read replica Prisma client" in r.getMessage() for r in caplog.records ) + + +@pytest.mark.asyncio +async def test_connect_degrades_writer_when_reader_available(): + """A writer connect failure with a healthy reader must NOT abort proxy + startup (LIT-3792): startup swallows the raise when + allow_requests_on_db_unavailable is set, leaving the proxy with no Prisma + client at all, so DB-stored models never load and every request 400s. + Degrading instead keeps reads (key auth, model loads) on the replica.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.connect = AsyncMock(side_effect=RuntimeError("primary unreachable")) + reader_inner.connect = AsyncMock() + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + + # Must not raise — writer failure is non-fatal while the reader is up. + await routing.connect() + + assert routing.writer_unavailable is True + assert routing.reader_unavailable is False + writer_inner.connect.assert_awaited_once() + reader_inner.connect.assert_awaited_once() + + # Reads keep routing to the reader. + assert routing.query_raw is reader_inner.query_raw + + +@pytest.mark.asyncio +async def test_connect_raises_when_writer_and_reader_both_fail(): + """Full DB outage: with neither side reachable the wrapper must raise the + writer's error so existing allow_requests_on_db_unavailable startup + handling applies unchanged.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.connect = AsyncMock(side_effect=RuntimeError("primary down")) + reader_inner.connect = AsyncMock(side_effect=RuntimeError("replica down")) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + + with pytest.raises(RuntimeError, match="primary down"): + await routing.connect() + + +@pytest.mark.asyncio +async def test_connect_logs_writer_degradation(caplog): + """Operators need a clear signal that the proxy booted without a writer.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.connect = AsyncMock(side_effect=RuntimeError("primary unreachable")) + reader_inner.connect = AsyncMock() + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await routing.connect() + + assert any( + "Failed to connect to primary (writer) DB" in r.getMessage() + for r in caplog.records + ) + + +@pytest.mark.asyncio +async def test_recreate_clears_writer_unavailable(): + """A successful writer recreate (health watchdog reconnect once the + primary is back) must clear the degraded-writer flag.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer = MagicMock() + writer.recreate_prisma_client = AsyncMock(return_value=True) + reader = MagicMock() + reader.iam_token_db_auth = False + reader.recreate_prisma_client = AsyncMock() + + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = True + + with patch.dict(os.environ, {"DATABASE_URL_READ_REPLICA": "reader-url"}): + await routing.recreate_prisma_client("writer-url") + + assert routing.writer_unavailable is False + + +@pytest.mark.asyncio +async def test_recreate_keeps_writer_unavailable_when_writer_recreate_fails(): + """While the primary is still down, a failed writer recreate must leave + the degraded flag set so the watchdog keeps retrying.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer = MagicMock() + writer.recreate_prisma_client = AsyncMock( + side_effect=RuntimeError("primary still down") + ) + reader = MagicMock() + reader.iam_token_db_auth = False + reader.recreate_prisma_client = AsyncMock() + + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = True + + with ( + patch.dict(os.environ, {"DATABASE_URL_READ_REPLICA": "reader-url"}), + pytest.raises(RuntimeError, match="primary still down"), + ): + await routing.recreate_prisma_client("writer-url") + + assert routing.writer_unavailable is True