mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
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:
parent
07b9ea8c3b
commit
cbb64111fa
4 changed files with 261 additions and 6 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue