mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(proxy-extras): recover the v2 migration resolver from concurrent migrate deploy deadlocks
Two instances racing prisma migrate deploy on one database deadlock on CREATE INDEX CONCURRENTLY: the victim gets P3018 with 40P01 and the survivor then sees the failed ledger row as P3009. Both were treated as unrecoverable, so neither instance came up. Roll the deadlocked migration's ledger row back and retry the deploy on P3018, consult the failed row's logs in _prisma_migrations to do the same on P3009, and retry a deadlock reported without a Prisma error code. Genuinely broken migrations still fail fast.
This commit is contained in:
mateo-berri
parent
6028afa2ab
commit
692f3b513c
2 changed files with 228 additions and 4 deletions
|
|
@ -40,6 +40,8 @@ def _get_prisma_env() -> dict:
|
|||
|
||||
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
|
||||
|
||||
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
|
||||
|
||||
_SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
|
||||
_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
|
||||
r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
|
||||
|
|
@ -262,6 +264,48 @@ class ProxyExtrasDBManager:
|
|||
env=prisma_env,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _roll_back_migration_best_effort(migration_name: str) -> None:
|
||||
"""Mark a migration rolled back, tolerating a concurrent resolver
|
||||
having already done it."""
|
||||
try:
|
||||
ProxyExtrasDBManager._roll_back_migration(migration_name)
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _failed_migration_logs(migration_name: str) -> str:
|
||||
"""Logs recorded on the migration's failed _prisma_migrations row.
|
||||
|
||||
P3009 stderr does not carry the original failure, so this is the only
|
||||
way to tell a migration that lost a deadlock race against a concurrent
|
||||
migrate deploy from one whose SQL is genuinely broken. Returns "" when
|
||||
psycopg is missing, the DB is unreachable, or no failed row exists.
|
||||
"""
|
||||
database_url = os.getenv("DATABASE_URL")
|
||||
if not database_url:
|
||||
return ""
|
||||
|
||||
try:
|
||||
import psycopg
|
||||
except ImportError:
|
||||
return ""
|
||||
|
||||
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
|
||||
try:
|
||||
with psycopg.connect(
|
||||
cleaned_url, connect_timeout=10, autocommit=True
|
||||
) as conn:
|
||||
row = conn.execute(
|
||||
"SELECT logs FROM _prisma_migrations "
|
||||
"WHERE migration_name = %s AND finished_at IS NULL "
|
||||
"AND rolled_back_at IS NULL",
|
||||
(migration_name,),
|
||||
).fetchone()
|
||||
except (psycopg.OperationalError, psycopg.DatabaseError):
|
||||
return ""
|
||||
return (row[0] or "") if row else ""
|
||||
|
||||
@staticmethod
|
||||
def _resolve_specific_migration(migration_name: str):
|
||||
"""Mark a specific migration as applied"""
|
||||
|
|
@ -658,7 +702,8 @@ class ProxyExtrasDBManager:
|
|||
v2 migration resolver (opt-in via --use_v2_migration_resolver).
|
||||
|
||||
Runs `prisma migrate deploy` and handles standard recovery paths
|
||||
(P3005 baseline, P3009/P3018 idempotent errors). Critically, it does
|
||||
(P3005 baseline, P3009/P3018 idempotent errors, deadlocks against a
|
||||
concurrent migrate deploy). Critically, it does
|
||||
NOT call `_resolve_all_migrations` — the diff-and-force recovery that
|
||||
caused schema thrashing when two LiteLLM versions contended for the
|
||||
same DB during rolling deploys.
|
||||
|
|
@ -764,6 +809,22 @@ class ProxyExtrasDBManager:
|
|||
f"Detail: {resolve_err}"
|
||||
) from resolve_err
|
||||
continue
|
||||
if migration_match and _MIGRATION_DEADLOCK_MARKER in (
|
||||
ProxyExtrasDBManager._failed_migration_logs(
|
||||
migration_match.group(1)
|
||||
)
|
||||
):
|
||||
logger.info(
|
||||
"Migration %s lost a deadlock race against a "
|
||||
"concurrent migrate deploy, rolling its ledger "
|
||||
"row back and retrying",
|
||||
migration_match.group(1),
|
||||
)
|
||||
ProxyExtrasDBManager._roll_back_migration_best_effort(
|
||||
migration_match.group(1)
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
continue
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
|
|
@ -809,11 +870,33 @@ class ProxyExtrasDBManager:
|
|||
) from resolve_err
|
||||
continue
|
||||
|
||||
if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr:
|
||||
logger.info(
|
||||
"Migration %s deadlocked against a concurrent "
|
||||
"migrate deploy, rolling its ledger row back "
|
||||
"and retrying",
|
||||
migration_match.group(1),
|
||||
)
|
||||
ProxyExtrasDBManager._roll_back_migration_best_effort(
|
||||
migration_match.group(1)
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
continue
|
||||
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
if _MIGRATION_DEADLOCK_MARKER in stderr:
|
||||
logger.info(
|
||||
"prisma migrate deploy attempt %s deadlocked against "
|
||||
"a concurrent migrate deploy, retrying",
|
||||
attempt + 1,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
continue
|
||||
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
|
|
@ -821,9 +904,9 @@ class ProxyExtrasDBManager:
|
|||
|
||||
raise RuntimeError(
|
||||
"Database migration failed after 4 attempts (retry loop "
|
||||
"exhausted by timeouts or repeated idempotent-recovery "
|
||||
"continues). Check database connectivity, load, and "
|
||||
"_prisma_migrations ledger state."
|
||||
"exhausted by timeouts, deadlock retries, or repeated "
|
||||
"idempotent-recovery continues). Check database connectivity, "
|
||||
"load, and _prisma_migrations ledger state."
|
||||
)
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
|
|
|||
|
|
@ -240,3 +240,144 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
|
|||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"
|
||||
|
||||
|
||||
_DEADLOCK_P3018_STDERR = (
|
||||
"Error: P3018\n"
|
||||
"Migration name: 20260415120000_health_check_latest_per_model_index\n"
|
||||
"Database error code: 40P01\n"
|
||||
"deadlock detected"
|
||||
)
|
||||
|
||||
|
||||
def _stub_v2_env(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
monkeypatch.setattr("time.sleep", lambda _: None)
|
||||
|
||||
|
||||
def _succeed_after(failures: int, stderr: str):
|
||||
calls = {"n": 0}
|
||||
|
||||
class _OkResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
def _run(*args, **kwargs):
|
||||
if "deploy" not in args[0]:
|
||||
return _OkResult()
|
||||
calls["n"] += 1
|
||||
if calls["n"] <= failures:
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=args[0], stderr=stderr, output=""
|
||||
)
|
||||
return _OkResult()
|
||||
|
||||
return _run
|
||||
|
||||
|
||||
def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
|
||||
"""v2: losing the migrate deploy deadlock race against a concurrent
|
||||
instance rolls the ledger row back and retries instead of dying."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
|
||||
rolled_back = []
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_roll_back_migration",
|
||||
lambda name: rolled_back.append(name),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
|
||||
)
|
||||
monkeypatch.setattr("subprocess.run", _succeed_after(1, _DEADLOCK_P3018_STDERR))
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
|
||||
|
||||
|
||||
def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
|
||||
"""v2: a deadlock on every attempt still fails after the retry budget."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None)
|
||||
|
||||
with patch(
|
||||
"subprocess.run",
|
||||
side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR),
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="after 4 attempts"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path):
|
||||
"""v2: the surviving instance sees the victim's failed ledger row as P3009.
|
||||
When that row's logs show a deadlock, roll it back and retry."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\n"
|
||||
"migrate found failed migrations in the target database\n"
|
||||
"The `20260415120000_health_check_latest_per_model_index` migration "
|
||||
"started at 2026-09-01 18:46:13 UTC failed"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_failed_migration_logs",
|
||||
lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock",
|
||||
)
|
||||
rolled_back = []
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_roll_back_migration",
|
||||
lambda name: rolled_back.append(name),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
|
||||
)
|
||||
monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
|
||||
|
||||
|
||||
def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
|
||||
"""v2: a failed ledger row whose logs show a real SQL error stays fatal."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\n"
|
||||
"migrate found failed migrations in the target database\n"
|
||||
"The `20260101000000_genuinely_broken` migration started at "
|
||||
"2026-09-01 18:46:13 UTC failed"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_failed_migration_logs",
|
||||
lambda name: 'ERROR: syntax error at or near "BRKN"',
|
||||
)
|
||||
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path):
|
||||
"""v2: a deadlock reported without a Prisma error code (the advisory-lock
|
||||
waiter as victim) is retried, not fatal."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
monkeypatch.setattr(
|
||||
"subprocess.run", _succeed_after(1, "Database error: deadlock detected")
|
||||
)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue