diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 7f68b72a522..4eae8a95064 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -569,6 +569,7 @@ class ProxyInitializationHelpers: ssl_keyfile_path: str, max_requests_before_restart: int | None = None, max_requests_before_restart_jitter: int | None = None, + keepalive_timeout: int | None = None, ): """ Run litellm with `gunicorn` @@ -659,6 +660,9 @@ class ProxyInitializationHelpers: else: gunicorn_options["max_requests_jitter"] = max_requests_before_restart_jitter + if keepalive_timeout is not None: + gunicorn_options["keepalive"] = keepalive_timeout + # Clean up prometheus .db files when a worker exits (prevents ghost gauge values) if os.environ.get("PROMETHEUS_MULTIPROC_DIR"): from litellm.proxy.prometheus_cleanup import mark_worker_exit @@ -982,7 +986,8 @@ class ProxyInitializationHelpers: "--keepalive_timeout", default=None, type=int, - help="Set the uvicorn keepalive timeout in seconds (uvicorn timeout_keep_alive parameter)", + help="Idle keep-alive timeout in seconds (uvicorn timeout_keep_alive; gunicorn keepalive with --run_gunicorn, " + "default 2s). Behind a reverse proxy, set it above the proxy's idle-connection timeout.", envvar="KEEPALIVE_TIMEOUT", ) @click.option( @@ -1655,6 +1660,7 @@ def run_server( ssl_keyfile_path=ssl_keyfile_path, max_requests_before_restart=max_requests_before_restart, max_requests_before_restart_jitter=max_requests_before_restart_jitter, + keepalive_timeout=keepalive_timeout, ) elif run_hypercorn is True: ProxyInitializationHelpers.init_hypercorn_server( diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 89a35923be0..119688e9656 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -2118,6 +2118,66 @@ class TestProxyInitializationHelpers: assert worker.config.http == UvicornWorker.CONFIG_KWARGS["http"] assert captured["options"]["graceful_timeout"] == 8 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") + @pytest.mark.parametrize( + ("args", "env", "expected_keep_alive"), + [ + ([], {}, 2), + ([], {"KEEPALIVE_TIMEOUT": "120"}, 120), + (["--keepalive_timeout", "60"], {"KEEPALIVE_TIMEOUT": "120"}, 60), + ], + ) + def test_run_gunicorn_workers_use_keepalive_timeout(self, monkeypatch, args, env, expected_keep_alive): + """--keepalive_timeout / KEEPALIVE_TIMEOUT reaches the gunicorn uvicorn worker, unset keeps gunicorn's 2s""" + pytest.importorskip("gunicorn") + import gunicorn.glogging + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + captured: dict = {} + + def capture_run(self): + captured["options"] = dict(self.options) + captured["cfg"] = self.cfg + + clean_env = { + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "KEEPALIVE_TIMEOUT") + } + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + }, + ), + patch("gunicorn.app.base.BaseApplication.run", capture_run), + patch( # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership + "litellm.proxy.proxy_cli.start_query_engine_reaper" + ), + patch("uvicorn.run") as mock_uvicorn_run, + ): + result = CliRunner().invoke(run_server, ["--local", "--run_gunicorn", *args], env=env) + + assert result.exit_code == 0, result.output + mock_uvicorn_run.assert_not_called() + for name in ("uvicorn.error", "uvicorn.access"): + uvicorn_logger = logging.getLogger(name) + for attr in ("handlers", "level", "propagate"): + monkeypatch.setattr(uvicorn_logger, attr, getattr(uvicorn_logger, attr)) + cfg = captured["cfg"] + worker = captured["options"]["worker_class"]( + 0, os.getpid(), [], MagicMock(), 30, cfg, gunicorn.glogging.Logger(cfg) + ) + + assert worker.config.timeout_keep_alive == expected_keep_alive + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") def test_gunicorn_master_is_reserved_for_forking_before_it_runs(self, fork_reservation): """preload forks workers from the master, so native routes are forbidden there first"""