diff --git a/docs/my-website/docs/proxy/deploy.md b/docs/my-website/docs/proxy/deploy.md index 6a11d069fb..dc2da22684 100644 --- a/docs/my-website/docs/proxy/deploy.md +++ b/docs/my-website/docs/proxy/deploy.md @@ -715,6 +715,25 @@ docker run ghcr.io/berriai/litellm:main-stable ``` +### Restart Workers After N Requests + +Use this to mitigate memory growth by recycling workers after a fixed number of requests. When set, each worker restarts after completing the specified number of requests. Defaults to disabled when unset. + +Usage Examples: + +```shell showLineNumbers title="docker run (CLI flag)" +docker run ghcr.io/berriai/litellm:main-stable \ + --max_requests_before_restart 10000 +``` + +Or set via environment variable: + +```shell showLineNumbers title="Environment Variable" +export MAX_REQUESTS_BEFORE_RESTART=10000 +docker run ghcr.io/berriai/litellm:main-stable +``` + + ### 5. config.yaml file on s3, GCS Bucket Object/url Use this if you cannot mount a config file on your deployment service (example - AWS Fargate, Railway etc) diff --git a/docs/my-website/docs/proxy/prod.md b/docs/my-website/docs/proxy/prod.md index a45474f39e..2858132c8e 100644 --- a/docs/my-website/docs/proxy/prod.md +++ b/docs/my-website/docs/proxy/prod.md @@ -71,6 +71,16 @@ Use this Docker `CMD`. This will start the proxy with 1 Uvicorn Async Worker CMD ["--port", "4000", "--config", "./proxy_server_config.yaml"] ``` +> Optional: If you observe gradual memory growth under sustained load, consider recycling workers after a fixed number of requests to mitigate leaks. Set this via CLI or environment variable: + +```shell +# CLI +CMD ["--port", "4000", "--config", "./proxy_server_config.yaml", "--max_requests_before_restart", "10000"] + +# or ENV (for deployment manifests / containers) +export MAX_REQUESTS_BEFORE_RESTART=10000 +``` + ## 4. Use Redis 'port','host', 'password'. NOT 'redis_url' diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 867a395d62..21e2c54ec7 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -185,6 +185,7 @@ class ProxyInitializationHelpers: num_workers: int, ssl_certfile_path: str, ssl_keyfile_path: str, + max_requests_before_restart: Optional[int] = None, ): """ Run litellm with `gunicorn` @@ -265,6 +266,10 @@ class ProxyInitializationHelpers: "access_log_format": '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s', } + # Optional: recycle workers after N requests to mitigate memory growth + if max_requests_before_restart is not None: + gunicorn_options["max_requests"] = max_requests_before_restart + if ssl_certfile_path is not None and ssl_keyfile_path is not None: print( # noqa f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" # noqa @@ -486,6 +491,13 @@ class ProxyInitializationHelpers: help="Set the uvicorn keepalive timeout in seconds (uvicorn timeout_keep_alive parameter)", envvar="KEEPALIVE_TIMEOUT", ) +@click.option( + "--max_requests_before_restart", + default=None, + type=int, + help="Restart worker after this many requests (uvicorn: limit_max_requests, gunicorn: max_requests)", + envvar="MAX_REQUESTS_BEFORE_RESTART", +) def run_server( # noqa: PLR0915 host, port, @@ -524,6 +536,7 @@ def run_server( # noqa: PLR0915 use_prisma_db_push: bool, skip_server_startup, keepalive_timeout, + max_requests_before_restart, ): args = locals() if local: @@ -813,6 +826,9 @@ def run_server( # noqa: PLR0915 log_config=log_config, keepalive_timeout=keepalive_timeout, ) + # Optional: recycle uvicorn workers after N requests + if max_requests_before_restart is not None: + uvicorn_args["limit_max_requests"] = max_requests_before_restart if run_gunicorn is False and run_hypercorn is False: if ssl_certfile_path is not None and ssl_keyfile_path is not None: print( # noqa @@ -837,6 +853,7 @@ def run_server( # noqa: PLR0915 num_workers=num_workers, ssl_certfile_path=ssl_certfile_path, ssl_keyfile_path=ssl_keyfile_path, + max_requests_before_restart=max_requests_before_restart, ) elif run_hypercorn is True: ProxyInitializationHelpers._init_hypercorn_server( diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4235e5d3ad..90d958e711 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -314,6 +314,51 @@ class TestProxyInitializationHelpers: call_args = mock_uvicorn_run.call_args assert call_args[1]["timeout_keep_alive"] == 30 + @patch("uvicorn.run") + @patch("builtins.print") + def test_max_requests_before_restart_flag(self, mock_print, mock_uvicorn_run): + """Test that the max_requests_before_restart flag is passed to uvicorn as limit_max_requests""" + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + + mock_app = MagicMock() + mock_proxy_config = MagicMock() + mock_key_mgmt = MagicMock() + mock_save_worker_config = MagicMock() + + with patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=mock_app, + ProxyConfig=mock_proxy_config, + KeyManagementSettings=mock_key_mgmt, + save_worker_config=mock_save_worker_config, + ) + }, + ), patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args: + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, ["--local", "--max_requests_before_restart", "123"] + ) + + assert result.exit_code == 0 + mock_uvicorn_run.assert_called_once() + + # Check that uvicorn.run was called with limit_max_requests parameter + call_args = mock_uvicorn_run.call_args + assert call_args[1]["limit_max_requests"] == 123 + @patch.dict(os.environ, {}, clear=True) def test_construct_database_url_from_env_vars(self): """Test the construct_database_url_from_env_vars function with various scenarios"""