scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(proxy): requeue tool-usage and autorouter transactions on transient db error

Browse source 
When flush_tool_usage_transactions or flush_autorouter_turn_transactions
fails due to a transient database transport error (such as connection
refusal, connection resets, or pooler blips), requeue the popped batch at
the head of its respective queue under the lock instead of permanently
dropping the transactions. Non-transient data errors are still logged and
dropped to prevent poison loop retries.

Closes #43002

Signed-off-by: Liang Xu <755674130@qq.com>
This commit is contained in:
Liang Xu 2026-09-27 09:47:35 +08:00
parent 7aba77197d
commit f4f8974743
2 changed files with 188 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

66
litellm/proxy/utils.py
View file

@ -7543,6 +7543,10 @@ async def update_spend_logs_job(
await _run_spend_logs_job(prisma_client, db_writer_client, proxy_logging_obj)
def _is_transient_spend_log_write_error(e: Exception) -> bool:
return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
async def _run_spend_logs_job(
prisma_client: PrismaClient,
db_writer_client: AsyncHTTPHandler | None,
@ -7588,8 +7592,8 @@ async def _run_spend_logs_job(
)
# Tool usage tracking: drain the request-time queue into the tool index and the
# LiteLLM_DailyToolSpend rollup. Never retried; a dropped batch is permanently
# absent from the rollup, so failures log at error.
# LiteLLM_DailyToolSpend rollup. On transient DB connectivity errors the batch is
# requeued at the head of the queue; non-transient data rejections are logged and dropped.
async with prisma_client._tool_usage_transactions_lock:
tool_usage_to_process: Final = prisma_client.tool_usage_transactions[:MAX_LOGS_PER_INTERVAL]
prisma_client.tool_usage_transactions = prisma_client.tool_usage_transactions[len(tool_usage_to_process) :]
@ -7600,12 +7604,29 @@ async def _run_spend_logs_job(
prisma_client=prisma_client,
transactions=tool_usage_to_process,
)
except Exception as tool_tracking_err:
verbose_proxy_logger.error(
"Spend tracking - tool usage flush failed; %s tool usage transactions dropped: %s",
except asyncio.CancelledError:
async with prisma_client._tool_usage_transactions_lock:
prisma_client.tool_usage_transactions = tool_usage_to_process + prisma_client.tool_usage_transactions
verbose_proxy_logger.warning(
"Spend tracking - tool usage flush cancelled, requeued %d rows for the next flush",
len(tool_usage_to_process),
tool_tracking_err,
)
raise
except Exception as tool_tracking_err: # noqa: BLE001 # drain failure must not abort spend job
if _is_transient_spend_log_write_error(tool_tracking_err):
async with prisma_client._tool_usage_transactions_lock:
prisma_client.tool_usage_transactions = tool_usage_to_process + prisma_client.tool_usage_transactions
verbose_proxy_logger.warning(
"Spend tracking - tool usage flush hit transient DB error (%s); requeued %d rows for the next flush",
tool_tracking_err,
len(tool_usage_to_process),
)
else:
verbose_proxy_logger.error(
"Spend tracking - tool usage flush failed; %s tool usage transactions dropped: %s",
len(tool_usage_to_process),
tool_tracking_err,
)
await flush_baseline_accounting(prisma_client)
@ -7622,12 +7643,33 @@ async def _run_spend_logs_job(
prisma_client=prisma_client,
transactions=autorouter_turns_to_process,
)
except Exception as autorouter_tracking_err: # noqa: BLE001 # a drain bug must not abort the spend job
verbose_proxy_logger.error(
"Spend tracking - auto-router session rollup drain failed; %s turn transactions dropped: %s",
except asyncio.CancelledError:
async with prisma_client._autorouter_turn_transactions_lock:
prisma_client.autorouter_turn_transactions = (
autorouter_turns_to_process + prisma_client.autorouter_turn_transactions
)
verbose_proxy_logger.warning(
"Spend tracking - auto-router turn drain cancelled, requeued %d rows for the next flush",
len(autorouter_turns_to_process),
autorouter_tracking_err,
)
raise
except Exception as autorouter_tracking_err: # noqa: BLE001 # a drain bug must not abort the spend job
if _is_transient_spend_log_write_error(autorouter_tracking_err):
async with prisma_client._autorouter_turn_transactions_lock:
prisma_client.autorouter_turn_transactions = (
autorouter_turns_to_process + prisma_client.autorouter_turn_transactions
)
verbose_proxy_logger.warning(
"Spend tracking - auto-router session rollup drain hit transient DB error (%s); requeued %d rows for the next flush",
autorouter_tracking_err,
len(autorouter_turns_to_process),
)
else:
verbose_proxy_logger.error(
"Spend tracking - auto-router session rollup drain failed; %s turn transactions dropped: %s",
len(autorouter_turns_to_process),
autorouter_tracking_err,
)
try:
from litellm.proxy.db.shadow_eval_funnel import flush_shadow_eval_funnel
@ -7768,10 +7810,6 @@ async def _monitor_spend_logs_queue(
MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256
def _is_transient_spend_log_write_error(e: Exception) -> bool:
return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
async def _create_spend_logs_with_poison_isolation(
repo: SpendLogsRepository,
rows: Sequence[Mapping[str, object]],

136
tests/unit/proxy/test_update_spend.py
View file

@ -321,3 +321,139 @@ async def test_update_spend_logs_multiple_batches_with_failure():
# Verify all logs were cleared from transactions
assert len(prisma_client.spend_log_transactions) == 0
@pytest.mark.asyncio
async def test_tool_usage_transactions_requeued_on_transient_db_error():
"""
Test that when flush_tool_usage_transactions fails due to a transient database
transport error, the batch is requeued at the head of the queue instead of being permanently dropped.
"""
from litellm.proxy.utils import update_spend_logs_job
prisma_client = MockPrismaClient()
proxy_logging_obj = create_mock_proxy_logging()
# Pre-populate tool usage transactions
initial_transactions = [
{"id": "tool_1", "tool_name": "calculator"},
{"id": "tool_2", "tool_name": "web_search"},
]
prisma_client.tool_usage_transactions = list(initial_transactions)
with patch(
"litellm.proxy.db.spend_log_tool_index.flush_tool_usage_transactions",
new=AsyncMock(side_effect=httpx.ConnectError("Can't reach database server")),
):
await update_spend_logs_job(prisma_client, None, proxy_logging_obj)
# Tool usage transactions should be requeued at the head of the queue
assert len(prisma_client.tool_usage_transactions) == 2
assert prisma_client.tool_usage_transactions == initial_transactions
@pytest.mark.asyncio
async def test_tool_usage_transactions_dropped_on_permanent_error():
"""
Test that when flush_tool_usage_transactions fails due to a non-transient error,
the batch is dropped so it does not loop forever.
"""
from litellm.proxy.utils import update_spend_logs_job
prisma_client = MockPrismaClient()
proxy_logging_obj = create_mock_proxy_logging()
prisma_client.tool_usage_transactions = [
{"id": "tool_1", "tool_name": "poison_row"},
]
with patch(
"litellm.proxy.db.spend_log_tool_index.flush_tool_usage_transactions",
new=AsyncMock(side_effect=ValueError("Invalid data payload")),
):
await update_spend_logs_job(prisma_client, None, proxy_logging_obj)
# Permanent error drops the batch to avoid poison loops
assert len(prisma_client.tool_usage_transactions) == 0
@pytest.mark.asyncio
async def test_autorouter_turn_transactions_requeued_on_transient_db_error():
"""
Test that when flush_autorouter_turn_transactions fails due to a transient database
transport error, the batch is requeued at the head of the queue instead of being permanently dropped.
"""
from litellm.proxy.utils import update_spend_logs_job
prisma_client = MockPrismaClient()
proxy_logging_obj = create_mock_proxy_logging()
initial_turns = [
{"id": "turn_1", "session_id": "sess_123"},
{"id": "turn_2", "session_id": "sess_456"},
]
prisma_client.autorouter_turn_transactions = list(initial_turns)
with patch(
"litellm.proxy.db.autorouter_session_rollup.flush_autorouter_turn_transactions",
new=AsyncMock(side_effect=httpx.ConnectError("Connection refused")),
):
await update_spend_logs_job(prisma_client, None, proxy_logging_obj)
# Autorouter turn transactions should be requeued at the head of the queue
assert len(prisma_client.autorouter_turn_transactions) == 2
assert prisma_client.autorouter_turn_transactions == initial_turns
@pytest.mark.asyncio
async def test_autorouter_turn_transactions_dropped_on_permanent_error():
"""
Test that when flush_autorouter_turn_transactions fails due to a non-transient error,
the batch is dropped so it does not loop forever.
"""
from litellm.proxy.utils import update_spend_logs_job
prisma_client = MockPrismaClient()
proxy_logging_obj = create_mock_proxy_logging()
prisma_client.autorouter_turn_transactions = [
{"id": "turn_1", "session_id": "bad_turn"},
]
with patch(
"litellm.proxy.db.autorouter_session_rollup.flush_autorouter_turn_transactions",
new=AsyncMock(side_effect=ValueError("Invalid session data")),
):
await update_spend_logs_job(prisma_client, None, proxy_logging_obj)
# Permanent error drops the batch
assert len(prisma_client.autorouter_turn_transactions) == 0
@pytest.mark.asyncio
async def test_tool_usage_and_autorouter_turn_requeued_on_cancelled_error():
"""
Test that when flush is cancelled via asyncio.CancelledError, batches are requeued
and CancelledError is re-raised.
"""
from litellm.proxy.utils import update_spend_logs_job
prisma_client = MockPrismaClient()
proxy_logging_obj = create_mock_proxy_logging()
prisma_client.tool_usage_transactions = [
{"id": "tool_cancel", "tool_name": "calc"},
]
with patch(
"litellm.proxy.db.spend_log_tool_index.flush_tool_usage_transactions",
new=AsyncMock(side_effect=asyncio.CancelledError()),
):
with pytest.raises(asyncio.CancelledError):
await update_spend_logs_job(prisma_client, None, proxy_logging_obj)
# Must be requeued when cancelled
assert len(prisma_client.tool_usage_transactions) == 1
assert prisma_client.tool_usage_transactions[0]["id"] == "tool_cancel"