scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

Merge pull request #42105 from BerriAI/litellm_flip_v2_migration_resolver_default

Browse source 
feat(proxy): default to the v2 migration resolver
This commit is contained in:
yuneng-jiang 2026-09-21 10:05:59 -07:00 • committed by GitHub
commit da76ba83ad
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 165 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
.circleci/config.yml
View file

@ -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,11 @@ jobs:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Run v2 migration resolver proxy smoke test
name: Run both migration resolvers against Postgres
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 \
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver
helm_chart_testing:
machine:

61
litellm/proxy/proxy_cli.py
View file

@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final
import click
import httpx
from click.core import ParameterSource
from dotenv import load_dotenv
from pydantic import BaseModel, ConfigDict
@ -181,6 +182,23 @@ def append_query_params(url: str | None, params: dict) -> str:
return modified_url
def resolve_v2_migration_resolver(*, use_legacy_flag: bool, env_value: str | None) -> bool:
from litellm_proxy_extras.utils import str_to_bool
if use_legacy_flag:
return False
if env_value is None:
return True
return bool(str_to_bool(env_value))
def deprecated_v2_flag_passed_on_cli() -> bool:
ctx: Final = click.get_current_context(silent=True)
if ctx is None:
return False
return ctx.get_parameter_source("use_v2_migration_resolver") is ParameterSource.COMMANDLINE
class ProxyInitializationHelpers:
@staticmethod
def _echo_litellm_version():
@ -932,12 +950,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 +1035,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 +1377,29 @@ 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,
env_value=os.getenv("USE_V2_MIGRATION_RESOLVER"),
)
if deprecated_v2_flag_passed_on_cli() 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, because 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

12
tests/local_testing/test_basic_python_version.py
View file

@ -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.
Runs after the default variant in the CI job that provides a local
Postgres, so both resolvers get real-database proxy-boot coverage.
"""
_run_proxy_server_smoke_test(extra_proxy_args=["--use_v2_migration_resolver"])
_run_proxy_server_smoke_test(extra_proxy_args=["--use_legacy_migration_resolver"])

109
tests/test_litellm/proxy/test_proxy_cli.py
View file

@ -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")
@ -2203,12 +2203,13 @@ class TestRunServerDbSetup:
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
capsys,
):
"""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
@ -2248,6 +2249,100 @@ class TestRunServerDbSetup:
mock_setup_database.assert_called_once_with(
use_migrate=True, use_v2_resolver=True
)
assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out
@pytest.mark.parametrize(
"use_legacy_flag, env_value, expected",
[
(False, None, True),
(False, "true", True),
(False, "false", False),
(True, None, False),
(True, "true", False),
],
ids=[
"unset-env-defaults-to-v2",
"env-true-selects-v2",
"env-false-selects-v1",
"legacy-flag-selects-v1",
"legacy-flag-beats-env-true",
],
)
def test_resolve_v2_migration_resolver(self, use_legacy_flag, env_value, expected):
from litellm.proxy.proxy_cli import resolve_v2_migration_resolver
assert (
resolve_v2_migration_resolver(
use_legacy_flag=use_legacy_flag, env_value=env_value
)
is expected
)
def test_deprecated_v2_flag_not_reported_outside_a_cli_invocation(self):
from litellm.proxy.proxy_cli import deprecated_v2_flag_passed_on_cli
assert deprecated_v2_flag_passed_on_cli() is False
@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_legacy_resolver_flag_reaches_database_setup(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
):
"""--use_legacy_migration_resolver must reach the database setup call.
The resolver decision itself is covered mock-free above; this is the
one wiring check that the flag is threaded through run_server.
"""
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"
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",
"--use_legacy_migration_resolver",
],
standalone_mode=False,
)
mock_setup_database.assert_called_once_with(
use_migrate=True, use_v2_resolver=False
)
# --- Module-level helpers for worker startup hook tests ---