mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
feat(proxy): default to the v2 migration resolver, keep v1 as an opt-out
The v2 resolver skips the diff-and-force recovery that caused schema thrashing when two LiteLLM versions contend for one database during a rolling deploy. The standalone migration Job already defaulted to v2; this aligns the proxy-server path. v1 stays reachable two ways: --use_legacy_migration_resolver on the CLI, and USE_V2_MIGRATION_RESOLVER=false for containerised deploys, where prisma_migration.py calls run_server with a fixed argv and the env var is the only route in. --use_v2_migration_resolver still parses, so existing commands do not die on an unknown option. Because v2 fails fast where v1 retried every failed deploy, a database that is not accepting connections yet, or another instance holding the migration advisory lock, would now kill a boot that used to ride it out. Those two failures are retried, with Prisma's stderr logged each round, and still raise once the attempts are spent. Moves the resolver tests from litellm-proxy-extras/tests, which no CI job runs, into tests/litellm-proxy-extras, and repoints the dedicated Postgres CircleCI job at the legacy path so v1 keeps real-DB and proxy-boot coverage.
This commit is contained in:
Yuneng Jiang
parent
ca0b951a43
commit
4f6fd85ab1
8 changed files with 409 additions and 46 deletions
|
|
@ -1483,7 +1483,7 @@ jobs:
|
|||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
|
||||
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
|
||||
|
||||
installing_litellm_on_python_3_13:
|
||||
docker:
|
||||
|
|
@ -1507,9 +1507,9 @@ jobs:
|
|||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
|
||||
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
|
||||
|
||||
installing_litellm_on_python_v2_migration_resolver:
|
||||
installing_litellm_on_python_legacy_migration_resolver:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
|
|
@ -1536,10 +1536,10 @@ jobs:
|
|||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Run v2 migration resolver proxy smoke test
|
||||
name: Run legacy migration resolver proxy smoke test
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -vv \
|
||||
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_v2_resolver
|
||||
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver
|
||||
|
||||
helm_chart_testing:
|
||||
machine:
|
||||
|
|
@ -2918,7 +2918,8 @@ jobs:
|
|||
command: |
|
||||
if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \
|
||||
(grep -q "Database setup failed after multiple retries" docker_output.log || \
|
||||
grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then
|
||||
grep -q "ERROR: Application startup failed. Exiting." docker_output.log || \
|
||||
grep -q "Database migration cannot proceed" docker_output.log); then
|
||||
echo "Expected error found. Test passed."
|
||||
else
|
||||
echo "Expected error not found. Test failed."
|
||||
|
|
@ -3050,7 +3051,7 @@ workflows:
|
|||
filters: *main_branches
|
||||
- installing_litellm_on_python_3_13:
|
||||
filters: *main_branches
|
||||
- installing_litellm_on_python_v2_migration_resolver:
|
||||
- installing_litellm_on_python_legacy_migration_resolver:
|
||||
filters: *main_branches
|
||||
- helm_chart_testing:
|
||||
requires:
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
|
|||
|
||||
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
|
||||
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
|
||||
|
||||
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
|
||||
- don't use emojis
|
||||
|
|
|
|||
|
|
@ -7,7 +7,8 @@ import subprocess
|
|||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Optional
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.replica_identity import (
|
||||
|
|
@ -50,6 +51,35 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
|
|||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
_MIGRATE_DEPLOY_ATTEMPTS: Final = 4
|
||||
|
||||
_TRANSIENT_DEPLOY_FAILURES: Final = MappingProxyType(
|
||||
{
|
||||
"deadlock detected": "a deadlock on the migration advisory lock (a concurrent migrate deploy)",
|
||||
"P1001": "an unreachable database server",
|
||||
"P1002": "a database server that timed out",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
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` "
|
||||
|
|
@ -648,7 +678,7 @@ class ProxyExtrasDBManager:
|
|||
@staticmethod
|
||||
def _setup_database_v2(use_migrate: bool) -> bool:
|
||||
"""
|
||||
v2 migration resolver (opt-in via --use_v2_migration_resolver).
|
||||
v2 migration resolver (what the proxy CLI selects by default).
|
||||
|
||||
Runs `prisma migrate deploy` and handles standard recovery paths
|
||||
(P3005 baseline, P3009/P3018 idempotent errors). Critically, it does
|
||||
|
|
@ -692,7 +722,7 @@ class ProxyExtrasDBManager:
|
|||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
try:
|
||||
for attempt in range(4):
|
||||
for attempt in range(_MIGRATE_DEPLOY_ATTEMPTS):
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
|
|
@ -807,16 +837,36 @@ class ProxyExtrasDBManager:
|
|||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
transient = _transient_deploy_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:
|
||||
raise RuntimeError(
|
||||
f"Database migration failed after "
|
||||
f"{_MIGRATE_DEPLOY_ATTEMPTS} attempts on {transient}. "
|
||||
"Check database connectivity and load."
|
||||
f"\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
logger.info(
|
||||
"prisma migrate deploy attempt %s failed on %s, retrying. "
|
||||
"Prisma error:\n%s",
|
||||
attempt + 1,
|
||||
transient,
|
||||
stderr,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
continue
|
||||
|
||||
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."
|
||||
f"Database migration failed after {_MIGRATE_DEPLOY_ATTEMPTS} "
|
||||
"attempts (retry loop exhausted by timeouts or repeated "
|
||||
"idempotent-recovery continues). Check database connectivity, "
|
||||
"load, and _prisma_migrations ledger state."
|
||||
)
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
|
@ -864,10 +914,11 @@ class ProxyExtrasDBManager:
|
|||
|
||||
Args:
|
||||
use_migrate: Whether to use prisma migrate instead of db push
|
||||
use_v2_resolver: Opt into the v2 migration resolver (safer during
|
||||
use_v2_resolver: Run the v2 migration resolver (safer during
|
||||
rolling deploys; does not run the diff-and-force recovery
|
||||
that causes schema thrashing). Defaults to False for
|
||||
backwards compatibility.
|
||||
that causes schema thrashing). Defaults to False here so
|
||||
direct callers keep the old behavior; the proxy CLI passes
|
||||
True, so the proxy's runtime default is v2.
|
||||
|
||||
Returns:
|
||||
bool: True if setup was successful, False otherwise
|
||||
|
|
@ -885,7 +936,7 @@ class ProxyExtrasDBManager:
|
|||
@staticmethod
|
||||
def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
|
||||
if use_v2_resolver:
|
||||
logger.info("Using v2 migration resolver (--use_v2_migration_resolver)")
|
||||
logger.info("Using v2 migration resolver")
|
||||
return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate)
|
||||
|
||||
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
|
||||
|
|
|
|||
|
|
@ -913,13 +913,14 @@ class ProxyInitializationHelpers:
|
|||
envvar="ENFORCE_PRISMA_MIGRATION_CHECK",
|
||||
)
|
||||
@click.option(
|
||||
"--use_v2_migration_resolver",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
"--use_v2_migration_resolver/--use_legacy_migration_resolver",
|
||||
default=True,
|
||||
help=(
|
||||
"Opt into the v2 migration resolver. Avoids the diff-and-force recovery "
|
||||
"path that can cause schema thrashing during rolling deploys where two "
|
||||
"LiteLLM versions contend for the same DB. Default is the v1 resolver."
|
||||
"Which database migration resolver to run at startup. The default v2 "
|
||||
"resolver avoids the diff-and-force recovery path that can cause schema "
|
||||
"thrashing during rolling deploys where two LiteLLM versions contend for "
|
||||
"the same DB. Pass --use_legacy_migration_resolver, or set "
|
||||
"USE_V2_MIGRATION_RESOLVER=false, to fall back to v1."
|
||||
),
|
||||
envvar="USE_V2_MIGRATION_RESOLVER",
|
||||
)
|
||||
|
|
@ -1310,10 +1311,11 @@ def run_server(
|
|||
else:
|
||||
if not use_v2_migration_resolver:
|
||||
print(
|
||||
"\033[1;33mLiteLLM Proxy: Using default (v1) migration resolver. "
|
||||
"If your deployment has seen schema thrashing during rolling "
|
||||
"deploys, try --use_v2_migration_resolver (safer: avoids the "
|
||||
"diff-and-force recovery that caused the thrash).\033[0m"
|
||||
"\033[1;33mLiteLLM Proxy: Using the legacy (v1) migration resolver. "
|
||||
"The default v2 resolver is safer: it avoids the diff-and-force "
|
||||
"recovery that caused schema thrashing during rolling deploys. "
|
||||
"Remove --use_legacy_migration_resolver / "
|
||||
"USE_V2_MIGRATION_RESOLVER=false to switch back to it.\033[0m"
|
||||
)
|
||||
try:
|
||||
setup_ok: Final = PrismaManager.setup_database(
|
||||
|
|
|
|||
|
|
@ -1,15 +1,26 @@
|
|||
"""Regression tests for ProxyExtrasDBManager v2 migration resolver.
|
||||
"""Regression tests for ProxyExtrasDBManager's v2 migration resolver.
|
||||
|
||||
The v2 resolver is opt-in via `--use_v2_migration_resolver` / the
|
||||
`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1
|
||||
(default) behavior is unchanged from pre-fix.
|
||||
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.
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0,
|
||||
os.path.abspath(
|
||||
os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras")
|
||||
),
|
||||
)
|
||||
|
||||
from litellm_proxy_extras.utils import (
|
||||
ProxyExtrasDBManager,
|
||||
_max_migration_timestamp,
|
||||
|
|
@ -240,3 +251,222 @@ 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_STDERR = (
|
||||
"Error: ERROR: deadlock detected\n"
|
||||
"DETAIL: Process 277 waits for ExclusiveLock on advisory lock "
|
||||
"[17556,0,72707369,1]; blocked by process 278.\n"
|
||||
"Process 278 waits for ShareLock on virtual transaction 3/1041; "
|
||||
"blocked by process 277."
|
||||
)
|
||||
|
||||
|
||||
class _DeployApplied:
|
||||
stdout = "All migrations have been successfully applied."
|
||||
stderr = ""
|
||||
returncode = 0
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
deploys = {"n": 0}
|
||||
|
||||
def _run(*args, **kwargs):
|
||||
cmd = args[0] if args else kwargs.get("args", [])
|
||||
if list(cmd)[-2:] == ["migrate", "deploy"]:
|
||||
deploys["n"] += 1
|
||||
return deploy_side_effect(deploys["n"], cmd)
|
||||
return _DeployApplied()
|
||||
|
||||
return _run, deploys
|
||||
|
||||
|
||||
def _prepare_v2_resolver(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 *_a, **_k: None)
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
if n == 1:
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=cmd, stderr=_DEADLOCK_STDERR, output=""
|
||||
)
|
||||
return _DeployApplied()
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
assert ok is True
|
||||
assert deploys["n"] == 2, "the deadlocked deploy must be retried, not raised"
|
||||
|
||||
|
||||
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."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=cmd, stderr=_DEADLOCK_STDERR, output=""
|
||||
)
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
with pytest.raises(RuntimeError, match="after 4 attempts"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
assert deploys["n"] == 4
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"stderr",
|
||||
[
|
||||
"Error: P1001: Can't reach database server at `db`:`5432`",
|
||||
"Error: P1002: The database server was reached but timed out.",
|
||||
],
|
||||
)
|
||||
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.
|
||||
"""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
if n == 1:
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=cmd, stderr=stderr, output=""
|
||||
)
|
||||
return _DeployApplied()
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
assert ok is True
|
||||
assert deploys["n"] == 2, "an unreachable database must be retried, not raised"
|
||||
|
||||
|
||||
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."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd=cmd,
|
||||
stderr="Error: P1001: Can't reach database server at `db`:`5432`",
|
||||
output="",
|
||||
)
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
with pytest.raises(RuntimeError, match="after 4 attempts"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
assert deploys["n"] == 4
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
stderr = "Error: P1001: Can't reach database server at `wrong`:`5432`"
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=cmd, stderr=stderr, output=""
|
||||
)
|
||||
|
||||
run, _ = _deploy_only(_side_effect)
|
||||
with caplog.at_level("INFO", logger="litellm_proxy_extras"):
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
ProxyExtrasDBManager.setup_database(
|
||||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
|
||||
assert "P1001" in str(exc_info.value)
|
||||
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."""
|
||||
_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'
|
||||
)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=cmd, stderr=stderr, output=""
|
||||
)
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
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
|
||||
|
|
@ -305,14 +305,16 @@ def _run_proxy_server_smoke_test(extra_proxy_args=None):
|
|||
|
||||
|
||||
def test_litellm_proxy_server_config_no_general_settings():
|
||||
"""Exercises the default (v1) migration resolver."""
|
||||
"""Exercises the default (v2) migration resolver."""
|
||||
_run_proxy_server_smoke_test()
|
||||
|
||||
|
||||
def test_litellm_proxy_server_config_no_general_settings_v2_resolver():
|
||||
"""Exercises the opt-in v2 migration resolver.
|
||||
def test_litellm_proxy_server_config_no_general_settings_legacy_resolver():
|
||||
"""Exercises the legacy (v1) migration resolver against a real database.
|
||||
|
||||
Runs in a separate CI job against a local Postgres to avoid collisions
|
||||
with the v1 variant when they share a database.
|
||||
v2 is the default, so the no-arg test above already covers it. This one is
|
||||
the only place the v1 opt-out gets real-DB migration plus proxy-boot
|
||||
coverage, and it runs in a separate CI job against its own Postgres to
|
||||
avoid collisions with the default variant.
|
||||
"""
|
||||
_run_proxy_server_smoke_test(extra_proxy_args=["--use_v2_migration_resolver"])
|
||||
_run_proxy_server_smoke_test(extra_proxy_args=["--use_legacy_migration_resolver"])
|
||||
|
|
|
|||
|
|
@ -1787,7 +1787,7 @@ class TestRunServerDbSetup:
|
|||
# use_prisma_db_push should be False (default), so use_migrate should be True
|
||||
run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
|
||||
mock_setup_database.assert_called_with(
|
||||
use_migrate=True, use_v2_resolver=False
|
||||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
|
||||
# Reset mocks
|
||||
|
|
@ -1802,7 +1802,7 @@ class TestRunServerDbSetup:
|
|||
standalone_mode=False,
|
||||
)
|
||||
mock_setup_database.assert_called_with(
|
||||
use_migrate=False, use_v2_resolver=False
|
||||
use_migrate=False, use_v2_resolver=True
|
||||
)
|
||||
|
||||
@patch("subprocess.run")
|
||||
|
|
@ -1869,7 +1869,7 @@ class TestRunServerDbSetup:
|
|||
)
|
||||
assert exc_info.value.code == 1
|
||||
mock_setup_database.assert_called_once_with(
|
||||
use_migrate=True, use_v2_resolver=False
|
||||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
|
||||
@patch("subprocess.run")
|
||||
|
|
@ -1981,6 +1981,83 @@ class TestRunServerDbSetup:
|
|||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"argv, env_value, expected_v2",
|
||||
[
|
||||
([], None, True),
|
||||
(["--use_v2_migration_resolver"], None, True),
|
||||
(["--use_legacy_migration_resolver"], None, False),
|
||||
([], "false", False),
|
||||
([], "true", True),
|
||||
(["--use_v2_migration_resolver"], "false", True),
|
||||
],
|
||||
)
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_migration_resolver_default_and_opt_out(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
argv,
|
||||
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.
|
||||
"""
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
||||
mock_subprocess_run.return_value = MagicMock(returncode=0)
|
||||
mock_should_update_schema.return_value = True
|
||||
mock_setup_database.return_value = True
|
||||
|
||||
mock_proxy_module = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k
|
||||
not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER")
|
||||
}
|
||||
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
|
||||
if env_value is not None:
|
||||
clean_env["USE_V2_MIGRATION_RESOLVER"] = env_value
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"proxy_server": mock_proxy_module,
|
||||
"litellm.proxy.proxy_server": mock_proxy_module,
|
||||
},
|
||||
),
|
||||
):
|
||||
run_server.main(
|
||||
["--local", "--skip_server_startup", *argv], standalone_mode=False
|
||||
)
|
||||
|
||||
mock_setup_database.assert_called_once_with(
|
||||
use_migrate=True, use_v2_resolver=expected_v2
|
||||
)
|
||||
|
||||
|
||||
# --- Module-level helpers for worker startup hook tests ---
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue