diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c0246f234a8b..d41f37a82e04 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -45,6 +45,7 @@ def _build_db_connection_url_params( connect_timeout: Optional[Union[int, float]] = None, socket_timeout: Optional[Union[int, float]] = None, extra_params: Optional[dict] = None, + existing_url: Optional[str] = None, ) -> dict: """Build the Prisma DATABASE_URL query params controlling connection pool behavior. @@ -53,10 +54,39 @@ def _build_db_connection_url_params( omitted when None so Prisma's defaults apply. `extra_params` is an untyped passthrough — keys it provides win over the named arguments above, so it can be used to override any default we set here. + + The four `keepalives*` defaults are libpq query-string params that get + forwarded to the underlying Postgres driver. Without them, dead TCP + sessions from killed Prisma engine subprocesses sit on the DB side for + the OS keepalive default (~7200s) and accumulate toward max_connections + under worker churn. The defaults here cause the kernel to detect dead + peers in ~90s instead. See BerriAI/litellm#26619. + + `existing_url` is the DATABASE_URL these params will be appended to. Any + keepalive key already present in that URL's query string is left as-is, + so operators who have set explicit values (e.g. `keepalives=0`) keep + them. `extra_params` still wins last for both keepalives and + everything else. """ - params: dict = { - "connection_limit": connection_limit, + keepalive_defaults: dict = { + "keepalives": "1", + "keepalives_idle": "60", + "keepalives_interval": "10", + "keepalives_count": "3", } + existing_keys: set = set() + if existing_url: + try: + existing_keys = set( + urlparse.parse_qs(urlparse.urlparse(existing_url).query).keys() + ) + except Exception: + existing_keys = set() + + params: dict = {"connection_limit": connection_limit} + for k, v in keepalive_defaults.items(): + if k not in existing_keys: + params[k] = v if pool_timeout is not None: params["pool_timeout"] = pool_timeout if connect_timeout is not None: @@ -1087,24 +1117,33 @@ def run_server( # noqa: PLR0915 try: from litellm.secret_managers.main import get_secret - connection_url_params = _build_db_connection_url_params( - connection_limit=db_connection_pool_limit, - pool_timeout=db_connection_timeout, - connect_timeout=db_connect_timeout, - socket_timeout=db_socket_timeout, - extra_params=db_extra_connection_params, - ) if os.getenv("DATABASE_URL", None) is not None: database_url = get_secret("DATABASE_URL", default_value=None) + url_str = str(database_url) if database_url else None modified_url = append_query_params( - str(database_url) if database_url else None, - connection_url_params, + url_str, + _build_db_connection_url_params( + connection_limit=db_connection_pool_limit, + pool_timeout=db_connection_timeout, + connect_timeout=db_connect_timeout, + socket_timeout=db_socket_timeout, + extra_params=db_extra_connection_params, + existing_url=url_str, + ), ) os.environ["DATABASE_URL"] = modified_url if os.getenv("DIRECT_URL", None) is not None: - database_url = os.getenv("DIRECT_URL") + direct_url = os.getenv("DIRECT_URL") modified_url = append_query_params( - database_url, connection_url_params + direct_url, + _build_db_connection_url_params( + connection_limit=db_connection_pool_limit, + pool_timeout=db_connection_timeout, + connect_timeout=db_connect_timeout, + socket_timeout=db_socket_timeout, + extra_params=db_extra_connection_params, + existing_url=direct_url, + ), ) os.environ["DIRECT_URL"] = modified_url subprocess.run(["prisma"], capture_output=True) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 759534a32a1f..a06c847cbdf4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -966,7 +966,17 @@ async def _run_pw_migration(): except Exception as e: verbose_proxy_logger.error(f"Error stopping DB health watchdog task: {e}") - await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] + # Shutdown disconnect failures must not abort the uvicorn lifespan — if the + # lifespan exits with an exception, uvicorn skips the Python-level + # `atexit`/SIGTERM cleanup and the Prisma query-engine subprocess is left + # as an orphan whose Postgres connections sit on the DB side for the OS + # keepalive default (~2h). See BerriAI/litellm#26619. + try: + await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] + except Exception: + verbose_proxy_logger.exception( + "proxy_shutdown_event failed (continuing shutdown to avoid orphan subprocesses)" + ) def _generate_stable_operation_id(route: Any) -> str: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 032ab6c63b21..86994cc5497b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4147,14 +4147,6 @@ async def connect(self): ) raise e - # Define a retrying strategy with exponential backoff - @backoff.on_exception( - backoff.expo, - Exception, # base exception to catch for the backoff - max_tries=3, # maximum number of retries - max_time=10, # maximum total time to retry for - on_backoff=on_backoff, # specifying the function to call on backoff - ) async def disconnect(self): start_time = time.time() try: @@ -4175,7 +4167,12 @@ async def disconnect(self): traceback_str=error_traceback, ) ) - raise e + # Do NOT re-raise: propagating from disconnect() aborts the + # uvicorn lifespan shutdown and leaves the Prisma query-engine + # subprocess as an orphan with leaked Postgres connections. + # disconnect() is only called from shutdown paths; no caller + # benefits from the re-raise, and the lifespan is actively + # harmed by it. See BerriAI/litellm#26619. def _get_engine_pid(self) -> int: try: diff --git a/tests/proxy_unit_tests/test_aproxy_startup.py b/tests/proxy_unit_tests/test_aproxy_startup.py index 4dbf5b462a91..15c59bc6ab8a 100644 --- a/tests/proxy_unit_tests/test_aproxy_startup.py +++ b/tests/proxy_unit_tests/test_aproxy_startup.py @@ -13,6 +13,7 @@ 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import pytest, logging, asyncio +from unittest.mock import AsyncMock import litellm from litellm.proxy.proxy_server import ( router, @@ -94,3 +95,50 @@ async def test_proxy_gunicorn_startup_config_dict(): # test_proxy_gunicorn_startup() + + +@pytest.mark.asyncio +async def test_proxy_lifespan_swallows_shutdown_event_exception(monkeypatch, caplog): + """ + Regression test for BerriAI/litellm#26619. + + `proxy_shutdown_event()` is invoked inside the uvicorn lifespan + `proxy_startup_event` context manager. If it raises and the call site + has no try/except, the lifespan exits with an exception, uvicorn skips + the Python-level `atexit`/SIGTERM cleanup, and the Prisma query-engine + subprocess is left as an orphan whose Postgres connections sit on the + DB side for ~2h. This test pins the wrapped-in-try/except behavior. + """ + from litellm._logging import verbose_proxy_logger + + setattr(litellm.proxy.proxy_server, "prisma_client", None) + database_url = os.environ.pop("DATABASE_URL", None) + + filepath = os.path.dirname(os.path.abspath(__file__)) + config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" + os.environ["WORKER_CONFIG"] = config_fp + + raising_shutdown = AsyncMock( + side_effect=RuntimeError("simulated disconnect failure") + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_shutdown_event", raising_shutdown + ) + + try: + verbose_proxy_logger.setLevel(logging.DEBUG) + caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name) + + # Driving the full lifespan: if the patch is missing, RuntimeError + # propagates out of the __aexit__ here and the test fails. + async with proxy_startup_event(app=None) as _: + pass + + raising_shutdown.assert_awaited_once() + # The exception should have been logged, not swallowed silently. + assert any( + "proxy_shutdown_event failed" in rec.getMessage() for rec in caplog.records + ), "expected proxy_shutdown_event failure to be logged" + finally: + if database_url is not None: + os.environ["DATABASE_URL"] = database_url diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 580ed95062ba..695ff756da52 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -581,7 +581,55 @@ def test_build_db_connection_url_params_defaults(self): from litellm.proxy.proxy_cli import _build_db_connection_url_params params = _build_db_connection_url_params(connection_limit=10, pool_timeout=60) - assert params == {"connection_limit": 10, "pool_timeout": 60} + assert params == { + "connection_limit": 10, + "pool_timeout": 60, + "keepalives": "1", + "keepalives_idle": "60", + "keepalives_interval": "10", + "keepalives_count": "3", + } + + def test_build_db_connection_url_params_includes_keepalives(self): + """Default libpq keepalive params are emitted so the kernel reaps dead + Postgres sessions from killed Prisma engine subprocesses in ~90s + instead of the OS default (~7200s). See BerriAI/litellm#26619.""" + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params(connection_limit=10, pool_timeout=60) + assert params["keepalives"] == "1" + assert params["keepalives_idle"] == "60" + assert params["keepalives_interval"] == "10" + assert params["keepalives_count"] == "3" + + def test_build_db_connection_url_params_keepalives_overridable_via_extras(self): + """Operators can opt out of the keepalive defaults by passing + `extra_params` because `extra_params.update` runs last.""" + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + extra_params={"keepalives": "0"}, + ) + assert params["keepalives"] == "0" + + def test_build_db_connection_url_params_respects_existing_url_keepalives(self): + """If the DATABASE_URL already has a keepalive value, the default + for that key is not emitted — operators who set explicit + keepalive values in their URL keep them.""" + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + existing_url=("postgresql://u:p@h/db?keepalives=0&keepalives_idle=120"), + ) + assert "keepalives" not in params + assert "keepalives_idle" not in params + # Keys NOT in the URL still get the default. + assert params["keepalives_interval"] == "10" + assert params["keepalives_count"] == "3" def test_build_db_connection_url_params_omits_none_timeouts(self): from litellm.proxy.proxy_cli import _build_db_connection_url_params diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ae0996d16d5b..dbbea34bbce1 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -906,6 +906,90 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): assert master_key == test_resolved_key +@pytest.mark.asyncio +async def test_proxy_lifespan_runs_shutdown_event_on_clean_exit(monkeypatch): + """ + Happy-path coverage for the try/except around `await proxy_shutdown_event()` + inside the uvicorn lifespan (BerriAI/litellm#26619). With shutdown + mocked to a clean no-op, the lifespan exits without entering the + except branch. + """ + import litellm + from litellm.proxy.proxy_server import proxy_startup_event + + setattr(litellm.proxy.proxy_server, "prisma_client", None) + database_url = os.environ.pop("DATABASE_URL", None) + + filepath = os.path.dirname(os.path.abspath(__file__)) + # Use the existing repo test config that other proxy_startup_event tests use. + config_fp = os.path.normpath( + os.path.join( + filepath, + "..", + "..", + "proxy_unit_tests", + "test_configs", + "test_config_no_auth.yaml", + ) + ) + os.environ["WORKER_CONFIG"] = config_fp + + clean_shutdown = AsyncMock(return_value=None) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_shutdown_event", clean_shutdown + ) + + try: + async with proxy_startup_event(app=None) as _: + pass + clean_shutdown.assert_awaited_once() + finally: + if database_url is not None: + os.environ["DATABASE_URL"] = database_url + + +@pytest.mark.asyncio +async def test_prisma_disconnect_does_not_reraise_on_failure(): + """ + Regression test for BerriAI/litellm#26619. + + `PrismaClient.disconnect()` is only ever called from shutdown paths. + If it re-raises on failure, the uvicorn lifespan exits with an + exception, uvicorn skips the Python-level `atexit`/SIGTERM cleanup, + and the Prisma query-engine subprocess is left as an orphan whose + Postgres connections sit on the DB side for ~2h. This test pins the + "swallow + log + schedule failure_handler" behavior. + """ + from litellm.proxy.utils import PrismaClient + + # Bypass __init__ (which imports prisma binaries) — we're only testing + # the disconnect() failure path, which needs `self.db` and + # `self.proxy_logging_obj` and nothing else. + client = PrismaClient.__new__(PrismaClient) + + boom = AsyncMock(side_effect=RuntimeError("simulated engine disconnect failure")) + fake_db = MagicMock() + fake_db.disconnect = boom + client.db = fake_db + + failure_handler = AsyncMock() + fake_proxy_logging = MagicMock() + fake_proxy_logging.failure_handler = failure_handler + client.proxy_logging_obj = fake_proxy_logging + + # If the fix regresses, this call propagates RuntimeError and the test fails. + await client.disconnect() + + boom.assert_awaited() + # The failure_handler is scheduled via asyncio.create_task — yield once + # to let the scheduled coroutine run before asserting. + await asyncio.sleep(0) + failure_handler.assert_awaited_once() + call_kwargs = failure_handler.await_args.kwargs + assert call_kwargs["call_type"] == "disconnect" + assert isinstance(call_kwargs["original_exception"], RuntimeError) + + def test_team_info_masking(): """ Test that sensitive team information is properly masked @@ -5859,15 +5943,16 @@ async def fresh_lock(_counter_key): if call.kwargs.get("nx") is True ] assert len(nx_writes) == 2 - assert sorted(set_results) == [False, True], ( - f"expected exactly one SET NX winner and one loser, got {set_results}" - ) + assert sorted(set_results) == [ + False, + True, + ], f"expected exactly one SET NX winner and one loser, got {set_results}" # Loser path executed: after the winner's SET NX returned True, the # losing coalesced() call falls back to async_get_cache to read the # winner's value rather than re-seeding. - assert get_after_set_count >= 1, ( - "loser branch (else: read back winner's value) was never exercised" - ) + assert ( + get_after_set_count >= 1 + ), "loser branch (else: read back winner's value) was never exercised" @pytest.mark.asyncio