diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index d22484bc0e8..21f255b0666 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,10 +804,11 @@ 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, ) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue @@ -781,7 +819,14 @@ class ProxyExtrasDBManager: logger.info( "Schema exists but no migrations ledger — creating baseline" ) - ProxyExtrasDBManager._create_baseline_migration(schema_path) + baselined = ProxyExtrasDBManager._create_baseline_migration( + schema_path + ) + budget = ( + budget.after_recovery("baseline") + if baselined + else budget.spend() + ) continue if "P3009" in stderr: @@ -818,6 +863,7 @@ class ProxyExtrasDBManager: f"intervention may be required.\n\n" f"Detail: {resolve_err}" ) from resolve_err + budget = budget.after_recovery(f"resolved:{name}") continue if migration_match: migration_name = migration_match.group(1) @@ -831,6 +877,7 @@ class ProxyExtrasDBManager: migration_name, ) ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue raise RuntimeError( @@ -876,6 +923,7 @@ class ProxyExtrasDBManager: f"intervention may be required.\n\n" f"Detail: {resolve_err}" ) from resolve_err + budget = budget.after_recovery(f"resolved:{name}") continue if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr: @@ -888,6 +936,7 @@ class ProxyExtrasDBManager: ProxyExtrasDBManager._roll_back_migration_best_effort( migration_match.group(1) ) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue @@ -900,8 +949,9 @@ class ProxyExtrasDBManager: logger.info( "prisma migrate deploy attempt %s deadlocked against " "a concurrent migrate deploy, retrying", - attempt + 1, + budget.attempt_number, ) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue @@ -909,8 +959,9 @@ class ProxyExtrasDBManager: logger.info( "prisma migrate deploy attempt %s timed out waiting for " "the advisory lock a concurrent migrate deploy holds, retrying", - attempt + 1, + budget.attempt_number, ) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue @@ -920,9 +971,9 @@ class ProxyExtrasDBManager: ) from e 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." ) 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 == []