From fe77d11061bb31e646759e100f94887214854ab3 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 00:30:47 -0800 Subject: [PATCH] fix(accounting): drain a backlog in one worker run, one batch per transaction One run applied a single fresh batch, so at 500 per batch charges fell behind any load above that rate. The worker now keeps claiming fresh batches while they come back full, for up to five seconds, and each batch still commits on its own. Co-Authored-By: Claude Opus 5.5 --- litellm/accounting/worker.py | 20 +++++++++++++++---- .../database/test_spendlogs_ledger_parity.py | 18 ++++++++--------- 2 files changed, 25 insertions(+), 13 deletions(-) diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py index 0607c1647e1..32bf24d4ba5 100644 --- a/litellm/accounting/worker.py +++ b/litellm/accounting/worker.py @@ -14,6 +14,7 @@ from litellm.accounting.metrics import record_deliveries from litellm.accounting.projection import Contribution from litellm.repositories.accounting_delivery_repository import ( AccountingDeliveryRepository, + DeliveryBatchResult, DeliveryProjection, LedgerDatabase, LedgerTransaction, @@ -34,6 +35,14 @@ class _AccountingAlerts(Protocol): _JOB: Final = "accounting_delivery" +def _filled(result: DeliveryBatchResult, batch_size: int) -> bool: + return result.applied + len(result.parked) >= batch_size + + +def _progressed(result: DeliveryBatchResult) -> bool: + return bool(result.applied or result.failed) + + async def flush_accounting_deliveries( writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging" ) -> None: @@ -56,7 +65,7 @@ async def flush_accounting_deliveries( started: Final = time.monotonic() batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500")) - async def run_batch(*, retry_only: bool) -> bool: + async def run_batch(*, retry_only: bool) -> DeliveryBatchResult: result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only) if result.applied: await writer.invalidate_accounting_entities(applied, proxy_logging_obj) @@ -68,12 +77,15 @@ async def flush_accounting_deliveries( proxy_logging_obj, message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.", ) - return bool(result.applied or result.failed) + return result - await run_batch(retry_only=False) + fresh_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline: + break retry_deadline: Final = time.monotonic() + 5 for _ in range(batch_size): - if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True): + if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)): break await repository.cleanup_applied( retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")), diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index 78439e4317e..7ea0bc181fb 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -63,10 +63,10 @@ _DAILY_TABLES: Final = ( ) -@pytest.mark.parametrize("batch_size,row_count,expected,distinct_models", [(None, 251, 251, False), (None, 501, 500, False), ("200", 251, 200, False), ("10000", 10001, 10000, False), ("3001", 3001, 3001, True)]) +@pytest.mark.parametrize("batch_size,row_count,transactions,distinct_models", [(None, 251, 1, False), (None, 1001, 3, False), ("200", 251, 2, False), ("3001", 3001, 1, True)]) @pytest.mark.asyncio -async def test_worker_applies_the_configured_backlog_in_one_transaction( - parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, expected: int, distinct_models: bool +async def test_worker_drains_a_backlog_in_one_run_one_configured_batch_per_transaction( + parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, transactions: int, distinct_models: bool ): from litellm.accounting.worker import flush_accounting_deliveries @@ -92,13 +92,13 @@ async def test_worker_applies_the_configured_backlog_in_one_transaction( )) == row_count - 1 started = time.monotonic() await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) - print(f"Ledger batch applied={expected} queued={row_count} seconds={time.monotonic() - started:.3f}") + print(f"Ledger backlog queued={row_count} seconds={time.monotonic() - started:.3f}") with psycopg.connect(parity_database) as connection: - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (expected,) - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (row_count - expected,) - assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (1,) - assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (expected * 0.25,) - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (expected if distinct_models else 1,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (row_count,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (0,) + assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (transactions,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (row_count * 0.25,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (row_count if distinct_models else 1,) finally: await client.db.disconnect()