mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
fix(fallback): edit the config file's fallback list when it owns the key
A delete of a rule only the config file holds answers the config-owned 400 again instead of 404, and restating a rule the file already has answers 200 without copying the file's list into the database row. Adds a writer-read regression test and an integration test against a Redis whose ACL refuses DEL
This commit is contained in:
parent
42d54a9d98
commit
94cbf51e11
3 changed files with 111 additions and 28 deletions
|
|
@ -62,10 +62,31 @@ async def _stored_router_settings(prisma_client: PrismaClient) -> Mapping[str, J
|
|||
return _STORED_ROUTER_SETTINGS.validate_python(row.param_value)
|
||||
|
||||
|
||||
def _reject_config_owned_rules(proxy_config: "ProxyConfig", fallback_key: str, rules: Sequence[StoredFallback]) -> None:
|
||||
def _editable_rules(
|
||||
proxy_config: "ProxyConfig", stored_settings: Mapping[str, JsonValue], fallback_key: str
|
||||
) -> list[StoredFallback]:
|
||||
file_settings: Final = proxy_config.router_settings
|
||||
rules: Final = (
|
||||
file_settings.config_value(fallback_key)
|
||||
if file_settings.owned_by_config(fallback_key)
|
||||
else stored_settings.get(fallback_key)
|
||||
)
|
||||
return _STORED_FALLBACKS.validate_python(rules or [])
|
||||
|
||||
|
||||
async def _write_rules(
|
||||
prisma_client: PrismaClient,
|
||||
proxy_config: "ProxyConfig",
|
||||
stored_settings: Mapping[str, JsonValue],
|
||||
fallback_key: str,
|
||||
rules: Sequence[StoredFallback],
|
||||
) -> None:
|
||||
proxy_config.reject_config_owned_writes(
|
||||
section_name=ROUTER_SETTINGS_PARAM, changed_keys={fallback_key: _JSON_RULES.validate_python(rules)}
|
||||
)
|
||||
if proxy_config.router_settings.owned_by_config(fallback_key):
|
||||
return
|
||||
await _persist_router_settings(prisma_client, {**stored_settings, fallback_key: rules})
|
||||
|
||||
|
||||
async def _persist_router_settings(prisma_client: PrismaClient, router_settings: Mapping[str, object]) -> None:
|
||||
|
|
@ -174,7 +195,7 @@ async def create_fallback(
|
|||
fallback_key = "content_policy_fallbacks"
|
||||
|
||||
# Get existing fallbacks
|
||||
existing_fallbacks: Final = _STORED_FALLBACKS.validate_python(router_settings.get(fallback_key) or [])
|
||||
existing_fallbacks: Final = _editable_rules(proxy_config, router_settings, fallback_key)
|
||||
|
||||
# Update or add the fallback configuration
|
||||
fallback_updated = False
|
||||
|
|
@ -189,9 +210,7 @@ async def create_fallback(
|
|||
# Add new fallback
|
||||
existing_fallbacks.append({data.model: data.fallback_models})
|
||||
|
||||
_reject_config_owned_rules(proxy_config, fallback_key, existing_fallbacks)
|
||||
|
||||
await _persist_router_settings(prisma_client, {**router_settings, fallback_key: existing_fallbacks})
|
||||
await _write_rules(prisma_client, proxy_config, router_settings, fallback_key, existing_fallbacks)
|
||||
|
||||
# Update the in-memory router configuration
|
||||
setattr(llm_router, fallback_key, existing_fallbacks)
|
||||
|
|
@ -329,7 +348,7 @@ async def delete_fallback(
|
|||
fallback_key = "content_policy_fallbacks"
|
||||
|
||||
# Get existing fallbacks
|
||||
existing_fallbacks: Final = _STORED_FALLBACKS.validate_python(router_settings.get(fallback_key) or [])
|
||||
existing_fallbacks: Final = _editable_rules(proxy_config, router_settings, fallback_key)
|
||||
|
||||
updated_fallbacks: Final = [rule for rule in existing_fallbacks if not _rule_covers(rule, model)]
|
||||
|
||||
|
|
@ -339,9 +358,7 @@ async def delete_fallback(
|
|||
detail={"error": f"No {fallback_type} fallbacks configured for model '{model}'"},
|
||||
)
|
||||
|
||||
_reject_config_owned_rules(proxy_config, fallback_key, updated_fallbacks)
|
||||
|
||||
await _persist_router_settings(prisma_client, {**router_settings, fallback_key: updated_fallbacks})
|
||||
await _write_rules(prisma_client, proxy_config, router_settings, fallback_key, updated_fallbacks)
|
||||
|
||||
# Update the in-memory router configuration
|
||||
setattr(llm_router, fallback_key, updated_fallbacks)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from functools import partial
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
|
|
@ -10,6 +11,7 @@ import httpx
|
|||
import psutil
|
||||
import pytest
|
||||
from pydantic import JsonValue
|
||||
from redis import Redis
|
||||
|
||||
from tests.integration._support.client import (
|
||||
Gateway,
|
||||
|
|
@ -27,7 +29,7 @@ from tests.integration._support.process import (
|
|||
setup_only_proxy_run,
|
||||
worker_configuration,
|
||||
)
|
||||
from tests.integration._support.redis_process import owned_redis
|
||||
from tests.integration._support.redis_process import OwnedRedis, owned_redis
|
||||
|
||||
pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state]
|
||||
|
||||
|
|
@ -385,6 +387,39 @@ def test_fallback_write_survives_a_redis_outage(gateway: Gateway, tmp_path: Path
|
|||
cache.start()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _refusing_cache_deletes(cache: OwnedRedis) -> Iterator[None]:
|
||||
with Redis(host=cache.host, port=cache.port) as client:
|
||||
client.acl_setuser("default", enabled=True, commands=["-del", "-unlink"])
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
client.acl_setuser("default", enabled=True, commands=["+del", "+unlink"])
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("migrated_database")
|
||||
def test_consecutive_creates_keep_every_rule_when_redis_refuses_the_cache_eviction(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with owned_redis(tmp_path) as cache:
|
||||
overrides: Final = {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)}
|
||||
with owned_proxy_process(gateway, tmp_path, overrides, database_setup=()) as owned:
|
||||
with owned.gateway.scenario() as scenario:
|
||||
target: Final = scenario.model(model=f"openai/x3t-{uuid.uuid4().hex}")
|
||||
first: Final = scenario.model(model=f"openai/x3a-{uuid.uuid4().hex}")
|
||||
second: Final = scenario.model(model=f"openai/x3b-{uuid.uuid4().hex}")
|
||||
for model in (target, first, second):
|
||||
_wait_until_served(owned.gateway, model)
|
||||
with _refusing_cache_deletes(cache):
|
||||
created: Final = tuple(
|
||||
_create_fallback(owned.gateway, scenario, model, [target]) for model in (first, second)
|
||||
)
|
||||
assert all(response.status_code == 200 for response in created), [r.text for r in created]
|
||||
covered: Final = _covered_models(_stored_fallbacks())
|
||||
assert {first, second} <= covered, covered
|
||||
assert EVICTION_WARNING in owned.log.read_text(), owned.log.read_text()[-3000:]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("migrated_database")
|
||||
def test_fallback_write_survives_a_worker_kill(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {}, workers=2, database_setup=()) as owned:
|
||||
|
|
|
|||
|
|
@ -175,11 +175,11 @@ class TestCreateFallbackForTeamScopedModels:
|
|||
return client
|
||||
|
||||
@pytest.fixture
|
||||
def proxy_config(self) -> MagicMock:
|
||||
return MagicMock()
|
||||
def proxy_config(self) -> ProxyConfig:
|
||||
return ProxyConfig()
|
||||
|
||||
async def _create(
|
||||
self, request: FallbackCreateRequest, router: Router, prisma_client: MagicMock, proxy_config: MagicMock
|
||||
self, request: FallbackCreateRequest, router: Router, prisma_client: MagicMock, proxy_config: ProxyConfig
|
||||
) -> FallbackResponse:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router", router),
|
||||
|
|
@ -258,27 +258,29 @@ def _config_row(router_settings: RouterSettings) -> SimpleNamespace:
|
|||
class _StoredRouterSettings:
|
||||
"""The LiteLLM_Config router_settings row, with the proxy objects that read it through the config cache"""
|
||||
|
||||
def __init__(self, router_settings: RouterSettings | None, config_file: RouterSettings | None = None) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
router_settings: RouterSettings | None,
|
||||
config_file: RouterSettings | None = None,
|
||||
*,
|
||||
lagging_replica: bool = False,
|
||||
) -> None:
|
||||
self.row: SimpleNamespace | None = None if router_settings is None else _config_row(router_settings)
|
||||
self.config_file: Final = config_file or {}
|
||||
self.prisma_client: Final = MagicMock()
|
||||
self.prisma_client.get_generic_data = AsyncMock(side_effect=lambda **_: self.row)
|
||||
self.prisma_client.writer_db = self.prisma_client.db
|
||||
self.prisma_client.db.litellm_config.find_unique = AsyncMock(side_effect=lambda **_: copy.deepcopy(self.row))
|
||||
self.prisma_client.writer_db.litellm_config.find_unique = AsyncMock(
|
||||
side_effect=lambda **_: copy.deepcopy(self.row)
|
||||
)
|
||||
self.prisma_client.db.litellm_config.find_unique = AsyncMock(
|
||||
side_effect=lambda **_: None if lagging_replica else copy.deepcopy(self.row)
|
||||
)
|
||||
self.prisma_client.db.litellm_config.upsert = AsyncMock(side_effect=self._upsert)
|
||||
self.proxy_config: Final = MagicMock() if config_file is None else ProxyConfig()
|
||||
if config_file is not None:
|
||||
self.proxy_config.router_settings.load_yaml(_ROUTER_SETTINGS_JSON.validate_python(config_file))
|
||||
self.proxy_config.get_config = AsyncMock(side_effect=self._get_config)
|
||||
self.proxy_config: Final = ProxyConfig()
|
||||
self.proxy_config.router_settings.load_yaml(_ROUTER_SETTINGS_JSON.validate_python(config_file or {}))
|
||||
|
||||
async def _upsert(self, where: dict[str, str], data: dict[str, dict[str, str]]) -> None:
|
||||
self.row = _config_row(json.loads(data["update"]["param_value"]))
|
||||
|
||||
async def _get_config(self) -> dict[str, RouterSettings]:
|
||||
row: Final = await get_config_param(self.prisma_client, "router_settings")
|
||||
stored: Final = copy.deepcopy(row.param_value) if row is not None else {}
|
||||
return {"router_settings": {**stored, **self.config_file}}
|
||||
|
||||
def written_by_another_instance(self, router_settings: RouterSettings) -> None:
|
||||
self.row = _config_row(router_settings)
|
||||
|
||||
|
|
@ -341,6 +343,13 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
):
|
||||
return await delete_fallback(model, "general", MagicMock())
|
||||
|
||||
async def test_a_create_reads_the_rules_from_the_writer_while_the_replica_lags(self) -> None:
|
||||
stored: Final = _StoredRouterSettings({"fallbacks": [TEAM_RULE]}, lagging_replica=True)
|
||||
|
||||
await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored)
|
||||
|
||||
assert stored.stored_fallbacks() == [TEAM_RULE, GATEWAY_RULE]
|
||||
|
||||
async def test_a_second_create_keeps_the_first_rule(self) -> None:
|
||||
stored: Final = _StoredRouterSettings(None)
|
||||
|
||||
|
|
@ -490,6 +499,28 @@ class TestFallbackWritesHonourTheConfigFile:
|
|||
stored.prisma_client.db.litellm_config.upsert.assert_not_awaited()
|
||||
assert router.fallbacks == [CONFIG_RULE]
|
||||
|
||||
async def test_a_delete_of_a_rule_only_the_config_file_holds_is_refused(self, router: Router) -> None:
|
||||
stored: Final = _StoredRouterSettings(None, config_file={"fallbacks": [CONFIG_RULE]})
|
||||
router.fallbacks = [CONFIG_RULE]
|
||||
|
||||
with pytest.raises(HTTPException) as refused:
|
||||
await self._delete("gpt-5.4-mini", router, stored)
|
||||
|
||||
assert refused.value.status_code == 400
|
||||
assert "key 'fallbacks' is set in the config file" in str(refused.value.detail)
|
||||
assert stored.row is None
|
||||
assert router.fallbacks == [CONFIG_RULE]
|
||||
|
||||
async def test_a_create_that_restates_one_of_several_config_rules_writes_nothing(self, router: Router) -> None:
|
||||
stored: Final = _StoredRouterSettings(None, config_file={"fallbacks": [CONFIG_RULE, CONFIG_WINDOW_RULE]})
|
||||
request: Final = FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"])
|
||||
|
||||
response: Final = await self._create(request, router, stored)
|
||||
|
||||
assert response.fallback_models == ["team-fallback"]
|
||||
assert stored.row is None
|
||||
assert router.fallbacks == [CONFIG_RULE, CONFIG_WINDOW_RULE]
|
||||
|
||||
async def test_a_create_that_restates_the_config_rule_is_accepted(
|
||||
self, router: Router, stored: _StoredRouterSettings
|
||||
) -> None:
|
||||
|
|
@ -557,7 +588,7 @@ class TestCreateFallback:
|
|||
@pytest.fixture
|
||||
def mock_proxy_config(self):
|
||||
"""Create a mock proxy config"""
|
||||
return MagicMock()
|
||||
return ProxyConfig()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_user_api_key_dict(self):
|
||||
|
|
@ -866,7 +897,7 @@ class TestDeleteFallback:
|
|||
@pytest.fixture
|
||||
def mock_proxy_config(self):
|
||||
"""Create a mock proxy config"""
|
||||
return MagicMock()
|
||||
return ProxyConfig()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_user_api_key_dict(self):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue