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()