mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
feat(proxy): default to the v2 migration resolver
The migrations Job entrypoint (migrations/run.py) has defaulted to v2 with USE_V2_MIGRATION_RESOLVER=false as the opt-out, and the Helm chart documents that knob. Proxy startup still defaulted to v1, so the two paths disagreed about which resolver a deployment runs. Proxy startup now resolves the same way: v2 unless USE_V2_MIGRATION_RESOLVER is false or --use_legacy_migration_resolver is passed. - --use_v2_migration_resolver stays accepted as a no-op that warns, so existing commands and Helm values do not fail on an unknown option. - The dedicated Postgres smoke-test job is repointed at the legacy resolver so v1 keeps real-DB proxy-boot coverage, and the two jobs that deselected it by name are updated to match the rename. #39178 reverted an earlier flip because two replicas sharing a database deadlocked (40P01 / P3018) with neither answering /health/liveliness. That contention is what #40932 coordinates, which is why this builds on it.
This commit is contained in:
Yuneng Jiang
parent
e34fd201b3
commit
fac518dbf7
4 changed files with 135 additions and 25 deletions
|
|
@ -1508,7 +1508,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:
|
||||
|
|
@ -1532,7 +1532,7 @@ 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:
|
||||
docker:
|
||||
|
|
@ -1561,10 +1561,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:
|
||||
|
|
|
|||
|
|
@ -181,6 +181,14 @@ def append_query_params(url: str | None, params: dict) -> str:
|
|||
return modified_url
|
||||
|
||||
|
||||
def resolve_v2_migration_resolver(*, use_legacy_flag: bool) -> bool:
|
||||
from litellm_proxy_extras.utils import str_to_bool
|
||||
|
||||
if use_legacy_flag:
|
||||
return False
|
||||
return bool(str_to_bool(os.getenv("USE_V2_MIGRATION_RESOLVER", "true")))
|
||||
|
||||
|
||||
class ProxyInitializationHelpers:
|
||||
@staticmethod
|
||||
def _echo_litellm_version():
|
||||
|
|
@ -932,12 +940,24 @@ class ProxyInitializationHelpers:
|
|||
is_flag=True,
|
||||
default=False,
|
||||
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."
|
||||
"Deprecated and ignored: the v2 migration resolver is now the default, "
|
||||
"so this flag has no effect. It is still accepted so existing commands "
|
||||
"keep working. Pass --use_legacy_migration_resolver, or set "
|
||||
"USE_V2_MIGRATION_RESOLVER=false, to opt back into v1."
|
||||
),
|
||||
envvar="USE_V2_MIGRATION_RESOLVER",
|
||||
)
|
||||
@click.option(
|
||||
"--use_legacy_migration_resolver",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help=(
|
||||
"Fall back to the legacy v1 migration resolver. By default the proxy "
|
||||
"uses the v2 resolver, which avoids the diff-and-force recovery path "
|
||||
"that can cause schema thrashing during rolling deploys where two "
|
||||
"LiteLLM versions contend for the same DB."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--reload",
|
||||
is_flag=True,
|
||||
|
|
@ -1005,6 +1025,7 @@ def run_server(
|
|||
limit_concurrency: int | None,
|
||||
enforce_prisma_migration_check: bool,
|
||||
use_v2_migration_resolver: bool,
|
||||
use_legacy_migration_resolver: bool,
|
||||
reload: bool,
|
||||
prometheus_metrics_port: int | None,
|
||||
):
|
||||
|
|
@ -1346,17 +1367,28 @@ def run_server(
|
|||
if should_update_prisma_schema(general_settings.get("disable_prisma_schema_update")) is False:
|
||||
check_prisma_schema_diff(db_url=None)
|
||||
else:
|
||||
if not use_v2_migration_resolver:
|
||||
use_v2_resolver: Final = resolve_v2_migration_resolver(
|
||||
use_legacy_flag=use_legacy_migration_resolver
|
||||
)
|
||||
if use_v2_migration_resolver and use_v2_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: --use_v2_migration_resolver is "
|
||||
"deprecated and has no effect \u2014 the v2 migration resolver "
|
||||
"is now the default. You can safely remove it. To opt back "
|
||||
"into the legacy v1 resolver, pass "
|
||||
"--use_legacy_migration_resolver.\033[0m"
|
||||
)
|
||||
if not use_v2_resolver:
|
||||
print(
|
||||
"\033[1;33mLiteLLM Proxy: Using the legacy (v1) migration "
|
||||
"resolver. It performs the diff-and-force recovery that can "
|
||||
"cause schema thrashing during rolling deploys where two "
|
||||
"LiteLLM versions contend for the same DB.\033[0m"
|
||||
)
|
||||
try:
|
||||
setup_ok: Final = PrismaManager.setup_database(
|
||||
use_migrate=not use_prisma_db_push,
|
||||
use_v2_resolver=use_v2_migration_resolver,
|
||||
use_v2_resolver=use_v2_resolver,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
# Raised on unrecoverable migration errors: the v2
|
||||
|
|
|
|||
|
|
@ -305,14 +305,14 @@ 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 opt-out legacy (v1) migration resolver.
|
||||
|
||||
Runs in a separate CI job against a local Postgres to avoid collisions
|
||||
with the v1 variant when they share a database.
|
||||
with the default variant when they share a database.
|
||||
"""
|
||||
_run_proxy_server_smoke_test(extra_proxy_args=["--use_v2_migration_resolver"])
|
||||
_run_proxy_server_smoke_test(extra_proxy_args=["--use_legacy_migration_resolver"])
|
||||
|
|
|
|||
|
|
@ -1995,7 +1995,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
|
||||
|
|
@ -2010,7 +2010,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("atexit.register")
|
||||
|
|
@ -2070,7 +2070,7 @@ class TestRunServerDbSetup:
|
|||
|
||||
assert "prisma CLI is neither on PATH" not in capsys.readouterr().out
|
||||
mock_setup_database.assert_called_once_with(
|
||||
use_migrate=True, use_v2_resolver=False
|
||||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
|
||||
@patch("subprocess.run")
|
||||
|
|
@ -2137,7 +2137,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")
|
||||
|
|
@ -2204,11 +2204,11 @@ class TestRunServerDbSetup:
|
|||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
):
|
||||
"""USE_V2_MIGRATION_RESOLVER must select the v2 resolver.
|
||||
"""USE_V2_MIGRATION_RESOLVER=true must select the v2 resolver.
|
||||
|
||||
The Helm migrations Job runs `python litellm/proxy/prisma_migration.py`,
|
||||
which calls run_server with a fixed argv, so a deployment has no way to
|
||||
pass --use_v2_migration_resolver and an env var is the only route in.
|
||||
which calls run_server with a fixed argv, so a deployment reaches the
|
||||
resolver through the env var rather than a CLI flag.
|
||||
"""
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
||||
|
|
@ -2249,6 +2249,84 @@ class TestRunServerDbSetup:
|
|||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"argv_extra, env_extra, expected_v2",
|
||||
[
|
||||
([], {}, True),
|
||||
([], {"USE_V2_MIGRATION_RESOLVER": "false"}, False),
|
||||
(["--use_legacy_migration_resolver"], {}, False),
|
||||
(
|
||||
["--use_legacy_migration_resolver"],
|
||||
{"USE_V2_MIGRATION_RESOLVER": "true"},
|
||||
False,
|
||||
),
|
||||
(["--use_v2_migration_resolver"], {}, True),
|
||||
],
|
||||
ids=[
|
||||
"default-is-v2",
|
||||
"env-false-opts-out",
|
||||
"legacy-flag-opts-out",
|
||||
"legacy-flag-beats-env-true",
|
||||
"deprecated-v2-flag-still-accepted",
|
||||
],
|
||||
)
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
def test_migration_resolver_selection(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
argv_extra,
|
||||
env_extra,
|
||||
expected_v2,
|
||||
):
|
||||
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"
|
||||
clean_env.update(env_extra)
|
||||
|
||||
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_extra],
|
||||
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