mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix(proxy): respect store_model_in_db=False during periodic DB model sync
`ProxyConfig.add_deployment` runs on every loop and reads models from the
DB, then calls `_update_llm_router` which substitutes router state. The
only existing gate was `_should_load_db_object("models")`, which defaults
to "load all" when `supported_db_objects` is unset — meaning a stray row
in `LiteLLM_ProxyModelTable` (e.g. left over from a UI add) silently
replaces the models declared in `config.yaml` on every sync, even when
`store_model_in_db: False`.
The flag's documented purpose is "allow saving / managing models in DB".
When disabled, the DB must not be the source of truth for models —
neither for writes (already enforced in `/model/new`) nor for reads.
Fix: also gate the DB sync in `add_deployment` on `store_model_in_db`,
so `store_model_in_db: False` consistently means "config.yaml is the
sole source of truth for models". The `_should_load_db_object` gate is
preserved for fine-grained control when `store_model_in_db: True`.
Repro:
1. Set `general_settings.store_model_in_db: False` in config.yaml
2. Define N models in `model_list`
3. Have any row in `LiteLLM_ProxyModelTable` (or add one via UI before
toggling the flag)
4. Wait for one sync cycle — config.yaml models disappear from
`/v1/models`, only the DB rows remain.
After this fix, the YAML models are stable across syncs.
This commit is contained in:
AriOliv
parent
934ecdca78
commit
84c3e6e8f5
2 changed files with 226 additions and 3 deletions
|
|
@ -4964,7 +4964,7 @@ class ProxyConfig:
|
|||
- Check if model id's in router already
|
||||
- If not, add to router
|
||||
"""
|
||||
global llm_router, llm_model_list, master_key, general_settings
|
||||
global llm_router, llm_model_list, master_key, general_settings, store_model_in_db
|
||||
|
||||
try:
|
||||
# warm the config cache so the per-param reads below all hit
|
||||
|
|
@ -4980,8 +4980,21 @@ class ProxyConfig:
|
|||
],
|
||||
)
|
||||
|
||||
# Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set)
|
||||
if self._should_load_db_object(object_type="models"):
|
||||
# Only load models from DB if:
|
||||
# 1. store_model_in_db is enabled (the flag's documented purpose is
|
||||
# "allow saving / managing models in DB"); when disabled, the
|
||||
# DB must not be the source of truth for models, otherwise a
|
||||
# stale row in LiteLLM_ProxyModelTable can silently replace
|
||||
# models defined in config.yaml on every periodic sync.
|
||||
# 2. "models" is in supported_db_objects (or supported_db_objects
|
||||
# is not set, preserving the existing fine-grained gate).
|
||||
store_model_in_db_enabled = (
|
||||
general_settings.get("store_model_in_db", False) is True
|
||||
or store_model_in_db
|
||||
)
|
||||
if store_model_in_db_enabled and self._should_load_db_object(
|
||||
object_type="models"
|
||||
):
|
||||
new_models = await self._get_models_from_db(prisma_client=prisma_client)
|
||||
|
||||
# update llm router
|
||||
|
|
|
|||
|
|
@ -0,0 +1,210 @@
|
|||
"""
|
||||
Test that ProxyConfig.add_deployment honours store_model_in_db.
|
||||
|
||||
Regression test for: when store_model_in_db is False the periodic DB sync
|
||||
must not pull models from LiteLLM_ProxyModelTable. Previously the gate was
|
||||
only on `_should_load_db_object("models")` (default-on), so a stray UI-added
|
||||
row could replace config.yaml models on every sync — silently wiping out
|
||||
deployments declared in config.
|
||||
|
||||
Documented behaviour of `store_model_in_db`:
|
||||
"Allow saving / managing models in DB" — when disabled, the DB is not the
|
||||
source of truth for models.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
||||
class TestAddDeploymentStoreModelInDB:
|
||||
"""Test add_deployment respects store_model_in_db before reading DB models."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_db_model_load_when_store_model_in_db_false(self):
|
||||
"""When store_model_in_db=False (config.yaml), no DB sync of models."""
|
||||
proxy_config = ProxyConfig()
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_get_models_from_db",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_models,
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_update_llm_router",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_update_router,
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_init_non_llm_objects_in_db",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.prefetch_config_params",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.get_config_param",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"store_model_in_db": False},
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False),
|
||||
):
|
||||
await proxy_config.add_deployment(
|
||||
prisma_client=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
mock_get_models.assert_not_called()
|
||||
mock_update_router.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_loads_db_models_when_store_model_in_db_true(self):
|
||||
"""When store_model_in_db=True (config.yaml), DB sync runs as before."""
|
||||
proxy_config = ProxyConfig()
|
||||
|
||||
db_models = [MagicMock(model_id="db-id-1", model_name="some-model")]
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_get_models_from_db",
|
||||
new_callable=AsyncMock,
|
||||
return_value=db_models,
|
||||
) as mock_get_models,
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_update_llm_router",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_update_router,
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_init_non_llm_objects_in_db",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.prefetch_config_params",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.get_config_param",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"store_model_in_db": True},
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
):
|
||||
await proxy_config.add_deployment(
|
||||
prisma_client=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
mock_get_models.assert_called_once()
|
||||
mock_update_router.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_loads_db_models_when_module_global_true(self):
|
||||
"""
|
||||
Env var path (STORE_MODEL_IN_DB=True) sets the module-level
|
||||
`store_model_in_db` global. Must still trigger DB sync even when
|
||||
general_settings does not explicitly set the key.
|
||||
"""
|
||||
proxy_config = ProxyConfig()
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_get_models_from_db",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[],
|
||||
) as mock_get_models,
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_update_llm_router",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_update_router,
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_init_non_llm_objects_in_db",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.prefetch_config_params",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.get_config_param",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
):
|
||||
await proxy_config.add_deployment(
|
||||
prisma_client=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
mock_get_models.assert_called_once()
|
||||
mock_update_router.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_supported_db_objects_gate_still_applies(self):
|
||||
"""
|
||||
Even with store_model_in_db=True, an explicit supported_db_objects list
|
||||
that omits "models" must still skip the DB sync. This preserves the
|
||||
existing fine-grained gate.
|
||||
"""
|
||||
proxy_config = ProxyConfig()
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_get_models_from_db",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_models,
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_update_llm_router",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_update_router,
|
||||
patch.object(
|
||||
proxy_config,
|
||||
"_init_non_llm_objects_in_db",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.prefetch_config_params",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.get_config_param",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{
|
||||
"store_model_in_db": True,
|
||||
"supported_db_objects": ["mcp", "guardrails"],
|
||||
},
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
):
|
||||
await proxy_config.add_deployment(
|
||||
prisma_client=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
mock_get_models.assert_not_called()
|
||||
mock_update_router.assert_not_called()
|
||||
Loading…
Add table
Reference in a new issue