diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 22da9b834ab..a31016e0f17 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -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." diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index f637e139e86..deb78dcdb30 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -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 diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index ac1defe8d79..dac36965a6c 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -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)