mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(proxy): re-read UI settings on every config reload
The persisted UI settings were read once at startup, so a proxy admin flipping a runtime flag through PATCH /update/ui_settings only changed the pod that served the request. Every other pod kept serving the old value until it restarted. add_deployment, the reload the scheduler runs every 30s, now re-reads the row and applies the runtime flags before it takes the model reconcile lock, so a change made through one pod reaches the rest within one reload interval. The startup hook and the two settings endpoints share that helper instead of each repeating the flag copy. Claude-Session: https://claude.ai/code/session_018PUCupsaarVLJy4iDFx256
This commit is contained in:
parent
7743c3b4e2
commit
00ceffa827
4 changed files with 175 additions and 49 deletions
|
|
@ -650,6 +650,9 @@ from litellm.proxy.types_utils.utils import get_instance_fn
|
|||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
router as ui_crud_endpoints_router,
|
||||
)
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
sync_ui_settings_to_general_settings,
|
||||
)
|
||||
from litellm.proxy.ui_crud_endpoints.user_banner_endpoints import (
|
||||
router as user_banner_endpoints_router,
|
||||
)
|
||||
|
|
@ -1690,10 +1693,6 @@ class _SSOConfigRow(Protocol):
|
|||
sso_settings: MutableMapping[str, object]
|
||||
|
||||
|
||||
class _UISettingsRow(Protocol):
|
||||
ui_settings: Mapping[str, object] | str | None
|
||||
|
||||
|
||||
class _InvitationLinkRow(Protocol):
|
||||
user_id: str
|
||||
expires_at: datetime
|
||||
|
|
@ -7092,7 +7091,12 @@ class ProxyConfig:
|
|||
Returns what the reconcile saw, captured before the lock is released so a
|
||||
caller's verdict cannot be corrupted by the next reconcile's own in-flight
|
||||
window. See ReconcileOutcome.
|
||||
|
||||
Also re-reads the UI settings that back runtime flags. That runs before the lock, so a
|
||||
setting written through one pod reaches the others without waiting on a model reconcile.
|
||||
"""
|
||||
await sync_ui_settings_to_general_settings(prisma_client)
|
||||
|
||||
async with MODEL_RECONCILE_LOCK:
|
||||
return await self._add_deployment_locked(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
|
||||
|
||||
|
|
@ -9321,35 +9325,12 @@ class ProxyStartupEvent:
|
|||
|
||||
@classmethod
|
||||
async def _sync_ui_settings_to_general_settings(cls):
|
||||
"""
|
||||
Load persisted UI settings from the database and sync runtime flags
|
||||
into general_settings so they take effect immediately after startup.
|
||||
"""
|
||||
try:
|
||||
import json
|
||||
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
_RUNTIME_GENERAL_SETTINGS_FLAGS,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
return
|
||||
db_record: Final[_UISettingsRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime is a dict
|
||||
"_UISettingsRow | None",
|
||||
await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}),
|
||||
)
|
||||
if db_record and db_record.ui_settings:
|
||||
raw: Final = db_record.ui_settings
|
||||
ui_settings: Final = json.loads(raw) if isinstance(raw, str) else dict(raw)
|
||||
flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
|
||||
if flags_to_sync:
|
||||
general_settings.update(flags_to_sync)
|
||||
verbose_proxy_logger.info(
|
||||
"Synced UI settings to general_settings on startup: %s",
|
||||
list(flags_to_sync.keys()),
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("UI settings sync on startup skipped or failed: %s", e)
|
||||
"""Apply the persisted UI settings to general_settings before this pod serves traffic."""
|
||||
if prisma_client is None:
|
||||
return
|
||||
applied: Final = await sync_ui_settings_to_general_settings(prisma_client)
|
||||
if applied:
|
||||
verbose_proxy_logger.info("Synced UI settings to general_settings on startup: %s", list(applied))
|
||||
|
||||
@classmethod
|
||||
async def _load_heuristic_v1_tuning_baselines(
|
||||
|
|
@ -12501,7 +12482,6 @@ from litellm.repositories.table_repositories import (
|
|||
InvitationLinkRepository,
|
||||
PromptRepository,
|
||||
SSOConfigRepository,
|
||||
UISettingsRepository,
|
||||
)
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from typing import (
|
|||
from urllib.parse import urlparse
|
||||
|
||||
from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
|
||||
from pydantic import ConfigDict, JsonValue, ValidationError, create_model
|
||||
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model
|
||||
from pydantic.fields import FieldInfo
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
|
|
@ -1478,6 +1478,42 @@ async def get_ui_settings_cached() -> dict[str, JsonValue]:
|
|||
return ui_settings
|
||||
|
||||
|
||||
_UI_SETTINGS_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
def apply_runtime_general_settings_flags(ui_settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
|
||||
"""Copy the UI settings that gate runtime behavior into ``general_settings``. Returns what was applied."""
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
flags: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
|
||||
if flags:
|
||||
general_settings.update(flags)
|
||||
return MappingProxyType(flags)
|
||||
|
||||
|
||||
async def sync_ui_settings_to_general_settings(prisma_client: object) -> Mapping[str, JsonValue]:
|
||||
"""Re-read the persisted UI settings and apply the runtime flags to ``general_settings``.
|
||||
|
||||
Runs on startup and on every periodic config reload: the PATCH handler only updates the pod
|
||||
that served it, so every other pod needs its own read to pick up a change without a restart.
|
||||
Never raises. A read that fails leaves this pod on the flags it already had.
|
||||
"""
|
||||
try:
|
||||
db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
|
||||
where={"id": "ui_settings"}
|
||||
)
|
||||
stored: Final = (db_record.ui_settings if db_record else None) or "{}"
|
||||
parsed: Final = (
|
||||
_UI_SETTINGS_OBJECT.validate_json(stored)
|
||||
if isinstance(stored, str)
|
||||
else _UI_SETTINGS_OBJECT.validate_python(stored)
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning("Could not refresh UI settings from the database: %s", e)
|
||||
return MappingProxyType({})
|
||||
return apply_runtime_general_settings_flags(parsed)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/get/ui_settings",
|
||||
tags=["UI Settings"],
|
||||
|
|
@ -1506,13 +1542,7 @@ async def get_ui_settings():
|
|||
# Sanitize any unexpected keys from persisted config before returning
|
||||
ui_settings: Final = {k: v for k, v in parsed.items() if k in ALLOWED_UI_SETTINGS_FIELDS}
|
||||
|
||||
# Sync runtime flags into general_settings so the proxy picks them up
|
||||
# at runtime (covers server restart scenarios).
|
||||
_flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
|
||||
if _flags_to_sync:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
general_settings.update(_flags_to_sync)
|
||||
apply_runtime_general_settings_flags(ui_settings)
|
||||
|
||||
# Refresh DualCache so other code paths (e.g. /user/filter/ui) see fresh values
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
|
@ -1651,13 +1681,7 @@ async def update_ui_settings(
|
|||
},
|
||||
)
|
||||
|
||||
# Sync runtime flags to general_settings so the proxy picks them up
|
||||
# at runtime (general_settings is checked in pre-call utils).
|
||||
_flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
|
||||
if _flags_to_sync:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
general_settings.update(_flags_to_sync)
|
||||
apply_runtime_general_settings_flags(ui_settings)
|
||||
|
||||
# Invalidate + set DualCache so subsequent reads see the new values immediately
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
|
|
|||
|
|
@ -3722,3 +3722,58 @@ async def test_ProxyConfig__init_guardrails_in_db_skips_only_the_unloadable_row(
|
|||
|
||||
assert sorted(handler.IN_MEMORY_GUARDRAILS) == ["first", "last"]
|
||||
assert handler.reconciled_with == [{"first", "broken", "last"}]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# add_deployment: UI settings convergence
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_deployment_re_reads_ui_settings_so_other_pods_converge(monkeypatch):
|
||||
"""The periodic config reload picks up a UI setting written through another pod.
|
||||
|
||||
Startup used to be the only read, so a proxy admin flipping a runtime flag reached the pod
|
||||
that served the PATCH and nowhere else until every other pod restarted.
|
||||
"""
|
||||
general_settings: Dict[str, Any] = {"allow_agents_for_team_admins": False}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
|
||||
prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_uisettings.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
ui_settings=json.dumps({"allow_agents_for_team_admins": True, "enable_chat_ui": False})
|
||||
)
|
||||
)
|
||||
|
||||
config = ProxyConfig()
|
||||
config._should_load_db_object = MagicMock(return_value=False)
|
||||
config._init_non_llm_objects_in_db = AsyncMock()
|
||||
|
||||
await config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=MagicMock())
|
||||
|
||||
prisma_client.db.litellm_uisettings.find_unique.assert_awaited_once_with(where={"id": "ui_settings"})
|
||||
assert general_settings["allow_agents_for_team_admins"] is True
|
||||
assert "enable_chat_ui" not in general_settings
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_deployment_syncs_ui_settings_even_when_the_model_reconcile_fails(monkeypatch):
|
||||
"""A broken model reconcile must not strand every pod on stale settings."""
|
||||
general_settings: Dict[str, Any] = {"allow_agents_for_team_admins": False}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_uisettings.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(ui_settings={"allow_agents_for_team_admins": True})
|
||||
)
|
||||
|
||||
config = ProxyConfig()
|
||||
config._should_load_db_object = MagicMock(side_effect=RuntimeError("db down"))
|
||||
|
||||
await config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=MagicMock())
|
||||
|
||||
assert general_settings["allow_agents_for_team_admins"] is True
|
||||
|
|
|
|||
|
|
@ -3374,3 +3374,70 @@ class TestTeamAdminEditableTeamFieldsSetting:
|
|||
assert field_schema["type"] == "array"
|
||||
assert field_schema["items"]["type"] == "string"
|
||||
assert isinstance(field_schema["items"]["enum"], list)
|
||||
|
||||
|
||||
class TestSyncUiSettingsToGeneralSettings:
|
||||
"""The DB re-read each pod runs on startup and on every config reload."""
|
||||
|
||||
def _sync(self):
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
sync_ui_settings_to_general_settings,
|
||||
)
|
||||
|
||||
return sync_ui_settings_to_general_settings
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_applies_runtime_flags_and_leaves_other_ui_settings_alone(self, monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
general_settings: dict = {"allow_agents_for_team_admins": False}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
mock_prisma = MagicMock()
|
||||
record = MagicMock()
|
||||
record.ui_settings = json.dumps(
|
||||
{
|
||||
"allow_agents_for_team_admins": True,
|
||||
"team_admin_editable_team_fields": ["tpm_limit"],
|
||||
"enable_chat_ui": False,
|
||||
}
|
||||
)
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=record)
|
||||
|
||||
applied = await self._sync()(mock_prisma)
|
||||
|
||||
assert dict(applied) == {
|
||||
"allow_agents_for_team_admins": True,
|
||||
"team_admin_editable_team_fields": ["tpm_limit"],
|
||||
}
|
||||
assert general_settings["allow_agents_for_team_admins"] is True
|
||||
assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"]
|
||||
assert "enable_chat_ui" not in general_settings
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reads_a_row_the_prisma_client_already_deserialized(self, monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
general_settings: dict = {}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
mock_prisma = MagicMock()
|
||||
record = MagicMock()
|
||||
record.ui_settings = {"team_admin_editable_team_fields": ["rpm_limit"]}
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=record)
|
||||
|
||||
await self._sync()(mock_prisma)
|
||||
|
||||
assert general_settings["team_admin_editable_team_fields"] == ["rpm_limit"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_a_stored_row_general_settings_is_left_untouched(self, monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
general_settings: dict = {"allow_agents_for_team_admins": True}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
applied = await self._sync()(mock_prisma)
|
||||
|
||||
assert dict(applied) == {}
|
||||
assert general_settings == {"allow_agents_for_team_admins": True}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue