diff --git a/litellm/constants.py b/litellm/constants.py index a292b654778..a34c031443a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1823,6 +1823,13 @@ PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDUL PROXY_BATCH_WRITE_AT: Final = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) +# Per-step timeout (seconds) for DB calls made by the periodic config-sync job +# (ProxyConfig.add_deployment). A DB query stuck on a dead pooled connection +# would otherwise hang forever; with max_instances=1 that freezes the job and +# the in-memory model list never sees new deployments. +# <=0 disables the timeout (previous behavior). +PROXY_CONFIG_SYNC_DB_TIMEOUT_SECONDS: Final = get_env_int("PROXY_CONFIG_SYNC_DB_TIMEOUT_SECONDS", 60) + # APScheduler Configuration - MEMORY LEAK FIX # These settings prevent memory leaks in APScheduler's normalize() and _apply_jitter() functions APSCHEDULER_COALESCE: Final = os.getenv("APSCHEDULER_COALESCE", "True").lower() in [ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f842f2e1e4a..ec5db2b48f1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -39,6 +39,7 @@ from typing import ( Protocol, TypeAlias, TypedDict, + TypeVar, Union, cast, get_args, @@ -293,6 +294,7 @@ from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + PROXY_CONFIG_SYNC_DB_TIMEOUT_SECONDS, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG, @@ -2542,6 +2544,34 @@ heuristic_v1_tuning_baselines: Mapping[str, str] | None = None # Module-level rather than per-ProxyConfig because llm_router is a module global and a # second ProxyConfig instance must not get its own independent lock over it. MODEL_RECONCILE_LOCK: Final = asyncio.Lock() + +_T = TypeVar("_T") + + +async def _with_sync_db_timeout(awaitable: Awaitable[_T], step: str) -> _T: + """Guard a config-sync DB call with PROXY_CONFIG_SYNC_DB_TIMEOUT_SECONDS. + + A query that lands on a dead pooled connection never errors and never + returns; with APScheduler max_instances=1 that single hang freezes the + periodic config sync (and the in-memory model list) until process restart, + while holding MODEL_RECONCILE_LOCK blocks every other reconcile caller. + Timeout <= 0 disables the guard (previous behavior). + """ + if PROXY_CONFIG_SYNC_DB_TIMEOUT_SECONDS <= 0: + return await awaitable + try: + return await asyncio.wait_for(awaitable, timeout=PROXY_CONFIG_SYNC_DB_TIMEOUT_SECONDS) + except asyncio.TimeoutError: + verbose_proxy_logger.error( + "Config sync DB step '%s' timed out after %ss; " + "skipping this step. If this repeats, the DB connection pool " + "may have dead connections (check NAT keepalive / socket timeout).", + step, + PROXY_CONFIG_SYNC_DB_TIMEOUT_SECONDS, + ) + raise + + general_settings: dict = {} _GENERAL_SETTINGS_VIEW: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) @@ -7846,29 +7876,48 @@ class ProxyConfig: still_desired_ids: frozenset[str] | None = None + # Each DB step is wrapped with its own timeout so a query stuck on a + # dead pooled connection cancels and frees this job (and the reconcile + # lock) instead of hanging forever and freezing the in-memory model + # list until process restart. try: # warm the config cache so the per-param reads below all hit - await prefetch_config_params( - prisma_client, - [ - "general_settings", - "router_settings", - "litellm_settings", - "environment_variables", - "anthropic_beta_headers_reload_config", - ], + await _with_sync_db_timeout( + prefetch_config_params( + prisma_client, + [ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + "anthropic_beta_headers_reload_config", + ], + ), + step="prefetch_config_params", ) load_models: Final = self._should_load_db_object(object_type="models") - new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) if load_models else None - await self.get_credentials(prisma_client=prisma_client) + new_models: Final = ( + await _with_sync_db_timeout( + self._get_models_from_db(prisma_client=prisma_client), + step="_get_models_from_db", + ) + if load_models + else None + ) + await _with_sync_db_timeout( + self.get_credentials(prisma_client=prisma_client), + step="get_credentials", + ) if load_models: - still_desired_ids = await self._update_llm_router( - new_models=new_models, proxy_logging_obj=proxy_logging_obj + still_desired_ids = await _with_sync_db_timeout( + self._update_llm_router(new_models=new_models, proxy_logging_obj=proxy_logging_obj), + step="_update_llm_router", ) - db_general_settings: Final[_ConfigParamRow | None] = await get_config_param( - prisma_client, "general_settings" + db_general_settings: Final[_ConfigParamRow | None] = await _with_sync_db_timeout( + get_config_param(prisma_client, "general_settings"), + step="get_config_param(general_settings)", ) # update general settings @@ -7878,7 +7927,10 @@ class ProxyConfig: ) # initialize vector stores, guardrails, etc. table in db - await self._init_non_llm_objects_in_db(prisma_client=prisma_client) + await _with_sync_db_timeout( + self._init_non_llm_objects_in_db(prisma_client=prisma_client), + step="_init_non_llm_objects_in_db", + ) except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - %s", e) diff --git a/tests/test_litellm/proxy/proxy_server/test_config_sync_db_timeout.py b/tests/test_litellm/proxy/proxy_server/test_config_sync_db_timeout.py new file mode 100644 index 00000000000..51d3bcfdf8a --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_config_sync_db_timeout.py @@ -0,0 +1,145 @@ +"""Behavior tests for the config-sync DB timeout guard. + +``_with_sync_db_timeout`` bounds each DB call made by the periodic +``ProxyConfig.add_deployment`` job with ``PROXY_CONFIG_SYNC_DB_TIMEOUT_SECONDS``. + +Why this exists: a DB query that lands on a dead pooled connection neither +errors nor returns. With APScheduler ``max_instances=1`` the periodic config +sync hangs forever, every later run is skipped, and the in-memory model list +freezes -- requests for models added after the hang return 400 "Invalid model +name" until the process restarts. The timeout turns the hang into a normal +exception so the existing skip-and-retry-next-tick handling takes over. +""" + +from __future__ import annotations + +import asyncio +import time +from unittest.mock import AsyncMock + +import pytest + +import litellm.proxy.proxy_server as proxy_server_module +from litellm.proxy.proxy_server import ProxyConfig + + +def _patch_sync_timeout(monkeypatch: pytest.MonkeyPatch, seconds: float) -> None: + """Patch the timeout constant on the module that reads it. + + ``proxy_server`` imports the constant by value at module import time, so + patch the module attribute the function references. + """ + monkeypatch.setattr(proxy_server_module, "PROXY_CONFIG_SYNC_DB_TIMEOUT_SECONDS", seconds) + + +def _hang_forever() -> asyncio.Future: + """An awaitable that never completes: simulates a query stuck on a dead + pooled connection (no response, no error, no cancellation on its own).""" + return asyncio.get_running_loop().create_future() + + +class _FakePrismaClient: + """Minimal stand-in; every DB step is mocked per-test.""" + + +class TestWithSyncDbTimeout: + async def test_hanging_db_call_times_out(self, monkeypatch: pytest.MonkeyPatch): + _patch_sync_timeout(monkeypatch, 0.1) + with pytest.raises(asyncio.TimeoutError): + await proxy_server_module._with_sync_db_timeout(_hang_forever(), step="unit-test") + + async def test_fast_db_call_unaffected(self, monkeypatch: pytest.MonkeyPatch): + _patch_sync_timeout(monkeypatch, 5) + + async def quick() -> str: + await asyncio.sleep(0) + return "ok" + + assert await proxy_server_module._with_sync_db_timeout(quick(), step="unit-test") == "ok" + + async def test_timeout_disabled_passthrough(self, monkeypatch: pytest.MonkeyPatch): + _patch_sync_timeout(monkeypatch, 0) + + async def slow_but_finite() -> str: + await asyncio.sleep(0.2) + return "done" + + # With the timeout disabled the call runs to completion. + start = time.monotonic() + assert await proxy_server_module._with_sync_db_timeout(slow_but_finite(), step="unit-test") == "done" + assert time.monotonic() - start >= 0.2 + + +class TestAddDeploymentDoesNotHang: + async def test_add_deployment_returns_when_model_fetch_hangs(self, monkeypatch: pytest.MonkeyPatch): + """The incident scenario: the model-table query hangs on a dead + connection. ``add_deployment`` must return (not hang) so the next + scheduler tick can proceed.""" + _patch_sync_timeout(monkeypatch, 0.2) + proxy_config = ProxyConfig() + + async def _hang(self, prisma_client): + return await _hang_forever() + + monkeypatch.setattr(ProxyConfig, "_get_models_from_db", _hang) + # Non-DB steps still run normally. + monkeypatch.setattr( + proxy_server_module, + "prefetch_config_params", + AsyncMock(return_value=None), + ) + monkeypatch.setattr(proxy_server_module, "get_config_param", AsyncMock(return_value=None)) + monkeypatch.setattr( + ProxyConfig, + "_init_non_llm_objects_in_db", + AsyncMock(return_value=None), + ) + monkeypatch.setattr(ProxyConfig, "_should_load_db_object", lambda self, object_type: True) + + start = time.monotonic() + outcome = await asyncio.wait_for( + proxy_config.add_deployment( + prisma_client=_FakePrismaClient(), # type: ignore[arg-type] + proxy_logging_obj=None, + ), + timeout=10, + ) + elapsed = time.monotonic() - start + + assert outcome is not None + assert outcome.still_desired is None # no reconcile ran + assert elapsed < 5 # returned promptly instead of hanging forever + + async def test_reconcile_lock_released_after_timeout(self, monkeypatch: pytest.MonkeyPatch): + """After a timed-out run, MODEL_RECONCILE_LOCK must be free so other + reconcile callers (model writes, clear_cache) are not blocked.""" + _patch_sync_timeout(monkeypatch, 0.2) + proxy_config = ProxyConfig() + + async def _hang(self, prisma_client): + return await _hang_forever() + + monkeypatch.setattr(ProxyConfig, "_get_models_from_db", _hang) + monkeypatch.setattr( + proxy_server_module, + "prefetch_config_params", + AsyncMock(return_value=None), + ) + monkeypatch.setattr(ProxyConfig, "_should_load_db_object", lambda self, object_type: True) + + await asyncio.wait_for( + proxy_config.add_deployment( + prisma_client=_FakePrismaClient(), # type: ignore[arg-type] + proxy_logging_obj=None, + ), + timeout=10, + ) + + # The lock was held during the run; it must be released now. + assert not proxy_server_module.MODEL_RECONCILE_LOCK.locked() + # And a second run can acquire it immediately. + acquired = await asyncio.wait_for(proxy_server_module.MODEL_RECONCILE_LOCK.acquire(), timeout=1) + try: + assert acquired + finally: + proxy_server_module.MODEL_RECONCILE_LOCK.release()