diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 914c961b145..eaa03c5d7f7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1063,7 +1063,7 @@ class DBSpendUpdateWriter: await enqueue_spend_logs(prisma_client, (payload,)) if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: - request_spend_log_flush() + request_spend_log_flush(prisma_client) else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2c54b9a222d..44dbbc6cbdc 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3811,7 +3811,7 @@ class _StaleReadEngine: class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() - spend_log_flush_requested: ClassVar[asyncio.Event] = asyncio.Event() + spend_log_flush_requested: "asyncio.Event | None" = None spend_log_queue_bytes: ClassVar[int] = 0 spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None tool_usage_transactions: list["ToolUsageTransaction"] = [] @@ -6537,23 +6537,27 @@ async def enqueue_spend_logs( ) -def request_spend_log_flush() -> None: - """Wake the queue monitor now rather than leaving the rows for its next poll. +def request_spend_log_flush(prisma_client: PrismaClient) -> None: + """Wake this client's queue monitor now rather than leaving the rows for its next poll. The Responses API hands the client an id it can chain from straight away, and that lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. Repeated requests coalesce into the monitor's next pass, so the batching holds. + A request made before the monitor is running is dropped, and loses nothing: the + monitor reads the queue on its first pass, before it ever waits on a request. """ - PrismaClient.spend_log_flush_requested.set() + flush_requested: Final = prisma_client.spend_log_flush_requested + if flush_requested is not None: + flush_requested.set() -async def _wait_for_spend_log_flush_request(interval: float) -> bool: +async def _wait_for_spend_log_flush_request(flush_requested: asyncio.Event, interval: float) -> bool: """Wait out ``interval``, returning early and True when a flush was requested.""" try: - await asyncio.wait_for(PrismaClient.spend_log_flush_requested.wait(), timeout=interval) + await asyncio.wait_for(flush_requested.wait(), timeout=interval) except asyncio.TimeoutError: return False - PrismaClient.spend_log_flush_requested.clear() + flush_requested.clear() return True @@ -6980,6 +6984,8 @@ async def _monitor_spend_logs_queue( max_backoff: Final = 30.0 # Maximum backoff interval in seconds backoff_multiplier: Final = 1.5 # Exponential backoff multiplier current_interval = base_interval + flush_requested: Final = asyncio.Event() + prisma_client.spend_log_flush_requested = flush_requested # rebind-ok: the client owns its monitor's flush signal verbose_proxy_logger.info( "Starting spend logs queue monitor (threshold: %s, poll_interval: %ss)", threshold, base_interval @@ -7018,7 +7024,7 @@ async def _monitor_spend_logs_queue( # Exponential backoff when no logs to process current_interval = min(current_interval * backoff_multiplier, max_backoff) - if await _wait_for_spend_log_flush_request(current_interval): + if await _wait_for_spend_log_flush_request(flush_requested, current_interval): current_interval = base_interval except Exception as e: spend_log_error("Error in spend logs queue monitor: %s", str(e), exc=e) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 0bca7c9492c..5e977712a1e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2944,11 +2944,9 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_worker A `previous_response_id` chained straight off the previous turn reads the DB, so a Responses row cannot sit in this worker's queue until the monitor's next poll. """ - from litellm.proxy.utils import PrismaClient - db_writer = DBSpendUpdateWriter() prisma = _tool_usage_prisma() - PrismaClient.spend_log_flush_requested.clear() + prisma.spend_log_flush_requested = asyncio.Event() await db_writer._insert_spend_log_to_db( payload={"request_id": "req-1", "call_type": call_type}, @@ -2956,8 +2954,7 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_worker ) assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}] - assert PrismaClient.spend_log_flush_requested.is_set() is expects_flush - PrismaClient.spend_log_flush_requested.clear() + assert prisma.spend_log_flush_requested.is_set() is expects_flush def _batch_cost_payload() -> dict: diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index a1eb88a7834..c8b87bd671e 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -538,10 +538,9 @@ async def test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested( """ import litellm.constants as constants_mod import litellm.proxy.utils as utils_mod - from litellm.proxy.utils import PrismaClient, request_spend_log_flush + from litellm.proxy.utils import request_spend_log_flush monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) - PrismaClient.spend_log_flush_requested.clear() mock_prisma_client.spend_log_transactions = [] mock_prisma_client.tool_usage_transactions = [] @@ -562,16 +561,107 @@ async def test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested( try: await asyncio.sleep(0.05) assert not flushed.is_set() + assert isinstance(mock_prisma_client.spend_log_flush_requested, asyncio.Event) mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r1")) - request_spend_log_flush() + request_spend_log_flush(mock_prisma_client) await asyncio.wait_for(flushed.wait(), timeout=5.0) finally: monitor.cancel() with suppress(asyncio.CancelledError): await monitor - PrismaClient.spend_log_flush_requested.clear() + + +def test_monitor_spend_logs_queue_flush_survives_an_earlier_event_loop( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A second monitor, started in a fresh event loop, is still woken by a flush request, + so a worker whose first loop is gone keeps flushing Responses rows instead of stalling. + """ + import litellm.constants as constants_mod + import litellm.proxy.utils as utils_mod + from litellm.proxy.utils import request_spend_log_flush + + monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) + mock_prisma_client.tool_usage_transactions = [] + + async def _flush_once_under_a_monitor() -> None: + flushed: Final = asyncio.Event() + + async def _fake_job(*args: Any, **kwargs: Any) -> None: + flushed.set() + + monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job) + mock_prisma_client.spend_log_transactions = [] + + monitor: Final = asyncio.create_task( + _monitor_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + ) + try: + await asyncio.sleep(0.05) + assert not flushed.is_set() + + mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r1")) + request_spend_log_flush(mock_prisma_client) + + await asyncio.wait_for(flushed.wait(), timeout=5.0) + finally: + monitor.cancel() + with suppress(asyncio.CancelledError): + await monitor + + asyncio.run(_flush_once_under_a_monitor()) + asyncio.run(_flush_once_under_a_monitor()) + + +@pytest.mark.asyncio +async def test_flush_requested_before_the_monitor_starts_costs_the_row_nothing( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A Responses row enqueued before the monitor exists still reaches the DB on its first + pass, so dropping that early request delays nothing. + """ + import litellm.constants as constants_mod + import litellm.proxy.utils as utils_mod + from litellm.proxy.utils import request_spend_log_flush + + monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) + mock_prisma_client.spend_log_flush_requested = None + mock_prisma_client.tool_usage_transactions = [] + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")] + + flushed: Final = asyncio.Event() + + async def _fake_job(*args: Any, **kwargs: Any) -> None: + flushed.set() + + monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job) + + request_spend_log_flush(mock_prisma_client) + assert mock_prisma_client.spend_log_flush_requested is None + + monitor: Final = asyncio.create_task( + _monitor_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + ) + try: + await asyncio.wait_for(flushed.wait(), timeout=5.0) + finally: + monitor.cancel() + with suppress(asyncio.CancelledError): + await monitor def test_raise_failed_update_spend_exception_emits_failure_handler() -> None: