fix(proxy): give each spend-log queue monitor its own flush event

`PrismaClient.spend_log_flush_requested` was an `asyncio.Event` built at
import time, so it bound to whichever event loop first awaited it and every
later loop got `RuntimeError: ... is bound to a different event loop` out of
`_wait_for_spend_log_flush_request`. The queue monitor's blanket `except
Exception` swallowed that into its error logger, so the flush silently never
happened and the row sat in the worker's queue until the next poll.

The monitor now creates its own Event inside the loop that awaits it and
hands it to the client, and `request_spend_log_flush` signals through the
client instead of the class. A request that arrives before the monitor is
running is dropped and loses nothing, because the monitor reads the queue on
its first pass before it ever waits.

In CI this showed up as the proxy-endpoints shard flaking on
test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested whenever
--dist=loadscope put the health-endpoint tests, which boot a proxy TestClient
and start a monitor, on the same worker ahead of the spend-log tests.
This commit is contained in:
mateo-berri 2026-09-03 04:02:55 -07:00
parent 658f50663d
commit cedf35992b
4 changed files with 79 additions and 18 deletions

View file

@ -940,7 +940,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.")

View file

@ -3519,7 +3519,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"] = []
@ -6245,23 +6245,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
@ -6681,6 +6685,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
@ -6719,7 +6725,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)

View file

@ -2941,11 +2941,9 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(c
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},
@ -2953,8 +2951,7 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(c
)
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
@pytest.mark.asyncio

View file

@ -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,75 @@ 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())
def test_request_spend_log_flush_is_a_no_op_before_the_monitor_starts(mock_prisma_client: Any) -> None:
"""A Responses row enqueued before the monitor's first pass must not fail the request."""
from litellm.proxy.utils import request_spend_log_flush
mock_prisma_client.spend_log_flush_requested = None
request_spend_log_flush(mock_prisma_client)
assert mock_prisma_client.spend_log_flush_requested is None
def test_raise_failed_update_spend_exception_emits_failure_handler() -> None: