diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index d22484bc0e8..a4ea4789b49 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -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: """ diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index b3d457707b8..22558faedad 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -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 == []