mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
Merge eb57fa1c62 into b6f91a221b
This commit is contained in:
Prajjwal Chittori
GitHub
commit
80814ed2ad
2 changed files with 184 additions and 7 deletions
|
|
@ -11316,14 +11316,26 @@ class ProxyStartupEvent:
|
|||
# but YAML config has False.
|
||||
if store_model_in_db is not True and prisma_client is not None:
|
||||
try:
|
||||
_db_gs_record: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "general_settings"}
|
||||
_db_gs_record: Final[_ConfigParamRow | None] = await call_with_db_reconnect_retry(
|
||||
prisma_client,
|
||||
lambda: _config_param_table(prisma_client).find_first(where={"param_name": "general_settings"}),
|
||||
reason="startup_store_model_in_db_general_settings_lookup_failure",
|
||||
)
|
||||
if _db_gs_record is not None and isinstance(_db_gs_record.param_value, dict):
|
||||
_db_val: Final = _db_gs_record.param_value.get("store_model_in_db")
|
||||
if _db_val is True or (isinstance(_db_val, str) and _db_val.lower() == "true"):
|
||||
store_model_in_db = True
|
||||
verbose_proxy_logger.info("store_model_in_db=True loaded from DB, overriding config/env")
|
||||
if _db_gs_record is not None:
|
||||
_db_param_value = _db_gs_record.param_value
|
||||
# Prisma can return the JSON column as a raw string; normalize
|
||||
# it the way ConfigRepository.get_param does before reading, so
|
||||
# the override isn't silently skipped for string-valued rows.
|
||||
if isinstance(_db_param_value, str):
|
||||
try:
|
||||
_db_param_value = json.loads(_db_param_value)
|
||||
except json.JSONDecodeError:
|
||||
_db_param_value = {}
|
||||
if isinstance(_db_param_value, dict):
|
||||
_db_val: Final = _db_param_value.get("store_model_in_db")
|
||||
if _db_val is True or (isinstance(_db_val, str) and _db_val.lower() == "true"):
|
||||
store_model_in_db = True
|
||||
verbose_proxy_logger.info("store_model_in_db=True loaded from DB, overriding config/env")
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Failed to check DB for store_model_in_db: %s", str(e))
|
||||
|
||||
|
|
|
|||
|
|
@ -1160,6 +1160,171 @@ async def test_initialize_scheduled_jobs_hydrates_mcp_when_store_model_in_db_fal
|
|||
mock_proxy_config.init_mcp_servers_from_db.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_scheduled_jobs_recovers_from_transient_read_error(monkeypatch):
|
||||
"""
|
||||
Regression (#32412): the startup `store_model_in_db` lookup reads
|
||||
general_settings straight off the Prisma HTTP engine. A transient transport
|
||||
blip (httpx.ReadError) on that first read must be healed by one
|
||||
reconnect-and-retry, like every other DB read path, instead of leaving the
|
||||
connection poisoned and the DB-persisted store_model_in_db=True ignored.
|
||||
"""
|
||||
import types
|
||||
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
record = types.SimpleNamespace(param_value={"store_model_in_db": True})
|
||||
find_first = AsyncMock(side_effect=[httpx.ReadError("Server disconnected"), record])
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_config.find_first = find_first
|
||||
mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.slack_alerting_instance = MagicMock()
|
||||
mock_proxy_logging.db_spend_update_writer = MagicMock()
|
||||
mock_proxy_config = _mock_scheduled_proxy_config()
|
||||
|
||||
scheduler = AsyncIOScheduler()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False),
|
||||
patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
|
||||
):
|
||||
await ProxyStartupEvent.initialize_scheduled_background_jobs(
|
||||
general_settings={},
|
||||
prisma_client=mock_prisma_client,
|
||||
proxy_budget_rescheduler_min_time=1,
|
||||
proxy_budget_rescheduler_max_time=2,
|
||||
proxy_batch_write_at=5,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
|
||||
assert scheduler.get_job("add_deployment_job") is not None
|
||||
mock_prisma_client.attempt_db_reconnect.assert_awaited_once()
|
||||
assert find_first.await_count == 2
|
||||
mock_proxy_config.add_deployment.assert_awaited()
|
||||
mock_proxy_config.init_mcp_servers_from_db.assert_not_called()
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_scheduled_jobs_normalizes_string_general_settings(monkeypatch):
|
||||
"""
|
||||
Regression (#32412): Prisma can return the general_settings JSON column as a
|
||||
raw string rather than a parsed dict. The startup store_model_in_db lookup
|
||||
must json.loads that string before reading the flag, otherwise a
|
||||
DB-persisted store_model_in_db=True is silently skipped for string-valued
|
||||
rows and the proxy comes up as if model storage were off.
|
||||
"""
|
||||
import types
|
||||
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
record = types.SimpleNamespace(param_value='{"store_model_in_db": true}')
|
||||
find_first = AsyncMock(return_value=record)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_config.find_first = find_first
|
||||
mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.slack_alerting_instance = MagicMock()
|
||||
mock_proxy_logging.db_spend_update_writer = MagicMock()
|
||||
mock_proxy_config = _mock_scheduled_proxy_config()
|
||||
|
||||
scheduler = AsyncIOScheduler()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False),
|
||||
patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
|
||||
):
|
||||
await ProxyStartupEvent.initialize_scheduled_background_jobs(
|
||||
general_settings={},
|
||||
prisma_client=mock_prisma_client,
|
||||
proxy_budget_rescheduler_min_time=1,
|
||||
proxy_budget_rescheduler_max_time=2,
|
||||
proxy_batch_write_at=5,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
|
||||
assert scheduler.get_job("add_deployment_job") is not None
|
||||
assert find_first.await_count == 1
|
||||
mock_prisma_client.attempt_db_reconnect.assert_not_awaited()
|
||||
mock_proxy_config.add_deployment.assert_awaited()
|
||||
mock_proxy_config.init_mcp_servers_from_db.assert_not_called()
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_scheduled_jobs_ignores_malformed_string_general_settings(monkeypatch):
|
||||
"""
|
||||
A string-valued general_settings row that is not valid JSON must be treated
|
||||
as empty, leaving store_model_in_db off, instead of failing the lookup.
|
||||
"""
|
||||
import types
|
||||
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
record = types.SimpleNamespace(param_value="{not json")
|
||||
find_first = AsyncMock(return_value=record)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_config.find_first = find_first
|
||||
mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.slack_alerting_instance = MagicMock()
|
||||
mock_proxy_logging.db_spend_update_writer = MagicMock()
|
||||
mock_proxy_config = _mock_scheduled_proxy_config()
|
||||
|
||||
scheduler = AsyncIOScheduler()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False),
|
||||
patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
|
||||
patch("litellm.proxy.proxy_server.verbose_proxy_logger") as mock_logger,
|
||||
):
|
||||
await ProxyStartupEvent.initialize_scheduled_background_jobs(
|
||||
general_settings={},
|
||||
prisma_client=mock_prisma_client,
|
||||
proxy_budget_rescheduler_min_time=1,
|
||||
proxy_budget_rescheduler_max_time=2,
|
||||
proxy_batch_write_at=5,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
|
||||
assert scheduler.get_job("add_deployment_job") is None
|
||||
assert find_first.await_count == 1
|
||||
assert not any(
|
||||
"Failed to check DB for store_model_in_db" in str(call.args[0]) for call in mock_logger.debug.call_args_list
|
||||
)
|
||||
mock_proxy_config.add_deployment.assert_not_called()
|
||||
mock_proxy_config.init_mcp_servers_from_db.assert_awaited_once()
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_mcp_servers_from_db_respects_supported_db_objects(monkeypatch):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue