diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py index 32bf24d4ba5..be18d8ac6b1 100644 --- a/litellm/accounting/worker.py +++ b/litellm/accounting/worker.py @@ -87,10 +87,15 @@ async def flush_accounting_deliveries( for _ in range(batch_size): 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")), - batch_size=batch_size * 2, - ) + retention_days: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")) + cleanup_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if ( + await repository.cleanup_applied(retention_days=retention_days, batch_size=batch_size * 2) + < batch_size * 2 + or time.monotonic() >= cleanup_deadline + ): + break metrics: Final = await repository.metrics() record_deliveries(metrics, time.monotonic() - started) if metrics.batch_count >= int( diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index 68ab76d694d..60f3c074cca 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -103,6 +103,30 @@ async def test_worker_drains_a_backlog_in_one_run_one_configured_batch_per_trans await client.db.disconnect() +@pytest.mark.asyncio +async def test_worker_removes_more_expired_receipts_than_one_cleanup_batch_in_one_run( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "100") + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, state, applied_at) ' + "SELECT 'expired-' || n, '{}'::jsonb, false, 'applied', CURRENT_TIMESTAMP - INTERVAL '2 days' " + "FROM generate_series(1, 501) AS n" + ) + await client.db.connect() + try: + await flush_accounting_deliveries(DBSpendUpdateWriter(), client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) + finally: + await client.db.disconnect() + + @pytest.mark.parametrize("ledger", [False, True]) @pytest.mark.parametrize("exclude_messages", [False, True]) @pytest.mark.asyncio