From a81f4e8037a02fec80f61202be9760feb9f90214 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:27:03 -0700 Subject: [PATCH] fix: share shutdown callback budget with HTTP draining --- helm/litellm-helm/README.md | 6 + litellm/litellm_core_utils/logging_worker.py | 21 +++- litellm/proxy/db/db_spend_update_writer.py | 1 - litellm/proxy/proxy_server.py | 37 ++++-- .../shutdown/graceful_shutdown_manager.py | 57 ++++++++-- litellm/responses/streaming_iterator.py | 2 +- .../integration/spend/test_shutdown_flush.py | 74 +++++++++++- .../test_graceful_shutdown_manager.py | 104 +++++++++++++++++ tests/unit/proxy/test_proxy_server.py | 107 ++++++++++++------ 9 files changed, 347 insertions(+), 62 deletions(-) diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index e059d45e037..a3488d6114c 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -259,3 +259,9 @@ itself. `lensWorker.mode` selects `disabled`, `bundled` or `external`. Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for complete values, verification, GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published + +## Shutdown accounting + +`GRACEFUL_SHUTDOWN_TIMEOUT` defaults to 30 seconds. Pending success callbacks share the time remaining since SIGTERM or the first `/health/drain` call. Callback flushing reserves the existing 10-second teardown headroom for cancellation, tracked accounting tasks and persistence. Exhausting this budget can discard callbacks, which is logged + +Set `terminationGracePeriodSeconds` to at least `GRACEFUL_SHUTDOWN_TIMEOUT + 10`, allowing additional time for preStop hooks and database writes. This is not a hard shutdown bound or a guarantee of persistence. Database retries retain their own limits, and cancellation-resistant callbacks can remain alive. The existing preStop hook can also be followed by Uvicorn's separate connection drain, as described in `values.yaml` diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index c897cd18684..b2156216416 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -473,12 +473,19 @@ class LoggingWorker: self.start() self.enqueue(async_coroutine) - async def stop(self, *, discard_pending: bool = False) -> None: + @property + def pending_count(self) -> int: + return (self._queue.qsize() if self._queue is not None else 0) + len(self._dequeued_tasks) + + async def stop(self, *, discard_pending: bool = False, timeout: float | None = None) -> None: """Stop the logging worker and clean up resources.""" if discard_pending and self._queue is not None: - for task in self._drain_pending(self._queue): + discarded: Final = self._drain_pending(self._queue) + for task in discarded: task["coroutine"].close() self._queue.task_done() + if discarded: + verbose_logger.warning("LoggingWorker shutdown: %d queued callbacks discarded", len(discarded)) if self._timeout_summary_task is not None: self._timeout_summary_task.cancel() self._timeout_summary_task = None @@ -498,16 +505,20 @@ class LoggingWorker: task.cancel() # Wait for cancellation to settle; ignore errors raised during shutdown. - await asyncio.gather(*tasks_to_cancel, return_exceptions=True) + done, pending = await asyncio.wait(tasks_to_cancel, timeout=timeout) + await asyncio.gather(*done, return_exceptions=True) + if pending: + verbose_logger.warning("LoggingWorker shutdown: %d tasks still alive after cancellation", len(pending)) if discard_pending and self._queue is not None: for task in self._unstarted_dequeued_tasks(): task["coroutine"].close() self._untrack_dequeued(task) self._queue.task_done() - self._worker_task = None + if self._worker_task is not None and self._worker_task.done(): + self._worker_task = None # Drop references to completed tasks so we can restart cleanly. - self._running_tasks.clear() + self._running_tasks.difference_update(done) async def flush(self) -> None: """Flush the logging queue. diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 8ac2597c99b..538779f11b4 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1900,7 +1900,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) - # NOTE: Daily tag spend is committed by a separate scheduler job. ################## Daily End-User Spend Update Transactions ################## # Aggregate all in memory daily end-user spend transactions and commit to db diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bf822876383..df9fe5cac5e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -79,6 +79,7 @@ from litellm.constants import ( DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, DEFAULT_SLACK_ALERTING_THRESHOLD, + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS, INTERNAL_CALLBACK_NAMES, LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS, LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, @@ -98,6 +99,7 @@ from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, initialize_sentry, ) +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, offload_token_count @@ -1278,21 +1280,29 @@ async def _flush_spend_on_shutdown() -> None: async def _flush_spend_around_logging_drain_on_shutdown( flush_spend: Callable[[], Awaitable[None]] = _flush_spend_on_shutdown, + *, + worker: LoggingWorker = GLOBAL_LOGGING_WORKER, + clock: Callable[[], float] = time.monotonic, + drain_tasks: Callable[[float], Awaitable[None]] = drain_logging_tasks, ) -> None: - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - await flush_spend() - verbose_proxy_logger.info("Draining success callbacks before the final spend flush") + remaining: Final = GracefulShutdownManager.remaining(clock=clock) + callback_timeout: Final = max(0.0, remaining - GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS) + verbose_proxy_logger.info( + "Draining success callbacks before the final spend flush remaining_s=%.3f callback_timeout_s=%.3f", + remaining, + callback_timeout, + ) try: - await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=GracefulShutdownManager.get_timeout()) + await asyncio.wait_for(worker.flush(), timeout=callback_timeout) except asyncio.TimeoutError: - verbose_proxy_logger.warning("Timed out draining success callbacks on shutdown; accounting may be incomplete") + verbose_proxy_logger.warning( + "Success callback shutdown budget exhausted with %d unfinished callbacks; accounting may be incomplete", + worker.pending_count, + ) finally: - try: - await asyncio.wait_for(GLOBAL_LOGGING_WORKER.stop(discard_pending=True), timeout=5.0) - except asyncio.TimeoutError: - verbose_proxy_logger.warning("Timed out cancelling success callbacks on shutdown") - await drain_logging_tasks() + await worker.stop(discard_pending=True, timeout=GracefulShutdownManager.remaining(2.0, clock=clock)) + await drain_tasks(GracefulShutdownManager.remaining(5.0, clock=clock)) await flush_spend() @@ -1455,6 +1465,13 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries: @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: + with GracefulShutdownManager.capture_sigterm(): + async with _proxy_startup_event(app) as state: + yield state + + +@asynccontextmanager +async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: global \ prisma_client, \ master_key, \ diff --git a/litellm/proxy/shutdown/graceful_shutdown_manager.py b/litellm/proxy/shutdown/graceful_shutdown_manager.py index d7a07ec34e5..fdc5bf07406 100644 --- a/litellm/proxy/shutdown/graceful_shutdown_manager.py +++ b/litellm/proxy/shutdown/graceful_shutdown_manager.py @@ -5,7 +5,8 @@ Kubernetes terminates a pod by sending ``SIGTERM`` and, after ``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. +runs this drain. The lifespan registers a SIGTERM wrapper before serving +requests so callback draining can subtract the preceding connection 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 @@ -20,8 +21,12 @@ granularity of ``InFlightRequestsMiddleware``. import asyncio import os +import signal +import threading import time -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from types import FrameType from typing import Final from litellm._logging import verbose_proxy_logger @@ -44,6 +49,7 @@ class GracefulShutdownManager: _is_shutting_down: bool = False _shutdown_started_at: float | None = None _drain_performed: bool = False + _sigterm_handler: Callable[[int, FrameType | None], None] | None = None @classmethod def is_shutting_down(cls) -> bool: @@ -71,7 +77,7 @@ class GracefulShutdownManager: return DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT @classmethod - def start_shutdown(cls) -> None: + def start_shutdown(cls, clock: Callable[[], float] = time.monotonic) -> None: """ Mark the worker as draining. Idempotent — repeated calls (e.g. SIGTERM followed by a preStop hit on /health/drain) do not reset the clock. @@ -79,12 +85,42 @@ class GracefulShutdownManager: if cls._is_shutting_down: return cls._is_shutting_down = True - cls._shutdown_started_at = time.monotonic() + cls._shutdown_started_at = clock() verbose_proxy_logger.info( "graceful_shutdown_started in_flight_requests=%s", get_in_flight_requests(), ) + @classmethod + def remaining(cls, cap: float | None = None, *, clock: Callable[[], float] = time.monotonic) -> float: + elapsed: Final = 0.0 if cls._shutdown_started_at is None else clock() - cls._shutdown_started_at + remaining: Final = max(0.0, cls.get_timeout() - elapsed) + return remaining if cap is None else min(max(0.0, cap), remaining) + + @classmethod + @contextmanager + def capture_sigterm(cls, clock: Callable[[], float] = time.monotonic) -> Iterator[None]: + if threading.current_thread() is not threading.main_thread(): + yield + return + previous: Final = signal.getsignal(signal.SIGTERM) + if not callable(previous) or previous is cls._sigterm_handler: + yield + return + + def handle_sigterm(signum: int, frame: FrameType | None) -> None: + cls.start_shutdown(clock) + previous(signum, frame) + + cls._sigterm_handler = handle_sigterm + signal.signal(signal.SIGTERM, handle_sigterm) + try: + yield + finally: + if signal.getsignal(signal.SIGTERM) is handle_sigterm: + signal.signal(signal.SIGTERM, previous) + cls._sigterm_handler = None + @classmethod async def wait_for_drain( cls, @@ -93,13 +129,14 @@ class GracefulShutdownManager: count_fn: Callable[[], int] | None = None, poll_interval: float = _DRAIN_POLL_INTERVAL, log_interval: float = _DRAIN_LOG_INTERVAL, + clock: Callable[[], float] = time.monotonic, ) -> int: """ Poll the in-flight request counter until it reaches the drain target or ``timeout`` seconds elapse. Args: - timeout: Max seconds to wait. Defaults to ``get_timeout()``. + timeout: Max seconds to wait. Defaults to the remaining shutdown budget. exclude_self: When the caller is itself an in-flight HTTP request (the /health/drain endpoint), set this so the caller's own request is not counted as outstanding work. @@ -120,7 +157,7 @@ class GracefulShutdownManager: cls._drain_performed = True if timeout is None: - timeout = cls.get_timeout() + timeout = cls.remaining(clock=clock) if count_fn is None: count_fn = get_in_flight_requests @@ -128,7 +165,7 @@ class GracefulShutdownManager: # and so counts itself; treat <=1 as "drained" in that case. target: Final = 1 if exclude_self else 0 - start: Final = time.monotonic() + start: Final = clock() initial: Final = count_fn() last_log = start @@ -142,11 +179,11 @@ class GracefulShutdownManager: verbose_proxy_logger.info( "graceful_shutdown_complete drained_requests=%s elapsed_s=%.2f", drained, - time.monotonic() - start, + clock() - start, ) return drained - elapsed = time.monotonic() - start + elapsed = clock() - start if elapsed >= timeout: verbose_proxy_logger.warning( "graceful_shutdown_timeout in_flight_requests=%s elapsed_s=%.2f " @@ -157,7 +194,7 @@ class GracefulShutdownManager: ) return max(0, initial - current) - now = time.monotonic() + now = clock() if now - last_log >= log_interval: verbose_proxy_logger.info( "drain_waiting in_flight_requests=%s elapsed_s=%.2f", diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 4e418a91156..c5f71aba343 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -222,7 +222,7 @@ async def drain_logging_tasks(timeout: float = 5.0) -> None: return _, pending = await asyncio.wait(tuple(_PENDING_LOGGING_TASKS), timeout=timeout) if pending: - verbose_logger.warning("Timed out draining %d logging tasks on shutdown, dropping them", len(pending)) + verbose_logger.warning("Timed out draining logging tasks on shutdown: %d tasks still alive", len(pending)) _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType( diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index 6a975fc02df..0dda707dd9c 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -1,12 +1,14 @@ import hashlib import json import os +import re import threading +import time import uuid from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import Final @@ -29,6 +31,7 @@ BURST_REQUESTS: Final = 30 _GATED_LOGGER: Final = """ import os +import re import httpx from litellm.integrations.custom_logger import CustomLogger @@ -139,7 +142,10 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects( ) -> None: release: Final = threading.Event() tags: Final = (f"backlog-a-{uuid.uuid4().hex}", f"backlog-b-{uuid.uuid4().hex}") - with wire_server(_provider) as wire, wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release)) as logger: + with ( + wire_server(_provider) as wire, + wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release)) as logger, + ): config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer, gated=True) with owned_proxy_process( gateway, @@ -176,6 +182,70 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects( release.set() +@pytest.mark.parametrize("gunicorn", [False, True]) +def test_sigterm_http_drain_consumes_callback_budget(gateway: Gateway, tmp_path: Path, gunicorn: bool) -> None: + release_request: Final = threading.Event() + release_callback: Final = threading.Event() + tags: Final = (f"deadline-a-{uuid.uuid4().hex}", f"deadline-b-{uuid.uuid4().hex}") + + def held_provider(request: Request) -> Reply: + response: Final = _provider(request) + return replace(response, chunks=(response.body,), gate_before_first=release_request) + + with ( + wire_server(held_provider) as wire, + wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release_callback)) as logger, + ): + config: Final = _tag_shutdown_config(tmp_path, wire.url, False, gated=True) + with ( + owned_proxy_process( + gateway, + tmp_path, + { + "DATABASE_URL": os.environ["DATABASE_URL"], + "SHUTDOWN_CALLBACK_URL": logger.url, + "GRACEFUL_SHUTDOWN_TIMEOUT": "30", + "LITELLM_LOG": "DEBUG", + }, + config=config, + extra_arguments=("--run_gunicorn",) if gunicorn else (), + ) as owned, + ThreadPoolExecutor(max_workers=1) as executor, + ): + key: Final = string_value( + owned.gateway.post("/key/generate", {"user_id": f"deadline-{uuid.uuid4().hex}"})["key"] + ) + request: Final = executor.submit(_tagged_chat, owned.gateway, key, tags) + try: + eventually(wire.received.qsize, lambda count: count == 1, seconds=10) + owned.begin_shutdown() + eventually( + lambda: "graceful_shutdown_started" in owned.log.read_text(), lambda stamped: stamped, seconds=5 + ) + assert not request.done(), "The shutdown clock must start before the HTTP request drains" + stamped: Final = time.monotonic() + eventually(time.monotonic, lambda now: now - stamped >= 2, seconds=5) + held_seconds: Final = time.monotonic() - stamped + release_request.set() + request.result(timeout=10) + budget: Final = eventually( + lambda: re.search(r"remaining_s=([\d.]+) callback_timeout_s=([\d.]+)", owned.log.read_text()), + lambda match: match is not None, + seconds=15, + ) + assert budget is not None + remaining: Final = float(budget[1]) + allowance: Final = float(budget[2]) + assert remaining <= 30 - held_seconds + 0.05 + assert allowance == pytest.approx(max(0, remaining - 10), abs=0.002) + release_callback.set() + owned.process.wait(timeout=45) + _assert_shutdown_accounting(key, tags, 1) + finally: + release_request.set() + release_callback.set() + + def _api_requests(table: str, column: str, identity: str) -> int: rows: Final = read_rows( f'SELECT coalesce(sum(api_requests), 0)::int AS total FROM "{table}" WHERE {column}=%s', (identity,) diff --git a/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py b/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py index d38852617b9..fdf1adda1f4 100644 --- a/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py +++ b/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py @@ -179,3 +179,107 @@ async def test_emits_periodic_drain_waiting_log_while_waiting(): timeout=10, count_fn=count_fn, poll_interval=0, log_interval=0 ) assert drained == 3 + + +@pytest.mark.parametrize("elapsed, cap, expected", [(4, None, 26), (25, 2, 2), (29, 5, 1), (31, None, 0)]) +def test_remaining_budget_includes_time_before_lifespan( + monkeypatch: pytest.MonkeyPatch, elapsed: float, cap: float | None, expected: float +) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0 + elapsed) + assert GracefulShutdownManager.remaining(cap, clock=lambda: 100.0 + elapsed) == expected + + +def test_remaining_budget_before_shutdown(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + assert GracefulShutdownManager.remaining(clock=lambda: 100.0) == 30.0 + assert GracefulShutdownManager.remaining(-1, clock=lambda: 100.0) == 0.0 + + +@pytest.mark.parametrize("prestop", [False, True]) +def test_sigterm_stamps_before_chaining_and_restores_server_handler( + monkeypatch: pytest.MonkeyPatch, prestop: bool +) -> None: + import signal + from queue import SimpleQueue + from types import FrameType + from typing import Final + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + calls: Final[SimpleQueue[float]] = SimpleQueue() + + def server_handler(signum: int, frame: FrameType | None) -> None: + assert signum == signal.SIGTERM + calls.put(GracefulShutdownManager.remaining(clock=lambda: 105.0)) + + original: Final = signal.signal(signal.SIGTERM, server_handler) + try: + if prestop: + GracefulShutdownManager.start_shutdown(clock=lambda: 95.0) + with GracefulShutdownManager.capture_sigterm(clock=lambda: 100.0): + installed: Final = signal.getsignal(signal.SIGTERM) + with GracefulShutdownManager.capture_sigterm(clock=lambda: 102.0): + assert signal.getsignal(signal.SIGTERM) is installed + signal.raise_signal(signal.SIGTERM) + signal.raise_signal(signal.SIGTERM) + assert calls.get_nowait() == (20.0 if prestop else 25.0) + assert calls.get_nowait() == (20.0 if prestop else 25.0) + assert calls.empty() + assert signal.getsignal(signal.SIGTERM) is server_handler + finally: + signal.signal(signal.SIGTERM, original) + + +def test_sigterm_capture_preserves_a_replacement_handler() -> None: + import signal + from types import FrameType + from typing import Final + + def server_handler(signum: int, frame: FrameType | None) -> None: + pass + + original: Final = signal.signal(signal.SIGTERM, server_handler) + try: + with GracefulShutdownManager.capture_sigterm(): + signal.signal(signal.SIGTERM, signal.SIG_IGN) + assert signal.getsignal(signal.SIGTERM) is signal.SIG_IGN + finally: + signal.signal(signal.SIGTERM, original) + + +@pytest.mark.parametrize("handler", ["SIG_DFL", "SIG_IGN"]) +def test_sigterm_capture_preserves_default_and_ignored_dispositions(handler: str) -> None: + import signal + from typing import Final + + disposition: Final = signal.SIG_DFL if handler == "SIG_DFL" else signal.SIG_IGN + original: Final = signal.signal(signal.SIGTERM, disposition) + try: + with GracefulShutdownManager.capture_sigterm(): + assert signal.getsignal(signal.SIGTERM) is disposition + finally: + signal.signal(signal.SIGTERM, original) + + +def test_sigterm_capture_in_worker_thread_leaves_signal_ownership_unchanged() -> None: + import signal + from concurrent.futures import ThreadPoolExecutor + from typing import Final + + original: Final = signal.getsignal(signal.SIGTERM) + + def run_lifespan() -> None: + with GracefulShutdownManager.capture_sigterm(): + assert signal.getsignal(signal.SIGTERM) is original + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(run_lifespan).result() + assert signal.getsignal(signal.SIGTERM) is original + + +@pytest.mark.asyncio +async def test_lifespan_request_drain_does_not_restart_an_exhausted_budget(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) + assert await GracefulShutdownManager.wait_for_drain(count_fn=lambda: 2, clock=lambda: 130.0) == 2 diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index ce04c67355d..5a064d3190b 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -263,9 +263,10 @@ token = MASTER_KEY @pytest.mark.asyncio async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning() -> None: - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.litellm_core_utils.logging_worker import LoggingWorker from litellm.responses.streaming_iterator import spawn_logging_task + worker: Final = LoggingWorker() started: Final = asyncio.Event() release: Final = asyncio.Event() first_flush: Final = asyncio.Event() @@ -285,10 +286,10 @@ async def test_shutdown_persists_accounting_from_worker_callbacks_before_returni persisted.put_nowait(pending.get_nowait()) first_flush.set() - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(callback()) + worker.ensure_initialized_and_enqueue(callback()) await started.wait() shutdown: Final = asyncio.create_task( - litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend) + litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker) ) try: await first_flush.wait() @@ -300,20 +301,23 @@ async def test_shutdown_persists_accounting_from_worker_callbacks_before_returni assert pending.empty() finally: release.set() - await GLOBAL_LOGGING_WORKER.flush() - await GLOBAL_LOGGING_WORKER.stop() + await worker.flush() + await worker.stop() @pytest.mark.asyncio async def test_shutdown_reports_worker_timeout_and_flushes_completed_accounting( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager - monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "0") + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + worker: Final = LoggingWorker() started: Final = asyncio.Event() cancelled: Final = asyncio.Event() flushed_after_cancellation: Final = asyncio.Event() + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) async def callback() -> None: started.set() @@ -326,63 +330,100 @@ async def test_shutdown_reports_worker_timeout_and_flushes_completed_accounting( if cancelled.is_set(): flushed_after_cancellation.set() - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(callback()) + worker.ensure_initialized_and_enqueue(callback()) await started.wait() try: with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend) + await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown( + flush_spend, worker=worker, clock=lambda: 120.0 + ) assert cancelled.is_set() assert flushed_after_cancellation.is_set() - assert "Timed out draining success callbacks" in caplog.text + assert "budget exhausted with 1 unfinished callbacks" in caplog.text finally: - await GLOBAL_LOGGING_WORKER.stop() + GracefulShutdownManager.reset() + await worker.stop() @pytest.mark.asyncio async def test_shutdown_bounds_callback_cancellation_and_flushes_accounting( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager - monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "0") + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + worker: Final = LoggingWorker() started: Final = asyncio.Event() cancelling: Final = asyncio.Event() - cancelled: Final = asyncio.Event() + release: Final = asyncio.Event() + finished: Final = asyncio.Event() flushed: Final = asyncio.Event() - loop: Final = asyncio.get_running_loop() - clock: Final = loop.time + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) async def callback() -> None: started.set() try: - await asyncio.Event().wait() + await release.wait() except asyncio.CancelledError: cancelling.set() - try: - await asyncio.Event().wait() - finally: - cancelled.set() + await release.wait() + finally: + finished.set() + + async def drain_tasks(timeout: float) -> None: + assert timeout == 0.0 + await cancelling.wait() async def flush_spend() -> None: - if cancelled.is_set(): + if cancelling.is_set(): flushed.set() - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(callback()) + worker.ensure_initialized_and_enqueue(callback()) await started.wait() try: - with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - shutdown: Final = asyncio.create_task( - litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown( + flush_spend, worker=worker, clock=lambda: 130.0, drain_tasks=drain_tasks ) - await cancelling.wait() - monkeypatch.setattr(loop, "time", lambda: clock() + 6) - await shutdown - assert cancelled.is_set() + assert not finished.is_set(), "The cancellation-resistant callback must remain observable" assert flushed.is_set() - assert "Timed out cancelling success callbacks" in caplog.text + assert "tasks still alive after cancellation" in caplog.text finally: - monkeypatch.setattr(loop, "time", clock) - await GLOBAL_LOGGING_WORKER.stop() + release.set() + await worker.flush() + await worker.stop() + GracefulShutdownManager.reset() + + +@pytest.mark.asyncio +async def test_shutdown_preserves_persistence_reserve_after_partial_http_drain( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + worker: Final = LoggingWorker() + persisted: Final = asyncio.Event() + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) + + async def flush_spend() -> None: + persisted.set() + + async def drain_tasks(timeout: float) -> None: + assert timeout == 5.0 + + try: + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown( + flush_spend, worker=worker, clock=lambda: 106.0, drain_tasks=drain_tasks + ) + assert "remaining_s=24.000 callback_timeout_s=14.000" in caplog.text + assert persisted.is_set() + finally: + await worker.stop() + GracefulShutdownManager.reset() headers = {"Authorization": f"Bearer {token}"}