diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index ff649953e69..c897cd18684 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -473,8 +473,12 @@ class LoggingWorker: self.start() self.enqueue(async_coroutine) - async def stop(self) -> None: + async def stop(self, *, discard_pending: bool = False) -> 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): + task["coroutine"].close() + self._queue.task_done() if self._timeout_summary_task is not None: self._timeout_summary_task.cancel() self._timeout_summary_task = None @@ -496,6 +500,11 @@ class LoggingWorker: # Wait for cancellation to settle; ignore errors raised during shutdown. await asyncio.gather(*tasks_to_cancel, return_exceptions=True) + 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 # Drop references to completed tasks so we can restart cleanly. self._running_tasks.clear() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b263412ce0f..3e302f5035e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1268,14 +1268,30 @@ async def _flush_spend_logs_queue_on_shutdown() -> None: async def _flush_spend_on_shutdown() -> None: + from litellm.proxy.utils import update_daily_tag_spend + await flush_spend_counters_on_shutdown() + if prisma_client is not None: + await update_daily_tag_spend(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) await _flush_spend_logs_queue_on_shutdown() async def _flush_spend_around_logging_drain_on_shutdown( flush_spend: Callable[[], Awaitable[None]] = _flush_spend_on_shutdown, ) -> 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") + try: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=GracefulShutdownManager.get_timeout()) + except TimeoutError: + verbose_proxy_logger.warning("Timed out draining success callbacks on shutdown; accounting may be incomplete") + finally: + try: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.stop(discard_pending=True), timeout=5.0) + except TimeoutError: + verbose_proxy_logger.warning("Timed out cancelling success callbacks on shutdown") await drain_logging_tasks() await flush_spend() diff --git a/litellm/utils.py b/litellm/utils.py index be1e11dd39f..3ce394ce280 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1526,14 +1526,12 @@ def _schedule_async_success_logging( def _enqueue_async_logging() -> None: logging_obj.mark_async_success_log_enqueued() - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, ) if not logging_obj.defer_async_logging: @@ -1544,13 +1542,13 @@ def _schedule_async_success_logging( logging_obj.enqueue_deferred_logging = _enqueue_async_logging -async def _client_async_logging_helper( +def _client_async_logging_helper( logging_obj: LiteLLMLoggingObject, - result, - start_time, - end_time, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, is_completion_with_fallbacks: bool, -): +) -> None: if ( is_completion_with_fallbacks is False ): # don't log the parent event litellm.completion_with_fallbacks as a 'log_success_event', this will lead to double logging the same call - https://github.com/BerriAI/litellm/issues/7477 diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index e96511511d9..6a975fc02df 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -1,3 +1,4 @@ +import hashlib import json import os import threading @@ -26,6 +27,154 @@ MODEL_INSERT_ARRIVED_LOG_LINE: Final = "path=/model/new" COMMIT_DELAY_SECONDS: Final = 15 BURST_REQUESTS: Final = 30 +_GATED_LOGGER: Final = """ +import os +import httpx +from litellm.integrations.custom_logger import CustomLogger + +class ShutdownLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async with httpx.AsyncClient(timeout=25, trust_env=False) as client: + await client.post(os.environ["SHUTDOWN_CALLBACK_URL"]) + +logger = ShutdownLogger() +""" + + +def _tag_shutdown_config(directory: Path, upstream: str, redis_buffer: bool, *, gated: bool = False) -> Path: + path: Final = directory / "tag-shutdown.yaml" + if gated: + (directory / "shutdown_callbacks.py").write_text(_GATED_LOGGER) + path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": "shutdown-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": upstream + "/v1", + "api_key": "synthetic-key", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "proxy_batch_write_at": 3600, + "use_redis_transaction_buffer": redis_buffer, + "disable_model_info_refresh": True, + }, + "litellm_settings": { + "cache": True, + "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, + "callbacks": ["shutdown_callbacks.logger"] if gated else [], + }, + } + ) + ) + return path + + +def _tagged_chat(gateway: Gateway, key: str, tags: tuple[str, str]) -> None: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": "shutdown-model", + "messages": [{"role": "user", "content": uuid.uuid4().hex}], + "metadata": {"tags": list(tags)}, + }, + key=key, + ) + assert response.status_code == 200, response.text + + +def _assert_shutdown_accounting(key: str, tags: tuple[str, str], requests: int) -> None: + hashed: Final = hashlib.sha256(key.encode()).hexdigest() + for table, count in (("LiteLLM_SpendLogs", "count(*)"), ("LiteLLM_DailyUserSpend", "sum(api_requests)")): + rows: Final = read_rows( + f'SELECT ({count})::int AS requests, sum(spend) AS spend FROM "{table}" WHERE api_key=%s', (hashed,) + ) + assert rows[0]["requests"] == requests, (table, rows) + assert rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (table, rows) + key_rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashed,)) + assert key_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9) + for tag in tags: + tag_rows: Final = read_rows( + 'SELECT sum(api_requests)::int AS requests, sum(spend) AS spend FROM "LiteLLM_DailyTagSpend" WHERE api_key=%s AND tag=%s', + (hashed, tag), + ) + assert tag_rows[0]["requests"] == requests, (tag, tag_rows) + assert tag_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (tag, tag_rows) + + +@pytest.mark.parametrize("redis_buffer", [False, True]) +def test_tagged_request_survives_shutdown_before_the_tag_timer( + gateway: Gateway, tmp_path: Path, redis_buffer: bool +) -> None: + tags: Final = (f"shutdown-a-{uuid.uuid4().hex}", f"shutdown-b-{uuid.uuid4().hex}") + overrides: Final = {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_LOG": "DEBUG"} + with wire_server(_provider) as wire: + config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer) + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned: + key: Final = string_value( + owned.gateway.post("/key/generate", {"user_id": f"shutdown-{uuid.uuid4().hex}"})["key"] + ) + _tagged_chat(owned.gateway, key, tags) + owned.begin_shutdown() + owned.process.wait(timeout=90) + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as restarted: + _assert_shutdown_accounting(key, tags, 1) + restarted.begin_shutdown() + restarted.process.wait(timeout=90) + _assert_shutdown_accounting(key, tags, 1) + + +@pytest.mark.parametrize("redis_buffer", [False, True]) +def test_callback_backlog_is_accounted_before_shutdown_disconnects( + gateway: Gateway, tmp_path: Path, redis_buffer: bool +) -> 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: + config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer, gated=True) + with owned_proxy_process( + gateway, + tmp_path, + { + "DATABASE_URL": os.environ["DATABASE_URL"], + "LOGGING_WORKER_CONCURRENCY": "2", + "SHUTDOWN_CALLBACK_URL": logger.url, + "GRACEFUL_SHUTDOWN_TIMEOUT": "15", + "LITELLM_LOG": "DEBUG", + }, + config=config, + ) as owned: + key: Final = string_value( + owned.gateway.post("/key/generate", {"user_id": f"backlog-{uuid.uuid4().hex}"})["key"] + ) + try: + for _ in range(3): + _tagged_chat(owned.gateway, key, tags) + eventually(logger.received.qsize, lambda count: count == 2, seconds=10) + owned.begin_shutdown() + eventually( + lambda: ( + owned.process.poll() is not None + or "Draining success callbacks before the final spend flush" in owned.log.read_text() + ), + lambda ready: ready, + seconds=30, + ) + release.set() + owned.process.wait(timeout=90) + _assert_shutdown_accounting(key, tags, 3) + finally: + release.set() + def _api_requests(table: str, column: str, identity: str) -> int: rows: Final = read_rows( diff --git a/tests/unit/litellm_core_utils/test_logging_worker.py b/tests/unit/litellm_core_utils/test_logging_worker.py index f48495d1062..d7824d612ca 100644 --- a/tests/unit/litellm_core_utils/test_logging_worker.py +++ b/tests/unit/litellm_core_utils/test_logging_worker.py @@ -82,6 +82,56 @@ class _RecordCollector(logging.Handler): self.records.append(record) +@pytest.mark.asyncio +@pytest.mark.parametrize("running", [False, True]) +async def test_stop_discards_waiting_callbacks_without_executing_them(running: bool) -> None: + worker: Final = LoggingWorker(concurrency=1) + started: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + executed: Final = asyncio.Event() + + async def blocked() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + async def waiting() -> None: + executed.set() + + if running: + worker.ensure_initialized_and_enqueue(blocked()) + await started.wait() + else: + worker._ensure_queue() + worker.enqueue(waiting()) + await worker.stop(discard_pending=True) + assert not executed.is_set() + assert cancelled.is_set() == running + assert worker._queue is not None + assert worker._queue.empty() + await asyncio.wait_for(worker._queue.join(), timeout=1) + + +@pytest.mark.asyncio +async def test_stop_does_not_replay_a_callback_dequeued_before_its_processor_started() -> None: + worker: Final = LoggingWorker(concurrency=1) + executed: Final = asyncio.Event() + + async def callback() -> None: + executed.set() + + worker.ensure_initialized_and_enqueue(callback()) + await asyncio.sleep(0) + await worker.stop(discard_pending=True) + assert not executed.is_set() + worker._flush_on_exit() + assert not executed.is_set() + assert worker._queue is not None + await asyncio.wait_for(worker._queue.join(), timeout=1) + + class TestLoggingWorker: """Test cases for LoggingWorker functionality.""" diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index a26a64c6fb0..ce04c67355d 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -260,6 +260,131 @@ async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_comm 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.responses.streaming_iterator import spawn_logging_task + + started: Final = asyncio.Event() + release: Final = asyncio.Event() + first_flush: Final = asyncio.Event() + pending: Final[asyncio.Queue[str]] = asyncio.Queue() + persisted: Final[asyncio.Queue[str]] = asyncio.Queue() + + async def enqueue_spend() -> None: + pending.put_nowait("accounted request") + + async def callback() -> None: + started.set() + await release.wait() + spawn_logging_task(asyncio.get_running_loop(), enqueue_spend(), task_name="accounting") + + async def flush_spend() -> None: + while not pending.empty(): + persisted.put_nowait(pending.get_nowait()) + first_flush.set() + + GLOBAL_LOGGING_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) + ) + try: + await first_flush.wait() + waited_for_callback: Final = not shutdown.done() + release.set() + await shutdown + assert waited_for_callback, "Shutdown returned while accounting callbacks were still running" + assert persisted.get_nowait() == "accounted request" + assert pending.empty() + finally: + release.set() + await GLOBAL_LOGGING_WORKER.flush() + await GLOBAL_LOGGING_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 + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "0") + started: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + flushed_after_cancellation: Final = asyncio.Event() + + async def callback() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + async def flush_spend() -> None: + if cancelled.is_set(): + flushed_after_cancellation.set() + + GLOBAL_LOGGING_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) + assert cancelled.is_set() + assert flushed_after_cancellation.is_set() + assert "Timed out draining success callbacks" in caplog.text + finally: + await GLOBAL_LOGGING_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 + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "0") + started: Final = asyncio.Event() + cancelling: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + flushed: Final = asyncio.Event() + loop: Final = asyncio.get_running_loop() + clock: Final = loop.time + + async def callback() -> None: + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelling.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + async def flush_spend() -> None: + if cancelled.is_set(): + flushed.set() + + GLOBAL_LOGGING_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) + ) + await cancelling.wait() + monkeypatch.setattr(loop, "time", lambda: clock() + 6) + await shutdown + assert cancelled.is_set() + assert flushed.is_set() + assert "Timed out cancelling success callbacks" in caplog.text + finally: + monkeypatch.setattr(loop, "time", clock) + await GLOBAL_LOGGING_WORKER.stop() + + headers = {"Authorization": f"Bearer {token}"} example_completion_result = { diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 146e2df6413..a27166041bb 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10882,3 +10882,43 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_ ) assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred", [False, True]) +async def test_success_logging_is_enqueued_before_the_worker_is_flushed(deferred: bool) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.utils import _schedule_async_success_logging + + completed: Final[asyncio.Queue[object]] = asyncio.Queue() + + async def callback(kwargs: object, response: object, start: object, end: object) -> None: + completed.put_nowait(response) + + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + logging_obj: Final = Logging( + model="synthetic-model", + messages=[], + stream=False, + call_type="acompletion", + start_time=now.timestamp(), + litellm_call_id="logging-enqueue", + function_id="logging-enqueue", + dynamic_async_success_callbacks=[callback], + ) + logging_obj.defer_async_logging = deferred + result: Final = litellm.ModelResponse(id="accounted-response") + try: + _schedule_async_success_logging(logging_obj, result, now, now, False) + if deferred: + assert completed.empty() + assert logging_obj.enqueue_deferred_logging is not None + logging_obj.enqueue_deferred_logging() + await GLOBAL_LOGGING_WORKER.flush() + assert completed.get_nowait() is result + assert completed.empty() + finally: + await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() + await GLOBAL_LOGGING_WORKER.stop()