scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(proxy-extras): retry transient db push failures, drop a vacuous test

Browse source 
`prisma db push` under v2 raised on the first failure while v1 retried it four
times, so making v2 the default silently cost --use_prisma_db_push its
retries. It now uses the same transient classification as migrate deploy.

The classifier moves onto ProxyExtrasDBManager next to _is_permission_error
and _is_idempotent_error, which do the same kind of stderr matching.

Replaces a test that claimed to pin the transient classification but fed it a
P3009 stderr, which an earlier branch catches, so it passed even when the
classifier was mutated to treat everything as transient. The replacement uses
an unclassified error and fails on that mutant. Drops a v1 test that duplicated
test_v1_default_still_calls_resolve_all_migrations.
This commit is contained in:
Yuneng Jiang 2026-08-28 01:31:44 -07:00
parent 4f6fd85ab1
commit 7b36bfb967
No known key found for this signature in database
3 changed files with 110 additions and 113 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

98
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -51,9 +51,9 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
re.IGNORECASE,
)
_MIGRATE_DEPLOY_ATTEMPTS: Final = 4
_PRISMA_ATTEMPTS: Final = 4
_TRANSIENT_DEPLOY_FAILURES: Final = MappingProxyType(
_TRANSIENT_PRISMA_FAILURES: Final = MappingProxyType(
{
"deadlock detected": "a deadlock on the migration advisory lock (a concurrent migrate deploy)",
"P1001": "an unreachable database server",
@ -62,24 +62,6 @@ _TRANSIENT_DEPLOY_FAILURES: Final = MappingProxyType(
)
def _transient_deploy_failure(stderr: str) -> str | None:
"""Describe why a failed `prisma migrate deploy` is worth retrying, or None.
These are environment failures, not migration failures: the database is not
up yet, or another instance holds the migration lock. v1 retried every
failed deploy and absorbed them; failing fast on them instead would turn a
database that is ten seconds late into a dead proxy.
"""
return next(
(
reason
for marker, reason in _TRANSIENT_DEPLOY_FAILURES.items()
if marker in stderr
),
None,
)
PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
"LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
"so its primary key must include the partition key (\"startTime\"). `prisma db push` "
@ -304,6 +286,23 @@ class ProxyExtrasDBManager:
env=prisma_env,
)
@staticmethod
def _transient_prisma_failure(stderr: str) -> str | None:
"""Why a failed prisma command is worth retrying, or None.
v1 retried every failure, so it absorbed a database that was not up yet
or another instance holding the migration lock. v2 fails fast, which is
right for a broken migration and wrong for these.
"""
return next(
(
reason
for marker, reason in _TRANSIENT_PRISMA_FAILURES.items()
if marker in stderr
),
None,
)
@staticmethod
def _is_permission_error(error_message: str) -> bool:
"""
@ -699,20 +698,43 @@ class ProxyExtrasDBManager:
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=prisma_command_timeout(),
check=True,
env=_get_prisma_env(),
for attempt in range(_PRISMA_ATTEMPTS):
try:
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
env=_get_prisma_env(),
)
return True
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as e:
stderr = e.stderr or ""
transient = ProxyExtrasDBManager._transient_prisma_failure(
stderr
)
# Re-raise as RuntimeError so proxy_cli.py's
# `except RuntimeError` catches it and exits cleanly.
if transient is None or attempt == _PRISMA_ATTEMPTS - 1:
raise RuntimeError(
f"prisma db push failed.\n\nDetail: {e}"
f"\n\nPrisma error:\n{stderr}"
) from e
logger.info(
"prisma db push attempt %s failed on %s, retrying. "
"Prisma error:\n%s",
attempt + 1,
transient,
stderr,
)
time.sleep(random.randrange(5, 15))
raise RuntimeError(
f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts."
)
return True
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as e:
# Re-raise as RuntimeError so proxy_cli.py's
# `except RuntimeError` catches it and exits cleanly.
raise RuntimeError(f"prisma db push failed.\n\nDetail: {e}") from e
finally:
os.chdir(original_dir)
@ -722,7 +744,7 @@ class ProxyExtrasDBManager:
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
for attempt in range(_MIGRATE_DEPLOY_ATTEMPTS):
for attempt in range(_PRISMA_ATTEMPTS):
try:
result = subprocess.run(
[_get_prisma_command(), "migrate", "deploy"],
@ -837,17 +859,17 @@ class ProxyExtrasDBManager:
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
transient = _transient_deploy_failure(stderr)
transient = ProxyExtrasDBManager._transient_prisma_failure(stderr)
if transient is None:
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
if attempt == _MIGRATE_DEPLOY_ATTEMPTS - 1:
if attempt == _PRISMA_ATTEMPTS - 1:
raise RuntimeError(
f"Database migration failed after "
f"{_MIGRATE_DEPLOY_ATTEMPTS} attempts on {transient}. "
f"{_PRISMA_ATTEMPTS} attempts on {transient}. "
"Check database connectivity and load."
f"\n\nPrisma error:\n{stderr}"
) from e
@ -863,7 +885,7 @@ class ProxyExtrasDBManager:
continue
raise RuntimeError(
f"Database migration failed after {_MIGRATE_DEPLOY_ATTEMPTS} "
f"Database migration failed after {_PRISMA_ATTEMPTS} "
"attempts (retry loop exhausted by timeouts or repeated "
"idempotent-recovery continues). Check database connectivity, "
"load, and _prisma_migrations ledger state."

113
tests/litellm-proxy-extras/test_setup_database_fail_fast.py
View file

@ -1,10 +1,7 @@
"""Regression tests for ProxyExtrasDBManager's v2 migration resolver.
v2 is what the proxy CLI selects by default; v1 stays reachable via
`--use_legacy_migration_resolver` or `USE_V2_MIGRATION_RESOLVER=false`. At the
library level the resolver is picked with the `use_v2_resolver` kwarg, which
still defaults to False so `migrations/run.py` and any direct caller keep their
own explicit choice.
v2 is the proxy CLI default; v1 stays reachable via the `use_v2_resolver`
kwarg, which still defaults to False for direct callers.
"""
import os
@ -271,9 +268,7 @@ class _DeployApplied:
def _deploy_only(deploy_side_effect):
"""subprocess.run stand-in that only intercepts `prisma migrate deploy`.
Everything else the resolver shells out to, the Prisma toolchain check
above all, succeeds untouched, so a mock meant for the deploy call cannot
be silently consumed by an earlier subprocess call.
Scoped by argv so the Prisma toolchain check cannot consume the mock first.
"""
deploys = {"n": 0}
@ -298,13 +293,8 @@ def _prepare_v2_resolver(monkeypatch, tmp_path):
def test_v2_retries_transient_advisory_lock_deadlock(monkeypatch, tmp_path):
"""A deadlock on Prisma's migration advisory lock is transient and retried.
Several proxy replicas booting against one database race `migrate deploy`,
and Postgres aborts one side. v1 retried any failed deploy, so it rode this
out; v2 classifies unrecognised stderr as unrecoverable and raises, which
with v2 as the default would take a replica's whole boot down.
"""
"""v2: replicas racing `migrate deploy` deadlock on Prisma's advisory
lock, which is transient and must be retried rather than kill the boot."""
_prepare_v2_resolver(monkeypatch, tmp_path)
def _side_effect(n, cmd):
@ -323,8 +313,8 @@ def test_v2_retries_transient_advisory_lock_deadlock(monkeypatch, tmp_path):
def test_v2_persistent_advisory_lock_deadlock_eventually_raises(monkeypatch, tmp_path):
"""The deadlock retry stays bounded: a deadlock that never clears still
raises rather than looping forever or reporting a successful migration."""
"""v2: the deadlock retry is bounded, so a deadlock that never clears
still raises instead of looping or reporting success."""
_prepare_v2_resolver(monkeypatch, tmp_path)
def _side_effect(n, cmd):
@ -348,13 +338,7 @@ def test_v2_persistent_advisory_lock_deadlock_eventually_raises(monkeypatch, tmp
],
)
def test_v2_retries_transient_database_connectivity_errors(monkeypatch, tmp_path, stderr):
"""A database that is not accepting connections yet is retried, not fatal.
A proxy and its database starting together race routinely, and v1 rode that
out by retrying every failed deploy. v2 treats unrecognised stderr as
unrecoverable, so without this the default flip would turn a database that
is a few seconds late into a dead proxy instead of a slow boot.
"""
"""v2: a database not accepting connections yet is retried, not fatal."""
_prepare_v2_resolver(monkeypatch, tmp_path)
def _side_effect(n, cmd):
@ -373,9 +357,8 @@ def test_v2_retries_transient_database_connectivity_errors(monkeypatch, tmp_path
def test_v2_unreachable_database_still_fails_after_the_retries(monkeypatch, tmp_path):
"""Retrying connectivity errors must not turn a genuinely unreachable
database into a silent success: after the attempts are spent it still
raises, so the proxy exits instead of serving without its database."""
"""v2: a genuinely unreachable database still raises once the attempts
are spent, rather than passing as a successful migration."""
_prepare_v2_resolver(monkeypatch, tmp_path)
def _side_effect(n, cmd):
@ -395,13 +378,8 @@ def test_v2_unreachable_database_still_fails_after_the_retries(monkeypatch, tmp_
def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, caplog):
"""Retrying must not swallow why the database was unreachable.
Prisma's stderr is captured, so if the retry path neither logs it nor puts
it in the final error, an operator (and CI's bad-DATABASE_URL job, which
greps the boot log for the P1001 line) sees four silent retries and no
cause.
"""
"""v2: retrying must not swallow Prisma's stderr, which is captured and is
the only place the cause appears for an operator or a boot-log grep."""
_prepare_v2_resolver(monkeypatch, tmp_path)
stderr = "Error: P1001: Can't reach database server at `wrong`:`5432`"
@ -422,21 +400,47 @@ def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, cap
assert "P1001" in caplog.text
def test_v2_migration_failure_is_not_treated_as_transient(monkeypatch, tmp_path):
"""The transient classification must stay narrow: a genuinely broken
migration still fails fast on the first attempt rather than being retried
into the same error four times."""
def test_v2_db_push_retries_transient_failures(monkeypatch, tmp_path):
"""v2: `prisma db push` retries a transient failure like v1 did, so the
default flip does not cost --use_prisma_db_push its retries."""
_prepare_v2_resolver(monkeypatch, tmp_path)
stderr = (
"Error: P3009\n"
"The `20260101000000_genuinely_broken` migration failed to apply.\n"
'Reason: syntax error at or near "BRKN" LINE 42'
pushes = {"n": 0}
def _run(*args, **kwargs):
cmd = list(args[0] if args else kwargs.get("args", []))
if cmd[-3:] != ["db", "push", "--accept-data-loss"]:
return _DeployApplied()
pushes["n"] += 1
if pushes["n"] == 1:
raise subprocess.CalledProcessError(
returncode=1,
cmd=cmd,
stderr="Error: P1001: Can't reach database server at `db`:`5432`",
output="",
)
return _DeployApplied()
monkeypatch.setattr(
ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False
)
with patch("subprocess.run", side_effect=_run):
ok = ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
assert ok is True
assert pushes["n"] == 2
def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path):
"""v2: an unrecognised deploy failure still raises on the first attempt."""
_prepare_v2_resolver(monkeypatch, tmp_path)
def _side_effect(n, cmd):
raise subprocess.CalledProcessError(
returncode=1, cmd=cmd, stderr=stderr, output=""
returncode=1,
cmd=cmd,
stderr="Error: relation \"LiteLLM_SpendLogs\" does not exist",
output="",
)
run, deploys = _deploy_only(_side_effect)
@ -447,26 +451,3 @@ def test_v2_migration_failure_is_not_treated_as_transient(monkeypatch, tmp_path)
assert deploys["n"] == 1
def test_v1_still_runs_the_diff_and_force_recovery(monkeypatch, tmp_path):
"""v1 remains the pre-existing diff-and-force resolver, unchanged by the
default flip: it still calls _resolve_all_migrations after a deploy that
applied something. Operators opting back in must get exactly the old path.
"""
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
class FakeResult:
stdout = "Applied migration.\n"
stderr = ""
resolve_called = {"n": 0}
def fake_resolve(*args, **kwargs):
resolve_called["n"] += 1
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=False)
assert ok is True
assert resolve_called["n"] == 1

12
tests/test_litellm/proxy/test_proxy_cli.py
View file

@ -2008,15 +2008,9 @@ class TestRunServerDbSetup:
env_value,
expected_v2,
):
"""The proxy defaults to the v2 resolver, and v1 stays reachable.
Both opt-out routes matter: --use_legacy_migration_resolver for a CLI
boot, and USE_V2_MIGRATION_RESOLVER=false for containerised deploys,
where litellm/proxy/prisma_migration.py calls run_server with a fixed
argv and an env var is the only way in. The deprecated
--use_v2_migration_resolver must still parse so existing commands do
not die on an unknown option, and an explicit flag still beats the env.
"""
"""The proxy defaults to v2, and both v1 opt-out routes work: the
flag for a CLI boot, USE_V2_MIGRATION_RESOLVER=false for deploys that
cannot pass one. An explicit flag beats the env var."""
from litellm.proxy.proxy_cli import run_server
mock_subprocess_run.return_value = MagicMock(returncode=0)