Merge pull request #39506 from BerriAI/litellm_fix_v2_migration_resolver_attempt_accounting

fix(proxy-extras): only spend a migrate-deploy attempt when a pass made no progress
This commit is contained in:
Mateo Wang 2026-09-03 08:36:31 -07:00 committed by GitHub
commit 45495e1ab5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 340 additions and 149 deletions

View file

@ -6,6 +6,7 @@ import shutil
import subprocess
import tempfile
import time
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Optional
@ -45,6 +46,38 @@ _MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
@dataclass(frozen=True)
class _MigrateAttemptBudget:
"""Retries left, and the recoveries already run.
A recovery that lands something new costs nothing, so a database full of
objects `prisma db push` created works through them one per pass. Anything
that made no progress spends an attempt, so a stuck run still gives up.
"""
attempts_left: int
recoveries: frozenset[str] = frozenset()
@property
def exhausted(self) -> bool:
return self.attempts_left <= 0
@property
def attempt_number(self) -> int:
return MAX_MIGRATE_DEPLOY_ATTEMPTS - self.attempts_left + 1
def spend(self) -> "_MigrateAttemptBudget":
return replace(self, attempts_left=self.attempts_left - 1)
def after_recovery(self, recovery: str) -> "_MigrateAttemptBudget":
if recovery in self.recoveries:
return self.spend()
return replace(self, recoveries=self.recoveries | {recovery})
_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
@ -716,6 +749,9 @@ class ProxyExtrasDBManager:
Ahead-of-HEAD state (DB has migrations newer than this build ships)
is logged as a warning, not a fatal error users whose DBs got into
weird shapes from the old thrashing should still be able to start.
The retry budget only counts attempts that made no progress: see
_MigrateAttemptBudget.
"""
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
@ -749,8 +785,9 @@ class ProxyExtrasDBManager:
original_dir = os.getcwd()
os.chdir(migrations_dir)
deploy_timeout = prisma_migrate_deploy_timeout()
budget = _MigrateAttemptBudget(attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS)
try:
for attempt in range(4):
while not budget.exhausted:
try:
result = subprocess.run(
[_get_prisma_command(), "migrate", "deploy"],
@ -767,168 +804,155 @@ class ProxyExtrasDBManager:
logger.warning(
"prisma migrate deploy attempt %s timed out after %ss, retrying. "
"Raise %s if this database needs longer to apply its pending migrations.",
attempt + 1,
budget.attempt_number,
deploy_timeout,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
)
time.sleep(random.randrange(5, 15))
continue
next_budget = budget.spend()
except subprocess.CalledProcessError as e:
stderr = e.stderr or ""
next_budget = ProxyExtrasDBManager._budget_after_deploy_failure(
e, budget, schema_path
)
if "P3005" in stderr and "database schema is not empty" in stderr:
logger.info(
"Schema exists but no migrations ledger — creating baseline"
)
ProxyExtrasDBManager._create_baseline_migration(schema_path)
continue
if "P3009" in stderr:
migration_match = re.search(r"`(\d+_\S+?)`", stderr)
if (
migration_match
and ProxyExtrasDBManager._is_idempotent_error(stderr)
):
name = migration_match.group(1)
logger.info(
f"Migration {name} failed idempotently — marking applied and retrying"
)
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
# We're already inside the outer
# `except CalledProcessError` handler —
# re-raising CalledProcessError from here
# would escape as itself, bypassing
# proxy_cli.py's `except RuntimeError`.
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
continue
if migration_match:
migration_name = migration_match.group(1)
ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
if ledger_logs is not None and (
ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs
):
logger.info(
"Migration %s failed in a concurrent migrate deploy "
"deadlock race, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name)
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 "P3018" in stderr:
if ProxyExtrasDBManager._is_permission_error(stderr):
raise RuntimeError(
"Database migration failed due to insufficient "
"permissions. Please grant the required privileges "
f"and retry.\n\nPrisma error:\n{stderr}"
) from e
migration_match = re.search(
r"Migration name: (\d+_\S+)", stderr
)
if (
migration_match
and ProxyExtrasDBManager._is_idempotent_error(stderr)
):
name = migration_match.group(1)
logger.info(
f"Migration {name} SQL hit idempotent error — marking applied and retrying"
)
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) 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
if "P1002" in stderr and "advisory lock" in stderr:
logger.info(
"prisma migrate deploy attempt %s timed out waiting for "
"the advisory lock a concurrent migrate deploy holds, 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}"
) from e
if next_budget.attempts_left < budget.attempts_left:
time.sleep(random.randrange(5, 15))
budget = next_budget # rebind-ok: the loop carries the budget from one migrate deploy pass to the next
raise RuntimeError(
"Database migration failed after 4 attempts (retry loop "
"exhausted by timeouts, deadlock retries, or repeated "
"idempotent-recovery continues). Check database connectivity, "
f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
"attempts that made no progress (timeouts, deadlock retries, or a "
"recovery that had already run once). Check database connectivity, "
"load, and _prisma_migrations ledger state, and raise "
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
)
finally:
os.chdir(original_dir)
@staticmethod
def _budget_after_deploy_failure(
error: subprocess.CalledProcessError,
budget: "_MigrateAttemptBudget",
schema_path: str,
) -> "_MigrateAttemptBudget":
"""Recover from one failed `prisma migrate deploy`, and price the pass.
Returns the budget the next pass runs under, or raises when the failure
is not one this resolver knows how to recover from.
"""
stderr = error.stderr or ""
if "P3005" in stderr and "database schema is not empty" in stderr:
logger.info("Schema exists but no migrations ledger — creating baseline")
if ProxyExtrasDBManager._create_baseline_migration(schema_path):
return budget.after_recovery("baseline")
return budget.spend()
if "P3009" in stderr:
migration_match = re.search(r"`(\d+_\S+?)`", stderr)
if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
name = migration_match.group(1)
logger.info(
f"Migration {name} failed idempotently — marking applied and retrying"
)
ProxyExtrasDBManager._mark_migration_applied(name)
return budget.after_recovery(f"resolved:{name}")
if migration_match:
migration_name = migration_match.group(1)
ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
if ledger_logs is not None and (
ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs
):
logger.info(
"Migration %s failed in a concurrent migrate deploy "
"deadlock race, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
if "P3018" in stderr:
if ProxyExtrasDBManager._is_permission_error(stderr):
raise RuntimeError(
"Database migration failed due to insufficient "
"permissions. Please grant the required privileges "
f"and retry.\n\nPrisma error:\n{stderr}"
) from error
migration_match = re.search(r"Migration name: (\d+_\S+)", stderr)
if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
name = migration_match.group(1)
logger.info(
f"Migration {name} SQL hit idempotent error — marking applied and retrying"
)
ProxyExtrasDBManager._mark_migration_applied(name)
return budget.after_recovery(f"resolved:{name}")
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)
)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
if _MIGRATION_DEADLOCK_MARKER in stderr:
logger.info(
"prisma migrate deploy attempt %s deadlocked against "
"a concurrent migrate deploy, retrying",
budget.attempt_number,
)
return budget.spend()
if "P1002" in stderr and "advisory lock" in stderr:
logger.info(
"prisma migrate deploy attempt %s timed out waiting for "
"the advisory lock a concurrent migrate deploy holds, retrying",
budget.attempt_number,
)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
@staticmethod
def _mark_migration_applied(name: str) -> None:
"""Roll a failed ledger row back if it is still there, then mark it applied."""
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
# We're called from inside an `except CalledProcessError` handler —
# re-raising CalledProcessError from here would escape as itself,
# bypassing proxy_cli.py's `except RuntimeError`.
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
@staticmethod
def apply_replica_identity_full_if_requested() -> bool:
"""

View file

@ -703,3 +703,170 @@ class TestSpendLogsPartitionDetectionMissingPsycopg:
assert any(
"psycopg is not installed" in record.message for record in caplog.records
)
_ATTEMPT_BUDGET = 4
_P3005_STDERR = """Error: P3005
The database schema is not empty. Read more about how to baseline an existing production database: https://pris.ly/d/migrate-baseline
"""
def _p3018_stderr(migration_name):
return f"""Error: P3018
A migration failed to apply. New migrations cannot be applied before the error is recovered from.
Migration name: {migration_name}
Database error code: 42P07
Database error:
ERROR: relation "SomeTable" already exists
"""
class _MigrateDeployHarness:
"""Drives _setup_database_v2 with a scripted sequence of
`prisma migrate deploy` outcomes, with every recovery command faked out so
nothing touches a database or the packaged migrations directory."""
def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False):
import subprocess as subprocess_module
import litellm_proxy_extras.utils as utils_module
self.deploy_calls = []
self.resolved = []
self.baselines = 0
self._outcomes = list(outcomes)
self._repeat_last = repeat_last
self._subprocess_module = subprocess_module
monkeypatch.delenv("DATABASE_URL", raising=False)
monkeypatch.setattr(
ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path))
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_create_baseline_migration",
staticmethod(self._fake_baseline),
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_roll_back_migration",
staticmethod(lambda name: None),
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_resolve_specific_migration",
staticmethod(self.resolved.append),
)
monkeypatch.setattr(utils_module.subprocess, "run", self._fake_run)
monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None)
self.baseline_succeeds = True
def _fake_baseline(self, *args, **kwargs):
self.baselines += 1
return self.baseline_succeeds
def _next_outcome(self):
if self._outcomes:
if self._repeat_last and len(self._outcomes) == 1:
return self._outcomes[0]
return self._outcomes.pop(0)
raise AssertionError("prisma migrate deploy called more times than scripted")
def _fake_run(self, cmd, **kwargs):
assert cmd[1:] == ["migrate", "deploy"], f"unexpected prisma command: {cmd}"
self.deploy_calls.append(cmd)
outcome = self._next_outcome()
if outcome == "ok":
return _FakeCompleted()
if outcome == "timeout":
raise self._subprocess_module.TimeoutExpired(cmd, 1)
raise self._subprocess_module.CalledProcessError(1, cmd, stderr=outcome)
def run(self):
return ProxyExtrasDBManager._setup_database_v2(use_migrate=True)
class TestMigrateDeployAttemptAccounting:
"""A `prisma db push` database has a full schema and no ledger, so the v2
resolver baselines it and then works through every migration whose objects
already exist. Those recoveries make progress, so they must not spend the
retry budget, which is there to stop a run that is getting nowhere."""
def test_a_push_created_database_finishes_bootstrapping(
self, monkeypatch, tmp_path
):
already_there = [
"20250329084805_new_cron_job_table",
"20250806095134_rename_alias_to_server_name_mcp_table",
"20260224203854_add_agent_object_permissions_table",
"20260301120000_fourth_table",
"20260302120000_fifth_table",
"20260303120000_sixth_table",
]
harness = _MigrateDeployHarness(
monkeypatch,
tmp_path,
[_P3005_STDERR]
+ [_p3018_stderr(name) for name in already_there]
+ ["ok"],
)
assert harness.run() is True
assert harness.baselines == 1
assert harness.resolved == already_there
assert len(harness.deploy_calls) == len(already_there) + 2
def test_repeated_recovery_of_one_migration_still_gives_up(
self, monkeypatch, tmp_path
):
harness = _MigrateDeployHarness(
monkeypatch,
tmp_path,
[_p3018_stderr("20250329084805_new_cron_job_table")],
repeat_last=True,
)
with pytest.raises(RuntimeError):
harness.run()
assert len(harness.deploy_calls) <= _ATTEMPT_BUDGET + 1
def test_timeouts_still_spend_the_budget(self, monkeypatch, tmp_path):
harness = _MigrateDeployHarness(
monkeypatch, tmp_path, ["timeout"], repeat_last=True
)
with pytest.raises(RuntimeError):
harness.run()
assert len(harness.deploy_calls) == _ATTEMPT_BUDGET
def test_a_baseline_that_never_lands_stops_after_the_budget(
self, monkeypatch, tmp_path
):
harness = _MigrateDeployHarness(
monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True
)
harness.baseline_succeeds = False
with pytest.raises(RuntimeError):
harness.run()
assert len(harness.deploy_calls) == _ATTEMPT_BUDGET
def test_an_unrecoverable_error_is_not_retried(self, monkeypatch, tmp_path):
harness = _MigrateDeployHarness(
monkeypatch,
tmp_path,
["Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near \"SLECT\"\n"],
repeat_last=True,
)
with pytest.raises(RuntimeError):
harness.run()
assert len(harness.deploy_calls) == 1
assert harness.resolved == []