fix(proxy): keep serving reads from the read replica when the primary DB is down at startup (#31951)

* fix(proxy): keep serving reads from the read replica when the primary DB is down at startup

RoutingPrismaWrapper.connect() connected the writer first and let a writer
failure propagate, so a proxy that started during a primary outage ended up
with no Prisma client at all (startup swallows the error under
allow_requests_on_db_unavailable): DB-stored models never loaded and every
inference request failed with 400 Invalid model name, even with a healthy
DATABASE_URL_READ_REPLICA. Workers recycled via MAX_REQUESTS_BEFORE_RESTART
hit this mid-outage and stayed broken for the rest of the outage.

connect() now degrades on a writer-only failure: reads (key auth, DB-stored
model loads) are served by the reader, writes fail at call time, and the DB
health watchdog keeps retrying the writer reconnect, which clears the
degraded flag once the primary recovers. A full outage (both sides down)
still raises as before.

Resolves LIT-4159

* fix(proxy): clear degraded-writer flag when the reconnect probe finds the writer already healthy

The direct-reconnect path returns early when the writer probe succeeds
(engine already reconnected by another path, e.g. an IAM token refresh),
skipping recreate_prisma_client, which was the only runtime path clearing
_writer_unavailable. The stale flag made the watchdog fire reconnect
attempts against a healthy writer on every cooldown cycle until restart.
Clear the flag in the early-return branch and cover it with a regression
test that fails without the change
This commit is contained in:
Yassin Kortam 2026-07-04 22:45:09 +03:00 committed by GitHub
parent 07b9ea8c3b
commit cbb64111fa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 261 additions and 6 deletions

View file

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

View file

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

View file

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

View file

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