diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5efc63140d3..616ca638a80 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1097,6 +1097,27 @@ async def _flush_spend_logs_queue_on_shutdown() -> None: verbose_proxy_logger.exception("Error flushing spend logs queue on shutdown: %s", e) +async def _drain_reports_then_flush_spend( + drain_reports: Callable[[], Awaitable[None]], + drain_spend_events: Callable[[], Awaitable[None]], + stop_scheduler_jobs: Callable[[], Awaitable[None]] | None, + flush_spend_counters: Callable[[], Awaitable[None]], + flush_spend_logs: Callable[[], Awaitable[None]], +) -> None: + try: + await drain_reports() + except Exception as e: # noqa: BLE001 # shutdown must continue when a report drain fails + verbose_proxy_logger.error("Error draining passthrough upstream error reports: %s", e) + await drain_spend_events() + if stop_scheduler_jobs is not None: + try: + await stop_scheduler_jobs() + except Exception as e: + verbose_proxy_logger.error("Error stopping in-flight scheduled jobs: %s", e) + await flush_spend_counters() + await flush_spend_logs() + + async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None) -> None: 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") @@ -1585,23 +1606,18 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: except Exception as e: verbose_proxy_logger.error("Error stopping the spend view setup task: %s", e) - try: - await drain_passthrough_upstream_error_reports() - except Exception as e: # noqa: BLE001 # shutdown must continue when a report drain fails - verbose_proxy_logger.error("Error draining passthrough upstream error reports: %s", e) - - await _drain_spend_event_producer_on_shutdown() - # Shutdown event - finish or cancel in-flight scheduled jobs before the shutdown flushes and the DB disconnect - if scheduler is not None and scheduler_executor is not None: - try: - await stop_in_flight_scheduler_jobs(scheduler, scheduler_executor) - except Exception as e: - verbose_proxy_logger.error("Error stopping in-flight scheduled jobs: %s", e) - - await flush_spend_counters_on_shutdown() - - await _flush_spend_logs_queue_on_shutdown() + await _drain_reports_then_flush_spend( + drain_reports=drain_passthrough_upstream_error_reports, + drain_spend_events=_drain_spend_event_producer_on_shutdown, + stop_scheduler_jobs=( + partial(stop_in_flight_scheduler_jobs, scheduler, scheduler_executor) + if scheduler is not None and scheduler_executor is not None + else None + ), + flush_spend_counters=flush_spend_counters_on_shutdown, + flush_spend_logs=_flush_spend_logs_queue_on_shutdown, + ) await proxy_config.stop_config_sync_subscriber() diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index 1fb2fe6153b..3cbbe140f2c 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -343,42 +343,94 @@ async def test_flush_spend_counters_on_shutdown_logs_and_swallows_commit_errors( assert "Error flushing spend counters on shutdown: db gone" in caplog.text -def test_shutdown_drains_passthrough_error_reports_before_spend_flushes(): - """Passthrough error report callbacks write spend rows through the logging - worker, so the drain must complete before the spend producer, counters and - spend-log queue are flushed or the delivered rows can be skipped. The drain - lives inside the ``proxy_startup_event`` lifespan teardown which cannot be - driven without running the whole startup, so assert the await order in - source: a revert of the ordering is what this guards. - """ - import ast +def _recorded_step(calls: List[str], name: str, report_done: "asyncio.Event") -> Callable[[], Awaitable[None]]: + async def _step() -> None: + calls.append(f"{name}:report_done={report_done.is_set()}") - parsed = ast.parse(inspect.getsource(ps)) - startup = next( - node - for node in parsed.body - if isinstance(node, (ast.AsyncFunctionDef, ast.FunctionDef)) and node.name == "proxy_startup_event" + return _step + + +@pytest.mark.asyncio +async def test_shutdown_flush_runs_after_in_flight_passthrough_error_reports(): + """A passthrough error report callback writes its spend row through the + logging worker, so every spend flush must run only after in-flight reports + have finished: a flush that ran while the report was still pending could + skip its row. + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _spawn_report_task, + drain_passthrough_upstream_error_reports, ) - awaited = tuple( - child.value.func.id - for child in sorted( - ( - node - for node in ast.walk(startup) - if isinstance(node, ast.Await) - and isinstance(node.value, ast.Call) - and isinstance(node.value.func, ast.Name) - ), - key=lambda node: node.lineno, - ) + + report_done = asyncio.Event() + calls: List[str] = [] + + async def report() -> None: + await asyncio.sleep(0.05) + report_done.set() + + _spawn_report_task(report()) + await ps._drain_reports_then_flush_spend( + drain_reports=drain_passthrough_upstream_error_reports, + drain_spend_events=_recorded_step(calls, "drain_spend_events", report_done), + stop_scheduler_jobs=_recorded_step(calls, "stop_scheduler_jobs", report_done), + flush_spend_counters=_recorded_step(calls, "flush_spend_counters", report_done), + flush_spend_logs=_recorded_step(calls, "flush_spend_logs", report_done), ) - drain_at = awaited.index("drain_passthrough_upstream_error_reports") - for flush in ( - "_drain_spend_event_producer_on_shutdown", - "flush_spend_counters_on_shutdown", - "_flush_spend_logs_queue_on_shutdown", - ): - assert drain_at < awaited.index(flush), f"report drain must run before {flush}" + + assert calls == [ + "drain_spend_events:report_done=True", + "stop_scheduler_jobs:report_done=True", + "flush_spend_counters:report_done=True", + "flush_spend_logs:report_done=True", + ], calls + + +@pytest.mark.asyncio +async def test_shutdown_flush_continues_when_report_drain_fails(): + report_done = asyncio.Event() + calls: List[str] = [] + + async def failing_drain() -> None: + raise RuntimeError("drain gone") + + await ps._drain_reports_then_flush_spend( + drain_reports=failing_drain, + drain_spend_events=_recorded_step(calls, "drain_spend_events", report_done), + stop_scheduler_jobs=_recorded_step(calls, "stop_scheduler_jobs", report_done), + flush_spend_counters=_recorded_step(calls, "flush_spend_counters", report_done), + flush_spend_logs=_recorded_step(calls, "flush_spend_logs", report_done), + ) + + assert calls == [ + "drain_spend_events:report_done=False", + "stop_scheduler_jobs:report_done=False", + "flush_spend_counters:report_done=False", + "flush_spend_logs:report_done=False", + ], calls + + +@pytest.mark.asyncio +async def test_shutdown_flush_skips_scheduler_step_when_no_scheduler(): + report_done = asyncio.Event() + calls: List[str] = [] + + async def noop_drain() -> None: + report_done.set() + + await ps._drain_reports_then_flush_spend( + drain_reports=noop_drain, + drain_spend_events=_recorded_step(calls, "drain_spend_events", report_done), + stop_scheduler_jobs=None, + flush_spend_counters=_recorded_step(calls, "flush_spend_counters", report_done), + flush_spend_logs=_recorded_step(calls, "flush_spend_logs", report_done), + ) + + assert calls == [ + "drain_spend_events:report_done=True", + "flush_spend_counters:report_done=True", + "flush_spend_logs:report_done=True", + ], calls # ---------------------------------------------------------------------------