diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c44f5602c95..4fea3757b94 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -210,6 +210,7 @@ def generate_feedback_box(): import contextlib +import dataclasses from collections import defaultdict from contextlib import asynccontextmanager from functools import lru_cache @@ -783,9 +784,58 @@ def cleanup_router_config_variables(): prisma_client = None +async def _stop_spend_background_jobs() -> None: + """Stop the scheduler and spend-logs queue monitor so nothing writes to the DB + while (or after) the shutdown flush runs.""" + if scheduler is not None: + try: + scheduler.shutdown(wait=False) + except Exception as e: # noqa: BLE001 # shutdown must never be blocked by scheduler teardown + verbose_proxy_logger.exception(f"Error stopping APScheduler on shutdown: {e}") + + monitor_task = spend_logs_queue_monitor.task + if monitor_task is not None: + monitor_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await monitor_task + spend_logs_queue_monitor.task = None + + +async def _flush_spend_buffers_on_shutdown() -> None: + """Commit in-memory spend buffers before the DB connection goes away. + + Without this, everything accumulated since the last scheduler run (key/team/user/ + end_user spend, spend logs, daily tag spend) is dropped on every restart.""" + if prisma_client is None: + return + + from litellm.proxy.utils import update_daily_tag_spend, update_spend + + try: + await update_spend( + prisma_client=prisma_client, + db_writer_client=db_writer_client, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # a failed flush must not stop the rest of shutdown + verbose_proxy_logger.exception(f"Error flushing spend updates on shutdown: {e}") + + try: + await update_daily_tag_spend( + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # a failed flush must not stop the rest of shutdown + verbose_proxy_logger.exception(f"Error flushing daily tag spend on shutdown: {e}") + + async def proxy_shutdown_event(): global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") + + await _stop_spend_background_jobs() + await _flush_spend_buffers_on_shutdown() + if prisma_client: verbose_proxy_logger.debug("Disconnecting from Prisma") await prisma_client.disconnect() @@ -2028,6 +2078,14 @@ celery_fn = None # Redis Queue for handling requests scheduler = None last_model_cost_map_reload = None + +@dataclasses.dataclass(slots=True) +class _SpendLogsQueueMonitorHandle: + task: Optional["asyncio.Task[None]"] = None + + +spend_logs_queue_monitor = _SpendLogsQueueMonitorHandle() + # Global variable for anthropic beta headers reload scheduling last_anthropic_beta_headers_reload = None @@ -8002,7 +8060,7 @@ class ProxyStartupEvent: from litellm.proxy.utils import _monitor_spend_logs_queue # Start background task to monitor spend logs queue size - asyncio.create_task( + spend_logs_queue_monitor.task = asyncio.create_task( _monitor_spend_logs_queue( prisma_client=prisma_client, db_writer_client=db_writer_client, diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a3f5049ef1d..908af273584 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -123,6 +123,116 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): } +@pytest.mark.asyncio +async def test_proxy_shutdown_event_flushes_spend_buffers_before_disconnect(monkeypatch): + """Regression for #34805: in-memory spend buffers must be committed before the + prisma connection is torn down, and the scheduler / queue monitor must be stopped + first so nothing is mid-flight against a closing connection.""" + calls: List[str] = [] + + fake_prisma = MagicMock() + + async def _disconnect(): + calls.append("disconnect") + + fake_prisma.disconnect = _disconnect + monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) + + fake_jwt = MagicMock() + fake_jwt.close = AsyncMock() + monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + + fake_scheduler = MagicMock() + fake_scheduler.shutdown = MagicMock(side_effect=lambda **_: calls.append("scheduler_shutdown")) + monkeypatch.setattr(ps, "scheduler", fake_scheduler, raising=False) + + monitor_started = asyncio.Event() + + async def _monitor(): + monitor_started.set() + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + calls.append("monitor_cancelled") + raise + + monitor_task = asyncio.create_task(_monitor()) + await monitor_started.wait() + monkeypatch.setattr(ps.spend_logs_queue_monitor, "task", monitor_task, raising=False) + + import litellm.proxy.utils as proxy_utils + + async def _update_spend(**kwargs): + calls.append("update_spend") + + async def _update_daily_tag_spend(**kwargs): + calls.append("update_daily_tag_spend") + + monkeypatch.setattr(proxy_utils, "update_spend", _update_spend, raising=False) + monkeypatch.setattr(proxy_utils, "update_daily_tag_spend", _update_daily_tag_spend, raising=False) + + import litellm + + monkeypatch.setattr(litellm, "cache", None, raising=False) + monkeypatch.setattr(litellm, "success_callback", [], raising=False) + + await proxy_shutdown_event() + + assert calls == [ + "scheduler_shutdown", + "monitor_cancelled", + "update_spend", + "update_daily_tag_spend", + "disconnect", + ] + assert monitor_task.cancelled() + assert ps.spend_logs_queue_monitor.task is None + + +@pytest.mark.asyncio +async def test_proxy_shutdown_event_flush_failure_still_disconnects(monkeypatch): + """A failing flush must not block the rest of shutdown; both buffers are + attempted and prisma is still disconnected.""" + calls: List[str] = [] + + fake_prisma = MagicMock() + + async def _disconnect(): + calls.append("disconnect") + + fake_prisma.disconnect = _disconnect + monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) + monkeypatch.setattr(ps, "scheduler", None, raising=False) + monkeypatch.setattr(ps.spend_logs_queue_monitor, "task", None, raising=False) + + fake_jwt = MagicMock() + fake_jwt.close = AsyncMock() + monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + + import litellm.proxy.utils as proxy_utils + + async def _update_spend(**kwargs): + calls.append("update_spend") + raise RuntimeError("db unreachable") + + async def _update_daily_tag_spend(**kwargs): + calls.append("update_daily_tag_spend") + + monkeypatch.setattr(proxy_utils, "update_spend", _update_spend, raising=False) + monkeypatch.setattr(proxy_utils, "update_daily_tag_spend", _update_daily_tag_spend, raising=False) + + import litellm + + monkeypatch.setattr(litellm, "cache", None, raising=False) + monkeypatch.setattr(litellm, "success_callback", [], raising=False) + + await proxy_shutdown_event() + + assert calls == ["update_spend", "update_daily_tag_spend", "disconnect"] + + @pytest.mark.asyncio async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): fake_prisma = MagicMock()