diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 48358010485..11954d121a5 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -440,10 +440,10 @@ db: # Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks # the pod NotReady and blocks only until in-flight requests actually finish # (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always -# waiting the worst-case duration. The drain runs once (the preStop hook and -# the SIGTERM handler share it), so set terminationGracePeriodSeconds a few -# seconds above GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before -# SIGKILL. +# waiting the worst-case duration. After SIGTERM the server waits up to +# GRACEFUL_SHUTDOWN_TIMEOUT again for connections still open and then cancels +# them, so set terminationGracePeriodSeconds a few seconds above twice +# GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before SIGKILL. # # /health/drain is off by default; enable it with # general_settings.enable_drain_endpoint: true. The kubelet calls preStop diff --git a/litellm/constants.py b/litellm/constants.py index 6a6816eb4fa..a7ea4a2975b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -104,6 +104,7 @@ S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY" MAX_FILE_LIST_LIMIT: Final = 10000 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) +GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10 budget_reservation_disabled_info_emitted = False DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)) DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 7764d877b98..3c5f55d87a8 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -308,12 +308,14 @@ class ProxyInitializationHelpers: import litellm from litellm._logging import _get_uvicorn_json_log_config, resolve_log_level + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager uvicorn_args: Final = { "app": "litellm.proxy.proxy_server:app", "host": host, "port": port, "server_header": False, + "timeout_graceful_shutdown": GracefulShutdownManager.get_timeout(), } if log_config is not None: print(f"Using log_config: {log_config}") @@ -571,6 +573,12 @@ class ProxyInitializationHelpers: """ Run litellm with `gunicorn` """ + import math + + from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + from litellm.proxy.shutdown.graceful_uvicorn_worker import build_graceful_shutdown_uvicorn_worker + if os.name == "nt": pass else: @@ -627,10 +635,12 @@ class ProxyInitializationHelpers: return self.application print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} with {num_workers} workers\033[0m\n") + graceful_shutdown_timeout: Final = GracefulShutdownManager.get_timeout() gunicorn_options: Final = { "bind": f"{host}:{port}", "workers": num_workers, # default is 1 - "worker_class": "uvicorn.workers.UvicornWorker", + "worker_class": build_graceful_shutdown_uvicorn_worker(graceful_shutdown_timeout), + "graceful_timeout": math.ceil(graceful_shutdown_timeout) + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS, "preload": True, # Add the preload flag, "accesslog": "-", # Log to stdout "timeout": 600, # default to very high number, bedrock/anthropic.claude-v2:1 can take 30+ seconds for the 1st chunk to come in diff --git a/litellm/proxy/shutdown/graceful_shutdown_manager.py b/litellm/proxy/shutdown/graceful_shutdown_manager.py index d07ef05abf7..d7a07ec34e5 100644 --- a/litellm/proxy/shutdown/graceful_shutdown_manager.py +++ b/litellm/proxy/shutdown/graceful_shutdown_manager.py @@ -2,9 +2,10 @@ Application-level graceful shutdown coordination for the LiteLLM proxy. Kubernetes terminates a pod by sending ``SIGTERM`` and, after -``terminationGracePeriodSeconds``, ``SIGKILL``. By default LiteLLM delegates -the signal to uvicorn and tears down immediately, dropping any in-flight -requests (streaming, batch inference, long-lived calls). +``terminationGracePeriodSeconds``, ``SIGKILL``. On ``SIGTERM`` uvicorn stops +accepting connections and waits up to ``GRACEFUL_SHUTDOWN_TIMEOUT`` for open +ones, cancelling whatever is still running, before the lifespan teardown that +runs this drain. A fixed ``preStop`` sleep can not solve this: it has to be sized for the *worst-case* request, so it either wastes time on every routine shutdown or is diff --git a/litellm/proxy/shutdown/graceful_uvicorn_worker.py b/litellm/proxy/shutdown/graceful_uvicorn_worker.py new file mode 100644 index 00000000000..6ae5ae60c3f --- /dev/null +++ b/litellm/proxy/shutdown/graceful_uvicorn_worker.py @@ -0,0 +1,10 @@ +from uvicorn.workers import UvicornWorker + + +def build_graceful_shutdown_uvicorn_worker(timeout_graceful_shutdown: float) -> type[UvicornWorker]: + class GracefulShutdownUvicornWorker(UvicornWorker): + def __init__(self, *args: object) -> None: + self.CONFIG_KWARGS = {**UvicornWorker.CONFIG_KWARGS, "timeout_graceful_shutdown": timeout_graceful_shutdown} + super().__init__(*args) + + return GracefulShutdownUvicornWorker diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 33f0216ff85..9566b6e30a2 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -1,4 +1,5 @@ import inspect +import logging import os from contextlib import nullcontext from pathlib import Path @@ -22,8 +23,10 @@ import yaml from uvicorn.config import LOOP_FACTORIES from uvicorn.importer import import_from_string +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS from litellm.proxy import proxy_cli from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server +from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT @pytest.fixture(autouse=True) @@ -153,6 +156,22 @@ class TestProxyInitializationHelpers: ) assert args["timeout_worker_healthcheck"] == 15 + @pytest.mark.parametrize( + ("raw_timeout", "expected"), + [("7.5", 7.5), ("0", 0.0), (None, DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT)], + ) + def test_uvicorn_graceful_shutdown_is_bounded_by_graceful_shutdown_timeout( + self, monkeypatch, raw_timeout, expected + ): + if raw_timeout is None: + monkeypatch.delenv("GRACEFUL_SHUTDOWN_TIMEOUT", raising=False) + else: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", raw_timeout) + + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000) + + assert uvicorn.Config(**args).timeout_graceful_shutdown == expected + @staticmethod def _uvicorn_access_info_enabled(args: dict) -> bool: import logging @@ -2050,6 +2069,49 @@ class TestProxyInitializationHelpers: assert captured["options"]["max_requests"] == 1000 assert captured["options"]["max_requests_jitter"] == 50 + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") + def test_gunicorn_workers_bound_graceful_shutdown_by_graceful_shutdown_timeout(self, monkeypatch): + pytest.importorskip("gunicorn") + import gunicorn.config + import gunicorn.glogging + from uvicorn.workers import UvicornWorker + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "7.5") + captured: dict = {} + + def capture_run(self): + captured["options"] = dict(self.options) + + with ( + 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" + ), + ): + ProxyInitializationHelpers._run_gunicorn_server( + host="127.0.0.1", + port=4011, + app=MagicMock(), + num_workers=2, + ssl_certfile_path=None, + ssl_keyfile_path=None, + ) + + 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)) + gunicorn_cfg = gunicorn.config.Config() + worker = captured["options"]["worker_class"]( + 0, os.getpid(), [], MagicMock(), 30, gunicorn_cfg, gunicorn.glogging.Logger(gunicorn_cfg) + ) + + assert isinstance(worker, UvicornWorker) + assert worker.config.timeout_graceful_shutdown == 7.5 + assert worker.config.loop == UvicornWorker.CONFIG_KWARGS["loop"] + 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") 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"""