mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
e85cb4052e
* test: stabilize cleanup, lockout, and post-kill timing tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): bound the wedged Redis burst p95 by the socket timeout the test sets The burst p95 under a fixed 0.75s cap tracked how fast the box ran the 30 concurrent writes, so a loaded runner failed it at unchanged code. A handler that waits on the wedged publish takes at least the coordination Redis socket timeout, so the test now sets that timeout and asserts the p95 stays under it. * test: pin the lockout clock and bound the wedged Redis burst by the proxy's 2s wait cap The lockout tests now set the login throttle's clock instead of reading the real one, so Retry-After is pinned at 77 and a round-down fails on every run The wedged Redis burst p95 is now bounded by the 2s cap the proxy puts on its own Redis waits, so a burst wait capped under the 5s socket timeout still fails the test The fresh post-kill clients keep the harness's 2s keep-alive expiry, here and in the two tests #47080 changed * test: tell ceil from round in the lockout tests and bound the burst by the imported eviction cap The refused attempt now runs with 76.4s left on the 77s block, where ceil gives 77 but round and int give 76, so both lockout tests fail if Retry-After stops rounding up. The burst bound imports _BULK_EVICTION_WAIT_SECONDS instead of copying its 2s value, so lowering the cap moves the bound with it. * test: bound the wedged Redis burst by the shorter of the eviction cap and the socket timeout Any wait the proxy puts on the stopped Redis lasts at least the shorter of the two, so the bound now follows whichever one a later change lowers. * test: drive the tool index cleanup test with the manual clock and keep the post-kill block as on main The post-kill reset is already fixed on main by making kill() wait for the worker's last thread, so the fresh-client and keep-alive edits go back to main's version. The tool index cleanup test now advances the injected clock instead of timing the real one, and the lockout clock helper gets a type hint. * test: advance the spend log cleanup tests' manual clock through its own method * test: freeze the throttle clock in the login_utils lockout tests and drive the run budget tests with the manual clock --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2269 lines
89 KiB
Python
2269 lines
89 KiB
Python
"""
|
|
Test cases for spend log cleanup functionality
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
import math
|
|
import time
|
|
from collections.abc import Awaitable, Callable
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from litellm.constants import (
|
|
SPEND_LOG_CLEANUP_BATCH_SIZE,
|
|
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP,
|
|
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS,
|
|
)
|
|
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
|
|
SPEND_LOG_CLEANUP_BOUND_SETTINGS,
|
|
SizeCleanupState,
|
|
SpendLogCleanup,
|
|
SpendLogStorage,
|
|
StorageProbe,
|
|
TableCleanupResult,
|
|
advance_size_state,
|
|
)
|
|
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
|
|
SpendLogCleanupMetrics,
|
|
)
|
|
from tests.unit.proxy.db.fake_prisma_engine import engine_call
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _ManualClock:
|
|
now: float = 0.0
|
|
|
|
def __call__(self) -> float:
|
|
return self.now
|
|
|
|
def advance(self, seconds: float) -> None:
|
|
self.now += seconds
|
|
|
|
|
|
def _far_deadline() -> float:
|
|
"""A run deadline far enough out that only the other bounds can stop a batch loop."""
|
|
return time.monotonic() + 3600
|
|
|
|
|
|
def _wire_tx(db, rejected_timeout=None):
|
|
"""
|
|
Model the prisma seam the cleanup job actually uses.
|
|
|
|
Every statement the job issues runs inside db.tx() so it can carry a SET
|
|
LOCAL statement_timeout. Batch and probe statements are forwarded to
|
|
db.execute_raw and db.query_raw, which is what tests configure and assert
|
|
on, while the SET LOCAL statements are answered here so they neither consume
|
|
a side_effect entry nor show up in the recorded call list. Lookup is
|
|
deferred to call time so this can be wired before a test assigns its own
|
|
execute_raw.
|
|
"""
|
|
|
|
@asynccontextmanager
|
|
async def _tx(**kwargs):
|
|
if kwargs.get("timeout") == rejected_timeout:
|
|
raise TypeError("transaction timeout rejected by fake")
|
|
tx = MagicMock()
|
|
|
|
async def _execute_raw(sql, *args):
|
|
if sql.lstrip().upper().startswith("SET LOCAL"):
|
|
return 0
|
|
return await db.execute_raw(sql, *args)
|
|
|
|
async def _query_raw(sql, *args):
|
|
return await db.query_raw(sql, *args)
|
|
|
|
tx.execute_raw = _execute_raw
|
|
tx.query_raw = _query_raw
|
|
yield tx
|
|
|
|
db.tx = _tx
|
|
db.query_raw = AsyncMock(return_value=[{"remaining": 0}])
|
|
|
|
|
|
def test_spend_log_cleanup_cron_scheduling():
|
|
"""Test that cron expressions are correctly parsed for spend log cleanup scheduling"""
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
|
|
# Valid cron expressions
|
|
cron_expr = "0 4 * * *" # 4:00 AM daily
|
|
trigger = CronTrigger.from_crontab(cron_expr)
|
|
assert trigger is not None
|
|
|
|
# Every minute (useful for testing)
|
|
trigger_minute = CronTrigger.from_crontab("*/1 * * * *")
|
|
assert trigger_minute is not None
|
|
|
|
# Specific day and hour
|
|
trigger_weekly = CronTrigger.from_crontab("0 3 * * 0") # 3 AM every Sunday
|
|
assert trigger_weekly is not None
|
|
|
|
# Invalid cron expression should raise ValueError
|
|
with pytest.raises(ValueError, match="Wrong number of fields; got"):
|
|
CronTrigger.from_crontab("invalid cron")
|
|
|
|
with pytest.raises(ValueError, match="is higher than the maximum value"):
|
|
CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour
|
|
|
|
|
|
def test_spend_log_cleanup_cron_scheduler_integration():
|
|
"""
|
|
Integration test: Verify the proxy_server scheduler logic correctly adds
|
|
cron-based cleanup job when maximum_spend_logs_cleanup_cron is configured.
|
|
|
|
This tests the logic in proxy_server.py lines 4671-4717 without requiring
|
|
a real database connection.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
|
|
# Mock scheduler
|
|
mock_scheduler = MagicMock()
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_cleanup_instance = MagicMock()
|
|
|
|
# Test Case 1: Cron-based scheduling
|
|
general_settings_cron = {
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_cron": "0 4 * * *", # 4 AM daily
|
|
}
|
|
|
|
cleanup_cron = general_settings_cron.get("maximum_spend_logs_cleanup_cron")
|
|
assert cleanup_cron is not None
|
|
|
|
# Simulate the scheduler logic from proxy_server.py
|
|
cron_trigger = CronTrigger.from_crontab(cleanup_cron)
|
|
mock_scheduler.add_job(
|
|
mock_cleanup_instance.cleanup_old_spend_logs,
|
|
cron_trigger,
|
|
args=[mock_prisma_client],
|
|
id="spend_log_cleanup_job",
|
|
replace_existing=True,
|
|
misfire_grace_time=3600,
|
|
)
|
|
|
|
# Verify scheduler was called correctly
|
|
mock_scheduler.add_job.assert_called_once()
|
|
call_args = mock_scheduler.add_job.call_args
|
|
|
|
# Verify the trigger is a CronTrigger
|
|
assert isinstance(call_args[0][1], CronTrigger)
|
|
|
|
# Verify job ID
|
|
assert call_args[1]["id"] == "spend_log_cleanup_job"
|
|
assert call_args[1]["replace_existing"] is True
|
|
|
|
# Test Case 2: Interval-based scheduling (fallback)
|
|
mock_scheduler.reset_mock()
|
|
general_settings_interval = {
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
# No cron, so it should fall back to interval
|
|
}
|
|
|
|
cleanup_cron_fallback = general_settings_interval.get("maximum_spend_logs_cleanup_cron")
|
|
assert cleanup_cron_fallback is None # No cron configured
|
|
|
|
# Simulate interval-based scheduling fallback
|
|
retention_interval = general_settings_interval.get("maximum_spend_logs_retention_interval", "1d")
|
|
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
|
|
|
interval_seconds = duration_in_seconds(retention_interval)
|
|
|
|
mock_scheduler.add_job(
|
|
mock_cleanup_instance.cleanup_old_spend_logs,
|
|
"interval",
|
|
seconds=interval_seconds,
|
|
args=[mock_prisma_client],
|
|
id="spend_log_cleanup_job",
|
|
replace_existing=True,
|
|
)
|
|
|
|
# Verify interval scheduling was called
|
|
mock_scheduler.add_job.assert_called_once()
|
|
interval_call_args = mock_scheduler.add_job.call_args
|
|
assert interval_call_args[0][1] == "interval"
|
|
assert interval_call_args[1]["seconds"] == 86400 # 1 day in seconds
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_delete_spend_logs():
|
|
# Test case 1: No retention set
|
|
cleaner = SpendLogCleanup(general_settings={})
|
|
assert cleaner._should_delete_spend_logs() is False
|
|
|
|
# Test case 2: Valid seconds string
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "3600s"})
|
|
assert cleaner._should_delete_spend_logs() is True
|
|
|
|
# Test case 3: Valid days string
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
|
|
assert cleaner._should_delete_spend_logs() is True
|
|
|
|
# Test case 4: Valid hours string
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "24h"})
|
|
assert cleaner._should_delete_spend_logs() is True
|
|
|
|
# Test case 5: Invalid format
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "invalid"})
|
|
assert cleaner._should_delete_spend_logs() is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_batch_transaction_timeout_covers_statement_timeout() -> None:
|
|
transaction: Final = MagicMock()
|
|
transaction.execute_raw = AsyncMock(side_effect=[0, 0, 3])
|
|
|
|
@asynccontextmanager
|
|
async def tx(timeout: timedelta):
|
|
yield transaction
|
|
|
|
prisma_client: Final = MagicMock()
|
|
prisma_client.db.tx = MagicMock(side_effect=tx)
|
|
cleaner: Final = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
|
|
cleaner.batch_timeout_seconds = 12
|
|
|
|
await cleaner._execute_delete_batch(
|
|
prisma_client, "DELETE FROM table", datetime.now(), "LiteLLM_SpendLogs", _far_deadline()
|
|
)
|
|
|
|
assert prisma_client.db.tx.call_args.kwargs["timeout"] >= timedelta(seconds=12)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_count_remaining_transaction_timeout_covers_statement_timeout() -> None:
|
|
transaction: Final = MagicMock()
|
|
transaction.execute_raw = AsyncMock(return_value=0)
|
|
transaction.query_raw = AsyncMock(return_value=[{"remaining": 7}])
|
|
|
|
@asynccontextmanager
|
|
async def tx(timeout: timedelta):
|
|
yield transaction
|
|
|
|
prisma_client: Final = MagicMock()
|
|
prisma_client.db.tx = MagicMock(side_effect=tx)
|
|
cleaner: Final = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
|
|
cleaner.batch_timeout_seconds = 12
|
|
|
|
result: Final = await cleaner._count_remaining(
|
|
prisma_client, datetime.now(), "LiteLLM_SpendLogs", "startTime", _far_deadline()
|
|
)
|
|
|
|
assert result == 7
|
|
assert prisma_client.db.tx.call_args.kwargs["timeout"] >= timedelta(seconds=12)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_old_spend_logs_batch_deletion():
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
# Setup Prisma client
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
|
|
# Mock execute_raw to return deleted counts (3 spend-log batches, then the
|
|
# tool-index cleanup's first batch returning 0)
|
|
mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0, 0])
|
|
|
|
# Wire up mocks
|
|
mock_prisma_client.db = mock_db
|
|
|
|
# Mock Redis cache and pod_lock_manager
|
|
mock_redis_cache = MagicMock()
|
|
mock_pod_lock_manager = MagicMock()
|
|
mock_pod_lock_manager.redis_cache = mock_redis_cache
|
|
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
|
mock_pod_lock_manager.release_lock = AsyncMock()
|
|
|
|
# Run cleanup with mocked pod_lock_manager
|
|
test_settings = {"maximum_spend_logs_retention_period": "7d"}
|
|
cleaner = SpendLogCleanup(general_settings=test_settings)
|
|
cleaner.pod_lock_manager = mock_pod_lock_manager
|
|
assert cleaner._should_delete_spend_logs() is True
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
# Validate batching and deletion via raw SQL
|
|
assert mock_db.execute_raw.call_count == 4
|
|
|
|
# Check the first call argument
|
|
call_args_sql = mock_db.execute_raw.call_args_list[0][0][0]
|
|
assert 'DELETE FROM "LiteLLM_SpendLogs"' in call_args_sql
|
|
# must match on the full composite identity: on a partitioned table
|
|
# request_id alone is not unique, and deleting by it would let a client
|
|
# reusing x-litellm-call-id take out a fresh row alongside the expired one
|
|
assert 'WHERE ("request_id", "startTime") IN' in call_args_sql
|
|
|
|
# After spend logs, the derived tool index rows expire on the same cutoff
|
|
tool_index_sql = mock_db.execute_raw.call_args_list[3][0][0]
|
|
assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in tool_index_sql
|
|
|
|
# The LiteLLM_DailyToolSpend rollup must outlive spend-log retention: it is
|
|
# the only copy of tool spend history once its per-request sources expire,
|
|
# so spend-log cleanup must never touch it.
|
|
for call in mock_db.execute_raw.call_args_list:
|
|
assert "LiteLLM_DailyToolSpend" not in call[0][0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_old_spend_logs_retention_period_cutoff():
|
|
"""
|
|
Test that logs are filtered using correct cutoff based on retention
|
|
"""
|
|
# Setup Prisma client
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(return_value=0)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
# Mock Redis cache and pod_lock_manager
|
|
mock_redis_cache = MagicMock()
|
|
mock_pod_lock_manager = MagicMock()
|
|
mock_pod_lock_manager.redis_cache = mock_redis_cache
|
|
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
|
mock_pod_lock_manager.release_lock = AsyncMock()
|
|
|
|
# Run cleanup with mocked pod_lock_manager
|
|
test_settings = {"maximum_spend_logs_retention_period": "24h"}
|
|
cleaner = SpendLogCleanup(general_settings=test_settings)
|
|
cleaner.pod_lock_manager = mock_pod_lock_manager
|
|
assert cleaner._should_delete_spend_logs() is True
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
# Verify the cutoff date is correct
|
|
cutoff_date = mock_db.execute_raw.call_args[0][1]
|
|
expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400)
|
|
assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1 # Allow 1 second difference for test execution time
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_drops_partitions_when_enabled_and_partitioned():
|
|
"""
|
|
With use_spend_logs_partitioning enabled and a partitioned table, cleanup
|
|
must reclaim disk by dropping partitions AND still delete expired rows the
|
|
drops cannot reach (DEFAULT partition, cutoff-spanning partitions), so
|
|
retention is never bypassed.
|
|
"""
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_prisma_client.db.execute_raw = AsyncMock(return_value=0)
|
|
|
|
partition_manager = MagicMock()
|
|
partition_manager.is_partitioned = AsyncMock(return_value=True)
|
|
partition_manager.ensure_partitions = AsyncMock(return_value=["p1"])
|
|
partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"])
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"use_spend_logs_partitioning": True,
|
|
},
|
|
partition_manager=partition_manager,
|
|
)
|
|
cleaner.pod_lock_manager = MagicMock()
|
|
cleaner.pod_lock_manager.redis_cache = None
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
partition_manager.ensure_partitions.assert_awaited_once()
|
|
partition_manager.drop_partitions_older_than.assert_awaited_once()
|
|
delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0]
|
|
assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql
|
|
# Partition drops only reclaim spend logs; the tool index must still be
|
|
# cleaned row-wise on the same run
|
|
all_sql = [c[0][0] for c in mock_prisma_client.db.execute_raw.call_args_list]
|
|
assert any('DELETE FROM "LiteLLM_SpendLogToolIndex"' in s for s in all_sql)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_uses_delete_when_partitioning_not_enabled():
|
|
"""
|
|
Even against a partitioned table, the partition path must stay off until
|
|
use_spend_logs_partitioning is explicitly enabled, so existing deployments
|
|
see zero behavior change. The catalog must not even be queried.
|
|
"""
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0])
|
|
|
|
partition_manager = MagicMock()
|
|
partition_manager.is_partitioned = AsyncMock(return_value=True)
|
|
partition_manager.ensure_partitions = AsyncMock()
|
|
partition_manager.drop_partitions_older_than = AsyncMock()
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={"maximum_spend_logs_retention_period": "7d"},
|
|
partition_manager=partition_manager,
|
|
)
|
|
cleaner.pod_lock_manager = MagicMock()
|
|
cleaner.pod_lock_manager.redis_cache = None
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
partition_manager.is_partitioned.assert_not_awaited()
|
|
partition_manager.drop_partitions_older_than.assert_not_awaited()
|
|
delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0]
|
|
assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_uses_delete_when_not_partitioned():
|
|
"""
|
|
With the feature enabled but the table not actually partitioned (script not
|
|
run yet), cleanup must keep using the batched DELETE path.
|
|
"""
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0])
|
|
|
|
partition_manager = MagicMock()
|
|
partition_manager.is_partitioned = AsyncMock(return_value=False)
|
|
partition_manager.drop_partitions_older_than = AsyncMock()
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"use_spend_logs_partitioning": True,
|
|
},
|
|
partition_manager=partition_manager,
|
|
)
|
|
cleaner.pod_lock_manager = MagicMock()
|
|
cleaner.pod_lock_manager.redis_cache = None
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
partition_manager.drop_partitions_older_than.assert_not_awaited()
|
|
assert mock_prisma_client.db.execute_raw.await_count == 3
|
|
delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0]
|
|
assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_old_spend_logs_no_retention_period():
|
|
"""
|
|
Test that no logs are deleted when no retention period is set
|
|
"""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_prisma_client.db.execute_raw = AsyncMock()
|
|
|
|
cleaner = SpendLogCleanup(general_settings={}) # no retention
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
mock_prisma_client.db.execute_raw.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lock_not_released_when_not_acquired():
|
|
"""
|
|
Lock release should be skipped when _should_delete_spend_logs returns False
|
|
before the lock is ever acquired.
|
|
"""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_prisma_client.db.execute_raw = AsyncMock()
|
|
|
|
mock_redis_cache = MagicMock()
|
|
mock_pod_lock_manager = MagicMock()
|
|
mock_pod_lock_manager.redis_cache = mock_redis_cache
|
|
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
|
mock_pod_lock_manager.release_lock = AsyncMock()
|
|
|
|
# No retention setting → _should_delete_spend_logs() returns False before lock is acquired
|
|
cleaner = SpendLogCleanup(general_settings={})
|
|
cleaner.pod_lock_manager = mock_pod_lock_manager
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
mock_pod_lock_manager.acquire_lock.assert_not_called()
|
|
mock_pod_lock_manager.release_lock.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_integer_retention_treated_as_days():
|
|
"""
|
|
An integer value for maximum_spend_logs_retention_period should be treated
|
|
as days (e.g., 3 → '3d' → 259200 seconds).
|
|
"""
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": 3})
|
|
result = cleaner._should_delete_spend_logs()
|
|
assert result is True
|
|
assert cleaner.retention_seconds == 3 * 86400 # 3 days in seconds
|
|
|
|
|
|
def test_string_retention_still_works():
|
|
"""
|
|
String values like '3d', '24h', '3600s' should continue to parse correctly.
|
|
"""
|
|
cases = [
|
|
("3d", 3 * 86400),
|
|
("24h", 24 * 3600),
|
|
("3600s", 3600),
|
|
("2w", 2 * 604800),
|
|
]
|
|
for setting, expected_seconds in cases:
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": setting})
|
|
assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}"
|
|
assert cleaner.retention_seconds == expected_seconds, (
|
|
f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_old_logs_aborts_on_non_int_execute_raw_return():
|
|
"""should abort deletion loop immediately when execute_raw returns a non-int
|
|
(e.g. None or dict), preventing an infinite loop."""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(return_value=None)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
|
|
|
|
assert mock_db.execute_raw.call_count == 1
|
|
assert result.rows_deleted == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_old_logs_continues_on_valid_int_return():
|
|
"""should continue deletion loop across batches when execute_raw returns valid int counts."""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0])
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
|
|
|
|
assert mock_db.execute_raw.call_count == 3
|
|
assert result.rows_deleted == 800
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_old_rows_stops_at_max_batches():
|
|
"""The batch cap must halt a cleanup that keeps finding rows, so a huge
|
|
backlog is spread across scheduled runs instead of one unbounded loop, and
|
|
the operator-facing knob must mean exactly the number of statements it names."""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(return_value=1000)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_max_batches": 2,
|
|
}
|
|
)
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
|
|
|
|
assert mock_db.execute_raw.call_count == 2
|
|
assert result.rows_deleted == 2000
|
|
assert result.stop_reason == "batch_cap_reached"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_old_tool_index_rows_deletes_on_composite_key():
|
|
"""Tool index rows are derived from spend logs and expire on the same cutoff;
|
|
the delete must match on the table's composite primary key."""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(side_effect=[5, 0])
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline())
|
|
|
|
assert result.rows_deleted == 5
|
|
delete_sql = mock_db.execute_raw.call_args_list[0][0][0]
|
|
assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in delete_sql
|
|
assert 'WHERE ("request_id", "tool_name") IN' in delete_sql
|
|
assert '"start_time" <' in delete_sql
|
|
assert mock_db.execute_raw.call_args_list[0][0][1] == cutoff_date
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch):
|
|
"""A single batch failure (e.g. DB timeout) must not abort the whole run —
|
|
subsequent batches should still execute and their counts accumulate."""
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
# Zero out the failure backoff so the test doesn't take ~0.5s of real sleep.
|
|
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
# batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed,
|
|
# batch 5 returns 0 → loop exits naturally.
|
|
mock_db.execute_raw = AsyncMock(side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0])
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
|
|
|
|
# All 5 batches should have been attempted; 100 + 200 + 50 = 350 deleted.
|
|
assert mock_db.execute_raw.call_count == 5
|
|
assert result.rows_deleted == 350
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch):
|
|
"""If batch failures persist for SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES
|
|
in a row (e.g. DB is down), the loop must abort instead of hot-looping."""
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
# Lower the threshold so the test is fast and deterministic.
|
|
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)
|
|
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
# Every batch raises — must abort after exactly 3 attempts, not loop forever.
|
|
mock_db.execute_raw = AsyncMock(side_effect=ConnectionError("simulated persistent DB outage"))
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
|
|
|
|
assert mock_db.execute_raw.call_count == 3
|
|
assert result.rows_deleted == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatch):
|
|
"""A success between failures must reset the consecutive-failure counter so
|
|
intermittent timeouts don't trip the abort threshold."""
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)
|
|
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
# Pattern: fail, fail, success (resets counter), fail, fail, success, done.
|
|
# Without reset, three of these would trip abort; with reset, they don't.
|
|
mock_db.execute_raw = AsyncMock(
|
|
side_effect=[
|
|
TimeoutError("t1"),
|
|
TimeoutError("t2"),
|
|
100,
|
|
TimeoutError("t3"),
|
|
TimeoutError("t4"),
|
|
50,
|
|
0,
|
|
]
|
|
)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
|
|
|
|
assert mock_db.execute_raw.call_count == 7
|
|
assert result.rows_deleted == 150
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch):
|
|
"""The outer error handler must call logger.exception() (not .error(str(e)))
|
|
so Prisma/DB timeouts surface a full traceback and exception type."""
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
mock_logger = MagicMock()
|
|
monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
# Force the outer try/except to fire by making _should_delete_spend_logs raise.
|
|
cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
cleaner.pod_lock_manager = None
|
|
|
|
def boom():
|
|
raise RuntimeError("simulated prisma timeout")
|
|
|
|
cleaner._should_delete_spend_logs = boom # type: ignore[assignment]
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
assert mock_logger.exception.called, "expected logger.exception() to be called"
|
|
# The exception type name must appear in the formatted args so operators can
|
|
# tell *what* failed, not just "Error during cleanup:".
|
|
call_args = mock_logger.exception.call_args
|
|
formatted = call_args[0][0] % call_args[0][1:]
|
|
assert "RuntimeError" in formatted
|
|
assert "simulated prisma timeout" in formatted
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch):
|
|
"""Even when batch deletion aborts due to consecutive failures, the pod lock
|
|
must still be released so the next scheduled run isn't permanently blocked."""
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2)
|
|
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down"))
|
|
mock_prisma_client.db = mock_db
|
|
|
|
mock_pod_lock_manager = MagicMock()
|
|
mock_pod_lock_manager.redis_cache = MagicMock()
|
|
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
|
mock_pod_lock_manager.release_lock = AsyncMock()
|
|
|
|
cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
cleaner.pod_lock_manager = mock_pod_lock_manager
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
# Cleanup didn't crash; the abort-after-failures path returned cleanly.
|
|
mock_pod_lock_manager.release_lock.assert_awaited_once()
|
|
|
|
|
|
def test_cleanup_batch_size_env_var(monkeypatch):
|
|
"""Ensure batch size is configurable via environment variable"""
|
|
import importlib
|
|
|
|
import litellm.constants as constants_module
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
# Set env var and reload modules to pick up new value
|
|
monkeypatch.setenv("SPEND_LOG_CLEANUP_BATCH_SIZE", "25")
|
|
importlib.reload(constants_module)
|
|
importlib.reload(cleanup_module)
|
|
|
|
cleaner = cleanup_module.SpendLogCleanup(general_settings={})
|
|
assert cleaner.batch_size == 25
|
|
|
|
# Remove env var and reload to restore default for other tests
|
|
monkeypatch.delenv("SPEND_LOG_CLEANUP_BATCH_SIZE", raising=False)
|
|
importlib.reload(constants_module)
|
|
importlib.reload(cleanup_module)
|
|
|
|
|
|
def _mock_prisma_for_retention(side_effect: list) -> "MagicMock":
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
client = MagicMock()
|
|
_wire_tx(client.db, rejected_timeout=timedelta(seconds=10))
|
|
client.db.execute_raw = AsyncMock(side_effect=side_effect)
|
|
return client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_spend_logs_retention_alone_does_not_touch_the_session_rollup():
|
|
client = _mock_prisma_for_retention([0, 0])
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
cleaner.pod_lock_manager = None
|
|
await cleaner.cleanup_old_spend_logs(client)
|
|
tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
|
|
assert any('"LiteLLM_SpendLogs"' in sql for sql in tables)
|
|
assert not any('"LiteLLM_AutoRouterSession"' in sql for sql in tables)
|
|
assert not any('"LiteLLM_AutoRouterUserSession"' in sql for sql in tables)
|
|
assert not any('"LiteLLM_AutoRouterDailySpend"' in sql for sql in tables)
|
|
assert not any('"LiteLLM_HealthCheckTable"' in sql for sql in tables)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_session_retention_alone_cleans_both_session_rollups_and_the_daily_rollup():
|
|
client = _mock_prisma_for_retention([0, 0, 0])
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_autorouter_session_retention_period": "365d"})
|
|
cleaner.pod_lock_manager = None
|
|
await cleaner.cleanup_old_spend_logs(client)
|
|
calls = client.db.execute_raw.call_args_list
|
|
tables = [call[0][0] for call in calls]
|
|
assert len(tables) == 3
|
|
assert '"LiteLLM_AutoRouterSession"' in tables[0]
|
|
assert '"LiteLLM_AutoRouterUserSession"' in tables[1]
|
|
assert '"LiteLLM_AutoRouterDailySpend"' in tables[2]
|
|
assert calls[2][0][1] == calls[0][0][1].date().isoformat()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_check_retention_alone_cleans_only_the_health_check_table():
|
|
client = _mock_prisma_for_retention([0])
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_health_check_retention_period": "30d"})
|
|
cleaner.pod_lock_manager = None
|
|
await cleaner.cleanup_old_spend_logs(client)
|
|
tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
|
|
assert len(tables) == 1
|
|
assert '"LiteLLM_HealthCheckTable"' in tables[0]
|
|
assert '"health_check_id"' in tables[0]
|
|
assert '"checked_at"' in tables[0]
|
|
cutoff_date = client.db.execute_raw.call_args[0][1]
|
|
expected_cutoff = datetime.now(timezone.utc) - timedelta(days=30)
|
|
assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_daily_tag_spend_retention_alone_prunes_only_that_table_by_calendar_day():
|
|
client = _mock_prisma_for_retention([0])
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_daily_tag_spend_retention_period": "90d"})
|
|
cleaner.pod_lock_manager = None
|
|
await cleaner.cleanup_old_spend_logs(client)
|
|
tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
|
|
assert len(tables) == 1
|
|
assert '"LiteLLM_DailyTagSpend"' in tables[0]
|
|
cutoff_day = client.db.execute_raw.call_args[0][1]
|
|
assert cutoff_day == (datetime.now(timezone.utc) - timedelta(days=90)).date().isoformat()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_spend_logs_retention_alone_keeps_daily_tag_spend_forever():
|
|
client = _mock_prisma_for_retention([0, 0])
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
cleaner.pod_lock_manager = None
|
|
await cleaner.cleanup_old_spend_logs(client)
|
|
tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
|
|
assert not any('"LiteLLM_DailyTagSpend"' in sql for sql in tables)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_each_retention_key_cuts_off_at_its_own_horizon():
|
|
client = _mock_prisma_for_retention([0, 0, 0, 0, 0, 0])
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_autorouter_session_retention_period": "365d",
|
|
"maximum_health_check_retention_period": "30d",
|
|
}
|
|
)
|
|
cleaner.pod_lock_manager = None
|
|
await cleaner.cleanup_old_spend_logs(client)
|
|
cutoffs = {
|
|
(
|
|
"LiteLLM_AutoRouterSession"
|
|
if '"LiteLLM_AutoRouterSession"' in call[0][0]
|
|
else "LiteLLM_AutoRouterUserSession"
|
|
if '"LiteLLM_AutoRouterUserSession"' in call[0][0]
|
|
else "LiteLLM_AutoRouterDailySpend"
|
|
if '"LiteLLM_AutoRouterDailySpend"' in call[0][0]
|
|
else "LiteLLM_HealthCheckTable"
|
|
if '"LiteLLM_HealthCheckTable"' in call[0][0]
|
|
else "logs"
|
|
): call[0][1]
|
|
for call in client.db.execute_raw.call_args_list
|
|
}
|
|
now = datetime.now(timezone.utc)
|
|
assert (now - cutoffs["logs"]).days == 7
|
|
assert (now - cutoffs["LiteLLM_AutoRouterSession"]).days == 365
|
|
assert cutoffs["LiteLLM_AutoRouterUserSession"] == cutoffs["LiteLLM_AutoRouterSession"]
|
|
assert cutoffs["LiteLLM_AutoRouterDailySpend"] == cutoffs["LiteLLM_AutoRouterSession"].date().isoformat()
|
|
assert (now - cutoffs["LiteLLM_HealthCheckTable"]).days == 30
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_retention_keys_means_no_cleanup_at_all():
|
|
client = _mock_prisma_for_retention([])
|
|
cleaner = SpendLogCleanup(general_settings={})
|
|
cleaner.pod_lock_manager = None
|
|
await cleaner.cleanup_old_spend_logs(client)
|
|
assert client.db.execute_raw.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_budget_stops_the_loop_and_leaves_the_backlog_for_the_next_run():
|
|
"""
|
|
The wall-clock budget is the bound that keeps a large backlog from turning
|
|
into one multi-hour run. With rows always available, the loop must stop on
|
|
the deadline rather than on the batch cap, and must report that reason so
|
|
operators can tell a budgeted stop from a drained table.
|
|
"""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
clock: Final = _ManualClock()
|
|
|
|
async def _delete_batch(_query: str, *_args: object) -> int:
|
|
clock.advance(0.1)
|
|
return 1000
|
|
|
|
mock_db.execute_raw = AsyncMock(side_effect=_delete_batch)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
# Comfortably more batches than a sub-second budget can reach (each
|
|
# batch advances the clock 0.1s), but small enough that a broken
|
|
# deadline still ends the loop
|
|
"maximum_spend_logs_cleanup_max_batches": 50,
|
|
},
|
|
clock=clock,
|
|
)
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, clock() + 1.0)
|
|
|
|
assert result.stop_reason == "budget_exhausted"
|
|
assert clock.now < 1.15, f"budgeted run overran its deadline: {clock.now}s"
|
|
assert mock_db.execute_raw.call_count < 50
|
|
assert result.rows_deleted > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_budget_is_shared_across_tables_not_granted_per_table():
|
|
"""
|
|
A per-table budget would let a run take N times the configured bound. The
|
|
deadline is computed once per run, so once it is spent on the first table
|
|
the later tables must stop immediately rather than each getting a fresh one.
|
|
"""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
clock: Final = _ManualClock()
|
|
|
|
async def _delete_batch(_query: str, *_args: object) -> int:
|
|
clock.advance(0.1)
|
|
return 1000
|
|
|
|
mock_db.execute_raw = AsyncMock(side_effect=_delete_batch)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_autorouter_session_retention_period": "365d",
|
|
# Comfortably more batches than a sub-second budget can reach (each
|
|
# batch advances the clock 0.1s), but small enough that a broken
|
|
# deadline still ends the loop
|
|
"maximum_spend_logs_cleanup_max_batches": 50,
|
|
"maximum_spend_logs_cleanup_run_budget": "1s",
|
|
},
|
|
clock=clock,
|
|
)
|
|
cleaner.pod_lock_manager = None
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
# three tables are eligible; a per-table budget would push this past 3s
|
|
assert clock.now < 1.15, f"budget was granted per table, not per run: {clock.now}s"
|
|
tables_touched = {call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list}
|
|
assert "LiteLLM_SpendLogs" in tables_touched
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_groups_share_budget_so_health_checks_still_get_a_delete():
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
clock: Final = _ManualClock()
|
|
|
|
async def _delete_batch(_query: str, *_args: object) -> int:
|
|
clock.advance(0.1)
|
|
return 1000
|
|
|
|
mock_db.execute_raw = AsyncMock(side_effect=_delete_batch)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_health_check_retention_period": "30d",
|
|
"maximum_spend_logs_cleanup_max_batches": 50,
|
|
"maximum_spend_logs_cleanup_run_budget": "1s",
|
|
},
|
|
clock=clock,
|
|
)
|
|
cleaner.pod_lock_manager = None
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
tables_touched = {call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list}
|
|
assert "LiteLLM_SpendLogs" in tables_touched
|
|
assert "LiteLLM_HealthCheckTable" in tables_touched
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_each_batch_carries_a_statement_and_lock_timeout():
|
|
"""
|
|
A Prisma transaction timeout cannot interrupt a statement already running,
|
|
so the Postgres statement_timeout and lock_timeout are the only things
|
|
stopping one batch from holding row locks and a pooled connection
|
|
indefinitely. Both must be set, inside the batch's own transaction, and
|
|
scoped with SET LOCAL so the pooled connection is left unchanged.
|
|
"""
|
|
recorded: list[str] = []
|
|
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
|
|
@asynccontextmanager
|
|
async def _tx(**kwargs):
|
|
tx = MagicMock()
|
|
|
|
async def _execute_raw(sql, *args):
|
|
recorded.append(sql.strip())
|
|
return 0
|
|
|
|
tx.execute_raw = _execute_raw
|
|
yield tx
|
|
|
|
mock_db.tx = _tx
|
|
mock_db.query_raw = AsyncMock(return_value=[{"remaining": 0}])
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_batch_timeout": "12s",
|
|
}
|
|
)
|
|
|
|
await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline())
|
|
|
|
assert "SET LOCAL statement_timeout = 12000" in recorded
|
|
assert "SET LOCAL lock_timeout = 12000" in recorded
|
|
# the timeouts must precede the delete they are meant to bound
|
|
assert recorded.index("SET LOCAL statement_timeout = 12000") < next(
|
|
i for i, sql in enumerate(recorded) if sql.startswith("DELETE")
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"setting_value",
|
|
["inf", "-inf", "nan", "1e400", "0s", "-5m", "not-a-duration"],
|
|
)
|
|
def test_a_non_finite_or_non_positive_budget_falls_back_to_the_default(setting_value):
|
|
"""
|
|
The knob must not be able to remove the bound it exists to enforce.
|
|
|
|
'inf', 'nan' and '1e400' are the spellings that would turn the deadline
|
|
into no deadline at all, and '0s' and '-5m' would make every run stop before
|
|
deleting anything. All of them must land on the default rather than being
|
|
honoured, and the resulting budget must be usable arithmetic.
|
|
"""
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_run_budget": setting_value,
|
|
}
|
|
)
|
|
|
|
assert cleaner.run_budget_seconds == SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS
|
|
assert math.isfinite(cleaner.run_budget_seconds)
|
|
assert cleaner.run_budget_seconds > 0
|
|
|
|
|
|
@pytest.mark.parametrize("setting_value", [0, -1, "abc", "", 2.9])
|
|
def test_a_bad_batch_size_falls_back_to_the_default(setting_value):
|
|
"""A zero or negative batch size would make every DELETE a no-op and the
|
|
loop spin, so unusable values must fall back rather than be honoured."""
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_batch_size": setting_value,
|
|
}
|
|
)
|
|
|
|
assert cleaner.batch_size >= 1
|
|
|
|
|
|
def test_operator_knobs_override_the_env_defaults():
|
|
"""The knobs are meant to be reachable from general_settings (and therefore
|
|
from the admin UI), not only from environment variables."""
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_batch_size": 250,
|
|
"maximum_spend_logs_cleanup_max_batches": 7,
|
|
"maximum_spend_logs_cleanup_run_budget": "90s",
|
|
"maximum_spend_logs_cleanup_batch_timeout": "2m",
|
|
}
|
|
)
|
|
|
|
assert cleaner.batch_size == 250
|
|
assert cleaner.max_batches == 7
|
|
assert cleaner.run_budget_seconds == 90
|
|
assert cleaner.batch_timeout_seconds == 120
|
|
|
|
|
|
_BOUND_SETTING_CASES = (
|
|
("maximum_spend_logs_cleanup_batch_size", 137, "batch_size", 137),
|
|
("maximum_spend_logs_cleanup_max_batches", 9, "max_batches", 9),
|
|
("maximum_spend_logs_cleanup_run_budget", "45s", "run_budget_seconds", 45.0),
|
|
("maximum_spend_logs_cleanup_batch_timeout", "8s", "batch_timeout_seconds", 8.0),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("setting_name, setting_value, attribute, expected", _BOUND_SETTING_CASES)
|
|
@pytest.mark.asyncio
|
|
async def test_a_bound_changed_after_construction_reaches_the_next_run(
|
|
setting_name, setting_value, attribute, expected
|
|
):
|
|
"""The scheduler holds one long-lived instance and the config reload mutates
|
|
general_settings in place, so a bound captured at construction would leave
|
|
every dashboard change inert until the process restarts."""
|
|
settings = {"maximum_spend_logs_retention_period": "7d"}
|
|
cleaner = SpendLogCleanup(general_settings=settings)
|
|
cleaner.pod_lock_manager = None
|
|
assert getattr(cleaner, attribute) != expected
|
|
|
|
settings[setting_name] = setting_value
|
|
|
|
await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0]))
|
|
|
|
assert getattr(cleaner, attribute) == expected
|
|
|
|
|
|
@pytest.mark.parametrize("cleared_to_none", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_a_bound_cleared_after_construction_falls_back_to_its_default(cleared_to_none):
|
|
"""Blanking the field in the dashboard has to restore the shipped default
|
|
rather than leave the operator's old bound in force, whether the reload
|
|
spells the clear as an explicit None or as an absent key."""
|
|
settings = {"maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_batch_size": 137}
|
|
cleaner = SpendLogCleanup(general_settings=settings)
|
|
cleaner.pod_lock_manager = None
|
|
assert cleaner.batch_size == 137
|
|
|
|
if cleared_to_none:
|
|
settings["maximum_spend_logs_cleanup_batch_size"] = None
|
|
else:
|
|
del settings["maximum_spend_logs_cleanup_batch_size"]
|
|
|
|
await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0]))
|
|
|
|
assert cleaner.batch_size == SPEND_LOG_CLEANUP_BATCH_SIZE
|
|
|
|
|
|
def test_every_declared_bound_setting_is_covered_by_a_live_reread_case():
|
|
"""A bound added to the declared set without a live-reread case would be
|
|
propagated by the proxy and then ignored by the running job."""
|
|
assert {case[0] for case in _BOUND_SETTING_CASES} == set(SPEND_LOG_CLEANUP_BOUND_SETTINGS)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table():
|
|
"""The remaining-eligible-rows metric must never itself become the long
|
|
scan this job exists to avoid, so its probe carries a LIMIT."""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(return_value=0)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
|
|
await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline())
|
|
|
|
count_sql = mock_db.query_raw.call_args[0][0]
|
|
assert "count(*)" in count_sql
|
|
assert "LIMIT $2" in count_sql
|
|
assert mock_db.query_raw.call_args[0][2] == SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_run_skipped_because_another_pod_holds_the_lock_is_reported():
|
|
"""Operators need to tell "nothing to do" apart from "someone else is doing
|
|
it", so a lock-skipped run is recorded under its own outcome."""
|
|
recorded: list[str] = []
|
|
original_record_run = SpendLogCleanupMetrics.record_run
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
cleaner.pod_lock_manager = MagicMock()
|
|
cleaner.pod_lock_manager.redis_cache = MagicMock()
|
|
cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
|
|
cleaner.pod_lock_manager.release_lock = AsyncMock()
|
|
|
|
SpendLogCleanupMetrics.record_run = classmethod(lambda cls, outcome: recorded.append(outcome))
|
|
try:
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
finally:
|
|
SpendLogCleanupMetrics.record_run = original_record_run
|
|
|
|
assert recorded == ["skipped_locked"]
|
|
cleaner.pod_lock_manager.release_lock.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_the_outstanding_rows_probe_carries_a_statement_timeout():
|
|
"""
|
|
The probe is a statement like any other, so if it were issued bare a slow one
|
|
would hold a connection past the budget the job advertises, which is exactly
|
|
what the bounds exist to prevent. With budget to spare it carries the same
|
|
per-statement timeout the delete batches do.
|
|
"""
|
|
recorded: list[str] = []
|
|
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
|
|
@asynccontextmanager
|
|
async def _tx(**kwargs):
|
|
tx = MagicMock()
|
|
|
|
async def _execute_raw(sql, *args):
|
|
recorded.append(sql.strip())
|
|
return 0
|
|
|
|
async def _query_raw(sql, *args):
|
|
recorded.append(sql.strip())
|
|
return [{"remaining": 7}]
|
|
|
|
tx.execute_raw = _execute_raw
|
|
tx.query_raw = _query_raw
|
|
yield tx
|
|
|
|
mock_db.tx = _tx
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_batch_timeout": "8s",
|
|
}
|
|
)
|
|
|
|
remaining = await cleaner._count_remaining(
|
|
mock_prisma_client,
|
|
datetime.now(timezone.utc) - timedelta(days=7),
|
|
"LiteLLM_SpendLogs",
|
|
"startTime",
|
|
_far_deadline(),
|
|
)
|
|
|
|
assert remaining == 7
|
|
count_index = next(i for i, sql in enumerate(recorded) if sql.startswith("SELECT count(*)"))
|
|
assert "SET LOCAL statement_timeout = 8000" in recorded[:count_index], (
|
|
f"the probe ran without a statement timeout: {recorded}"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_statement_timeout_is_clamped_to_the_budget_that_is_left():
|
|
"""
|
|
Postgres has no 'stop at time T', only a per-statement duration, so a batch
|
|
issued just under the deadline would run a whole batch timeout past it and
|
|
the run budget would be advisory. Clamping the timeout to the remaining
|
|
budget is what makes the budget a real wall clock.
|
|
"""
|
|
recorded: list[str] = []
|
|
client = MagicMock()
|
|
|
|
@asynccontextmanager
|
|
async def _tx(**kwargs):
|
|
tx = MagicMock()
|
|
|
|
async def _execute_raw(sql, *args):
|
|
recorded.append(sql.strip())
|
|
return 0
|
|
|
|
tx.execute_raw = _execute_raw
|
|
tx.query_raw = AsyncMock(return_value=[{"remaining": 0}])
|
|
yield tx
|
|
|
|
client.db.tx = _tx
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_batch_timeout": "30s",
|
|
}
|
|
)
|
|
|
|
# Only 2s of budget left against a 30s batch timeout.
|
|
await cleaner._execute_delete_batch(
|
|
client, "DELETE FROM x", datetime.now(timezone.utc), "LiteLLM_SpendLogs", time.monotonic() + 2
|
|
)
|
|
|
|
timeouts = [sql for sql in recorded if "statement_timeout" in sql]
|
|
assert timeouts, f"no statement timeout was issued: {recorded}"
|
|
issued_ms = int(timeouts[0].split("=")[1].strip())
|
|
assert issued_ms <= 2000, f"the batch was given {issued_ms}ms with only 2000ms of budget left"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_statement_is_issued_once_the_budget_is_spent():
|
|
"""
|
|
Every table exits through _finish_table, including the ones a spent run never
|
|
started, so an unconditional probe there would put one more statement per
|
|
table past the bound.
|
|
"""
|
|
client = _mock_prisma_for_retention([0, 0])
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
|
|
result = await cleaner._finish_table(
|
|
client,
|
|
datetime.now(timezone.utc) - timedelta(days=7),
|
|
"LiteLLM_SpendLogs",
|
|
"startTime",
|
|
123,
|
|
"budget_exhausted",
|
|
time.monotonic() - 1,
|
|
)
|
|
|
|
assert result.rows_deleted == 123
|
|
assert result.stop_reason == "budget_exhausted"
|
|
client.db.query_raw.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_batch_cancelled_by_the_deadline_is_budget_exhaustion_not_a_failure(monkeypatch):
|
|
"""
|
|
Clamping the timeout means the last batch of a budget-exhausted run is
|
|
cancelled by the deadline itself. Counting that as a batch failure would
|
|
inflate the failure metric on every such run and walk it toward the abort
|
|
threshold, so it has to be classified as the bound working.
|
|
"""
|
|
failures: list[str] = []
|
|
client = MagicMock()
|
|
_wire_tx(client.db)
|
|
|
|
# The deadline has to pass DURING the batch, not before it: a deadline
|
|
# already spent is caught by the loop's own check and no batch is ever
|
|
# issued, which would exercise none of the classification under test.
|
|
async def _cancelled_after_the_deadline(sql, *args):
|
|
await asyncio.sleep(0.05)
|
|
raise Exception("canceling statement due to statement timeout")
|
|
|
|
client.db.execute_raw = _cancelled_after_the_deadline
|
|
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
monkeypatch.setattr(SpendLogCleanupMetrics, "record_batch_failure", lambda table: failures.append(table))
|
|
|
|
result = await cleaner._delete_old_logs(
|
|
client, datetime.now(timezone.utc) - timedelta(days=7), time.monotonic() + 0.02
|
|
)
|
|
|
|
assert result.stop_reason == "budget_exhausted"
|
|
assert failures == [], f"a deadline cancellation was recorded as a batch failure: {failures}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_partition_maintenance_is_skipped_once_the_run_budget_is_spent():
|
|
"""
|
|
Dropping a partition is DDL holding an ACCESS EXCLUSIVE lock, and unlike a
|
|
delete batch it cannot be cut short once it has started. A run whose budget is
|
|
already gone must therefore not start it at all; the next tick picks it up.
|
|
"""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(return_value=0)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
partition_manager = MagicMock()
|
|
partition_manager.is_partitioned = AsyncMock(return_value=True)
|
|
partition_manager.ensure_partitions = AsyncMock()
|
|
partition_manager.drop_partitions_older_than = AsyncMock(return_value=[])
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"use_spend_logs_partitioning": True,
|
|
},
|
|
partition_manager=partition_manager,
|
|
)
|
|
cleaner._should_delete_spend_logs()
|
|
|
|
# a deadline already in the past is what a run that spent its budget on an
|
|
# earlier table looks like
|
|
await cleaner._clean_spend_log_tables(mock_prisma_client, time.monotonic() - 1)
|
|
|
|
partition_manager.ensure_partitions.assert_not_awaited()
|
|
partition_manager.drop_partitions_older_than.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_partition_maintenance_still_runs_while_the_run_has_budget():
|
|
"""The skip above must be caused by the spent budget, not by breaking the
|
|
partition path outright."""
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(return_value=0)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
partition_manager = MagicMock()
|
|
partition_manager.is_partitioned = AsyncMock(return_value=True)
|
|
partition_manager.ensure_partitions = AsyncMock()
|
|
partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"])
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"use_spend_logs_partitioning": True,
|
|
},
|
|
partition_manager=partition_manager,
|
|
)
|
|
cleaner._should_delete_spend_logs()
|
|
|
|
await cleaner._clean_spend_log_tables(mock_prisma_client, _far_deadline())
|
|
|
|
partition_manager.ensure_partitions.assert_awaited_once()
|
|
partition_manager.drop_partitions_older_than.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"stop_reasons, expected",
|
|
[
|
|
(("exhausted",), "completed"),
|
|
(("exhausted", "exhausted"), "completed"),
|
|
(("exhausted", "batch_cap_reached"), "batch_cap_reached"),
|
|
(("batch_cap_reached", "exhausted"), "batch_cap_reached"),
|
|
(("exhausted", "budget_exhausted"), "budget_exhausted"),
|
|
(("budget_exhausted", "exhausted"), "budget_exhausted"),
|
|
(("batch_cap_reached", "budget_exhausted"), "budget_exhausted"),
|
|
(("budget_exhausted", "batch_cap_reached"), "budget_exhausted"),
|
|
(("exhausted", "aborted"), "aborted"),
|
|
(("aborted", "exhausted"), "aborted"),
|
|
(("budget_exhausted", "aborted"), "aborted"),
|
|
(("aborted", "budget_exhausted"), "aborted"),
|
|
(("aborted", "budget_exhausted", "batch_cap_reached"), "aborted"),
|
|
],
|
|
)
|
|
def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(stop_reasons, expected):
|
|
"""
|
|
The run outcome answers "why did this run stop", so a table that merely ran
|
|
dry must never mask one that hit a bound, and an abort must outrank both.
|
|
|
|
Both orders of every pair are covered because this folds several per-table
|
|
results into one answer: a first-match-wins implementation would pass on
|
|
whichever order happened to be written and fail on its mirror.
|
|
"""
|
|
results = tuple(
|
|
TableCleanupResult(table_name=f"t{i}", rows_deleted=0, stop_reason=reason)
|
|
for i, reason in enumerate(stop_reasons)
|
|
)
|
|
assert SpendLogCleanup._run_outcome(results) == expected
|
|
|
|
|
|
_OTHER_OUTCOMES: Final = ("completed", "budget_exhausted", "batch_cap_reached", "skipped_locked", "skipped_disabled")
|
|
|
|
|
|
def _runs_recorded(outcome: str) -> float:
|
|
"""The real ``litellm_spend_log_cleanup_runs_total`` sample for one outcome, 0 when unset"""
|
|
from prometheus_client import REGISTRY
|
|
|
|
return REGISTRY.get_sample_value("litellm_spend_log_cleanup_runs_total", {"outcome": outcome}) or 0.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_raising(monkeypatch):
|
|
"""A run cut short by shutdown must leave its outcome and how far it got behind"""
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
mock_logger = MagicMock()
|
|
monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
|
|
aborted_runs_before = _runs_recorded("aborted")
|
|
other_runs_before = {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES}
|
|
|
|
third_batch_reached = asyncio.Event()
|
|
|
|
async def _execute_raw(sql, *args):
|
|
if third_batch_reached.is_set():
|
|
raise AssertionError("no batch may be issued after the cancelled one")
|
|
if _execute_raw.calls < 2:
|
|
_execute_raw.calls += 1
|
|
return 150
|
|
third_batch_reached.set()
|
|
await asyncio.Event().wait()
|
|
|
|
_execute_raw.calls = 0
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_prisma_client.db.execute_raw = _execute_raw
|
|
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
cleaner.pod_lock_manager = MagicMock()
|
|
cleaner.pod_lock_manager.redis_cache = MagicMock()
|
|
cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
|
cleaner.pod_lock_manager.release_lock = AsyncMock()
|
|
|
|
run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(mock_prisma_client))
|
|
await asyncio.wait_for(third_batch_reached.wait(), timeout=5)
|
|
run.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await run
|
|
|
|
assert _runs_recorded("aborted") == aborted_runs_before + 1
|
|
assert {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} == other_runs_before
|
|
cleaner.pod_lock_manager.release_lock.assert_awaited_once()
|
|
mock_logger.exception.assert_not_called()
|
|
(error_call,) = mock_logger.error.call_args_list
|
|
rendered = error_call[0][0] % error_call[0][1:]
|
|
assert rendered.startswith("Spend log cleanup cancelled after ")
|
|
assert "s (rows_deleted=300, batches=2)" in rendered
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatch):
|
|
"""The scheduler holds one cleaner for the life of the process, so progress must not carry over"""
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
mock_logger = MagicMock()
|
|
monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
_wire_tx(mock_prisma_client.db)
|
|
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0])
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
cleaner.pod_lock_manager = None
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, asyncio.CancelledError()])
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
(error_call,) = mock_logger.error.call_args_list
|
|
rendered = error_call[0][0] % error_call[0][1:]
|
|
assert "(rows_deleted=150, batches=1)" in rendered
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_size_pass_that_fires_while_an_age_pass_is_running_in_the_same_pod_waits_for_it(monkeypatch):
|
|
"""Without Redis there is no pod lock, and the two jobs are scheduled on two SpendLogCleanup instances"""
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
mock_logger = MagicMock()
|
|
monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
|
|
|
|
first_batch_done = asyncio.Event()
|
|
second_run_done = asyncio.Event()
|
|
|
|
async def _slow_execute_raw(sql, *args):
|
|
first_batch_done.set()
|
|
await second_run_done.wait()
|
|
return 100
|
|
|
|
slow_client = MagicMock()
|
|
_wire_tx(slow_client.db)
|
|
slow_client.db.execute_raw = _slow_execute_raw
|
|
size_db = MagicMock()
|
|
_wire_tx(size_db)
|
|
size_db.execute_raw = AsyncMock(return_value=1000)
|
|
_wire_size_probes(size_db, [_storage(allocated=10 * 1024**3, live=1000, dead=0)], cutoff=datetime.now(timezone.utc))
|
|
size_client = MagicMock()
|
|
size_client.db = size_db
|
|
|
|
age_cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
age_cleaner.pod_lock_manager = None
|
|
size_cleaner = _size_cleaner({})
|
|
|
|
slow_run = asyncio.ensure_future(age_cleaner.cleanup_old_spend_logs(slow_client))
|
|
await asyncio.wait_for(first_batch_done.wait(), timeout=5)
|
|
size_run = asyncio.ensure_future(size_cleaner.cleanup_spend_logs_by_size(size_client))
|
|
await asyncio.sleep(0.05)
|
|
assert size_db.deletes.await_count == 0 and not size_run.done(), "the size pass ran while the age pass held the pod"
|
|
|
|
second_run_done.set()
|
|
await asyncio.sleep(0)
|
|
slow_run.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await slow_run
|
|
await asyncio.wait_for(size_run, timeout=5)
|
|
|
|
assert size_db.deletes.await_count > 0, "the size pass never ran once the age pass let go"
|
|
(error_call,) = mock_logger.error.call_args_list
|
|
rendered = error_call[0][0] % error_call[0][1:]
|
|
assert "(rows_deleted=100, batches=1)" in rendered
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup():
|
|
"""
|
|
Both spend-log tables share one run budget. Before the fix the spend-log
|
|
loop ran against the whole deadline, so a backlog that outlasted the budget
|
|
meant LiteLLM_SpendLogToolIndex never received a single delete batch, run
|
|
after run. The index table must still get its own share of the budget.
|
|
"""
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
clock: Final = _ManualClock()
|
|
|
|
async def _delete_batch(_query: str, *_args: object) -> int:
|
|
clock.advance(0.1)
|
|
return 1000
|
|
|
|
mock_db.execute_raw = AsyncMock(side_effect=_delete_batch)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_max_batches": 500,
|
|
"maximum_spend_logs_cleanup_run_budget": "1s",
|
|
},
|
|
clock=clock,
|
|
)
|
|
cleaner.pod_lock_manager = None
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
tables = [call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list]
|
|
assert tables.count("LiteLLM_SpendLogs") > 0
|
|
assert tables.count("LiteLLM_SpendLogToolIndex") > 0, "tool index cleanup was starved by the spend-log backlog"
|
|
assert clock.now < 1.15, f"splitting the budget must not extend the run: {clock.now}s"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_that_leaves_backlog_logs_a_warning_summary_naming_each_table(caplog):
|
|
"""
|
|
Operators running at warning or error level saw nothing when a run stopped
|
|
with expired rows still present. A run that ends on a bound must emit one
|
|
WARNING line that names every table, its rows deleted and its stop reason.
|
|
"""
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(return_value=1000)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_max_batches": 2,
|
|
}
|
|
)
|
|
cleaner.pod_lock_manager = None
|
|
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
summaries = [record for record in caplog.records if "Spend log cleanup run finished" in record.getMessage()]
|
|
assert len(summaries) == 1
|
|
summary = summaries[0]
|
|
assert summary.levelno == logging.WARNING
|
|
message = summary.getMessage()
|
|
assert "outcome=batch_cap_reached" in message
|
|
assert "LiteLLM_SpendLogs: deleted=2000 stop_reason=batch_cap_reached" in message
|
|
assert "LiteLLM_SpendLogToolIndex: deleted=2000 stop_reason=batch_cap_reached" in message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_that_drains_every_table_logs_the_summary_at_info_not_warning(caplog):
|
|
"""A healthy run must not page anyone: the summary stays at INFO."""
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
_wire_tx(mock_db)
|
|
mock_db.execute_raw = AsyncMock(return_value=0)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
|
|
cleaner.pod_lock_manager = None
|
|
|
|
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
summaries = [record for record in caplog.records if "Spend log cleanup run finished" in record.getMessage()]
|
|
assert len(summaries) == 1
|
|
assert summaries[0].levelno == logging.INFO
|
|
assert "outcome=completed" in summaries[0].getMessage()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_cleanup_delete_batch_renders_a_postgres_delete_span_for_its_table(
|
|
postgres_span_names: Callable[[], Awaitable[tuple[str, ...]]],
|
|
) -> None:
|
|
client = MagicMock()
|
|
_wire_tx(client.db)
|
|
client.db.execute_raw = engine_call(5)
|
|
|
|
await SpendLogCleanup(general_settings={})._execute_delete_batch(
|
|
client, "DELETE FROM x", datetime(2026, 1, 1, tzinfo=timezone.utc), "LiteLLM_SpendLogs", time.monotonic() + 2
|
|
)
|
|
|
|
assert await postgres_span_names() == ("postgres.delete LiteLLM_SpendLogs",)
|
|
|
|
|
|
# --- size rule (maximum_spend_logs_size) ---------------------------------------------------------------
|
|
|
|
|
|
def _storage(allocated: int, live: int, dead: int, inserted: int | None = None, deleted: int | None = None) -> dict:
|
|
return {
|
|
"allocated_bytes": allocated,
|
|
"live_rows": live,
|
|
"dead_rows": dead,
|
|
"inserted_total": live + dead if inserted is None else inserted,
|
|
"deleted_total": dead if deleted is None else deleted,
|
|
}
|
|
|
|
|
|
def _wire_size_probes(
|
|
db, storage_rows: list[dict], cutoff: datetime, saved_state: SizeCleanupState | None = None
|
|
) -> list[SizeCleanupState]:
|
|
"""
|
|
Answer the size pass's catalog reads: the storage probe walks `storage_rows` one measurement per call
|
|
(repeating the last), the step-cutoff query returns `cutoff`, the remaining-rows probe returns 0, and the
|
|
LiteLLM_Config row holds `saved_state`. Writes of that row land in the returned list instead of
|
|
db.execute_raw, so DELETE counts stay what tests configure.
|
|
"""
|
|
measurements = iter(storage_rows)
|
|
last = [storage_rows[-1]]
|
|
saves: list[SizeCleanupState] = []
|
|
deletes = db.execute_raw
|
|
|
|
async def _query_raw(sql, *args):
|
|
if "pg_total_relation_size" in sql:
|
|
row = next(measurements, None)
|
|
if row is None:
|
|
row = last[0]
|
|
return [row]
|
|
if "OFFSET" in sql:
|
|
return [{"startTime": cutoff}]
|
|
if '"LiteLLM_Config"' in sql:
|
|
return [{"state_json": saved_state.model_dump_json()}] if saved_state is not None else []
|
|
return [{"remaining": 0}]
|
|
|
|
async def _execute_raw(sql, *args):
|
|
if '"LiteLLM_Config"' in sql:
|
|
saves.append(SizeCleanupState.model_validate_json(args[1]))
|
|
return 1
|
|
return await deletes(sql, *args)
|
|
|
|
db.query_raw = AsyncMock(side_effect=_query_raw)
|
|
db.execute_raw = _execute_raw
|
|
db.deletes = deletes
|
|
return saves
|
|
|
|
|
|
def _size_cleaner(settings: dict, **kwargs) -> SpendLogCleanup:
|
|
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_size": "1GB", **settings}, **kwargs)
|
|
cleaner.pod_lock_manager = None
|
|
return cleaner
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("value", "expected"),
|
|
[
|
|
("180GB", 180 * 1024**3),
|
|
("20 GiB", 20 * 1024**3),
|
|
("512MB", 512 * 1024**2),
|
|
("1.5TB", int(1.5 * 1024**4)),
|
|
("4096", 4096),
|
|
(4096, 4096),
|
|
],
|
|
)
|
|
def test_size_in_bytes_reads_the_units_pg_size_pretty_prints(value, expected):
|
|
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import size_in_bytes
|
|
|
|
assert size_in_bytes(value) == expected
|
|
|
|
|
|
@pytest.mark.parametrize("value", ["", "big", "-1GB", "0", 0, "1PB"])
|
|
def test_size_in_bytes_rejects_what_cannot_be_a_size(value):
|
|
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import size_in_bytes
|
|
|
|
with pytest.raises(ValueError, match="size"):
|
|
size_in_bytes(value)
|
|
|
|
|
|
def test_an_unset_size_means_no_size_rule_and_a_bad_target_falls_back_to_the_size(caplog):
|
|
assert SpendLogCleanup(general_settings={})._size_rule() is None
|
|
rule = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "200MB"})._size_rule()
|
|
assert rule is not None
|
|
assert (rule.limit_bytes, rule.target_bytes) == (1024**3, 200 * 1024**2)
|
|
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
|
|
rule_without_room = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "2GB"})._size_rule()
|
|
assert rule_without_room is not None
|
|
assert rule_without_room.target_bytes == rule_without_room.limit_bytes == 1024**3
|
|
assert any("must be below maximum_spend_logs_size" in record.getMessage() for record in caplog.records)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_size_pass_under_the_limit_issues_no_delete():
|
|
db = MagicMock()
|
|
_wire_tx(db)
|
|
db.execute_raw = AsyncMock(return_value=0)
|
|
_wire_size_probes(db, [_storage(allocated=500 * 1024**2, live=1000, dead=0)], cutoff=datetime.now(timezone.utc))
|
|
client = MagicMock()
|
|
client.db = db
|
|
|
|
await _size_cleaner({}).cleanup_spend_logs_by_size(client)
|
|
|
|
assert db.deletes.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_size_pass_deletes_oldest_first_in_steps_until_the_live_estimate_is_under_the_target(caplog):
|
|
"""
|
|
Over the limit, the pass turns the next step of oldest rows into a cutoff and runs the retention DELETE
|
|
for the spend logs and their tool-index rows, re-measuring after each step; it stops on the first
|
|
measurement under the target, not when the file size (which deletes never shrink) moves.
|
|
"""
|
|
db = MagicMock()
|
|
_wire_tx(db)
|
|
# 3 GiB over 3000 rows is 1 MiB a row, so 2560 rows take it from 3 GiB to the 512 MiB target: one step
|
|
# of 3 batches; the tool index is cleaned on the same cutoff
|
|
db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 300, 0])
|
|
step_cutoff = datetime(2026, 9, 1, tzinfo=timezone.utc)
|
|
allocated = 3 * 1024**3
|
|
_wire_size_probes(db, [_storage(allocated, live=3000, dead=0)], cutoff=step_cutoff)
|
|
client = MagicMock()
|
|
client.db = db
|
|
cleaner = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "512MB"})
|
|
|
|
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
|
|
await cleaner.cleanup_spend_logs_by_size(client)
|
|
|
|
deletes = [(call[0][0], call[0][1:]) for call in db.deletes.call_args_list]
|
|
assert len(deletes) == 5
|
|
assert 'DELETE FROM "LiteLLM_SpendLogs"' in deletes[0][0]
|
|
assert 'WHERE "startTime" < $1' in deletes[0][0]
|
|
assert deletes[0][1][0] == step_cutoff
|
|
assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in deletes[3][0]
|
|
assert deletes[3][1][0] == step_cutoff
|
|
offset_query = next(call for call in db.query_raw.call_args_list if "OFFSET" in call[0][0])
|
|
assert offset_query[0][1] == 3000
|
|
summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage())
|
|
assert "outcome=completed" in summary
|
|
assert "LiteLLM_SpendLogs (size rule): deleted=3000 stop_reason=target_reached" in summary
|
|
assert "LiteLLM_SpendLogToolIndex (size rule): deleted=300 stop_reason=target_reached" in summary
|
|
assert f"allocated_before={allocated}" in summary
|
|
assert "estimated_live_before=3221225472" in summary
|
|
assert "estimated_live_after=0" in summary
|
|
assert "target_reached=true above_target=0" in summary
|
|
assert "VACUUM FULL or pg_repack" in summary
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_size_pass_stops_at_the_run_batch_cap_and_reports_how_far_above_the_target_it_is(caplog):
|
|
db = MagicMock()
|
|
_wire_tx(db)
|
|
# 4 GiB over 4000 rows: 3584 rows must go to reach the 512 MiB target, the cap allows 3 batches
|
|
db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 0])
|
|
allocated = 4 * 1024**3
|
|
saves = _wire_size_probes(db, [_storage(allocated, live=4000, dead=0)], cutoff=datetime.now(timezone.utc))
|
|
client = MagicMock()
|
|
client.db = db
|
|
cleaner = _size_cleaner(
|
|
{"maximum_spend_logs_cleanup_target_size": "512MB", "maximum_spend_logs_cleanup_max_batches": 3}
|
|
)
|
|
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
await cleaner.cleanup_spend_logs_by_size(client)
|
|
|
|
assert db.deletes.await_count == 4
|
|
summary = next(r for r in caplog.records if "Spend log cleanup run finished" in r.getMessage())
|
|
assert summary.levelno == logging.WARNING
|
|
assert "outcome=batch_cap_reached" in summary.getMessage()
|
|
assert "LiteLLM_SpendLogs (size rule): deleted=3000 stop_reason=batch_cap_reached" in summary.getMessage()
|
|
# 1000 of 4000 rows left is 1 GiB live, 512 MiB above the target
|
|
assert f"target_reached=false above_target={512 * 1024**2}" in summary.getMessage()
|
|
# the next check starts from what this run freed, with the deletes already counted out of n_tup_del
|
|
assert saves == [
|
|
SizeCleanupState(
|
|
allocated_bytes=allocated,
|
|
reclaimable_bytes=round(3000 * allocated / 4000),
|
|
bytes_per_row=allocated / 4000,
|
|
inserted_total=4000,
|
|
deleted_total=3000,
|
|
)
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_size_pass_analyzes_and_measures_again_when_the_row_statistics_were_reset(caplog):
|
|
"""A crash zeroes pg_stat_user_tables under a full file; the pass rebuilds the counters instead of skipping.
|
|
The first read after ANALYZE still shows zeros, as the stats collector publishes them a moment later"""
|
|
db = MagicMock()
|
|
_wire_tx(db)
|
|
statements = AsyncMock(return_value=0)
|
|
db.execute_raw = statements
|
|
saves = _wire_size_probes(
|
|
db,
|
|
[
|
|
_storage(allocated=10 * 1024**3, live=0, dead=0, inserted=0, deleted=0),
|
|
_storage(allocated=10 * 1024**3, live=0, dead=0, inserted=0, deleted=0),
|
|
_storage(10 * 1024**3, 1000, 0, 0, 0),
|
|
],
|
|
cutoff=datetime.now(timezone.utc),
|
|
saved_state=SizeCleanupState(
|
|
allocated_bytes=10 * 1024**3,
|
|
reclaimable_bytes=9 * 1024**3 + 512 * 1024**2,
|
|
bytes_per_row=1024.0,
|
|
inserted_total=50_000,
|
|
deleted_total=49_000,
|
|
),
|
|
)
|
|
client = MagicMock()
|
|
client.db = db
|
|
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
await _size_cleaner({}).cleanup_spend_logs_by_size(client)
|
|
|
|
issued = [call.args[0] for call in statements.await_args_list]
|
|
assert issued == ['ANALYZE "LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"'], (
|
|
"the only statement is the ANALYZE: no DELETE, the carried count puts the table under 1GB"
|
|
)
|
|
assert any("row statistics are stale" in record.getMessage() for record in caplog.records)
|
|
assert saves[0].reclaimable_bytes == 9 * 1024**3 + 512 * 1024**2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_size_pass_does_not_refire_after_autovacuum_forgets_the_rows_it_deleted():
|
|
"""
|
|
The run before this one deleted 15,000 of 20,000 rows and saved what it freed. Autovacuum has since
|
|
turned those dead tuples into free space, so pg_stat_user_tables reports 5,000 live and 0 dead rows while
|
|
the files still measure the full 32 MB. Trusting the live share would call all 32 MB live and empty the
|
|
table; the saved count says 8 MB are live, under the limit, so nothing is deleted.
|
|
"""
|
|
db = MagicMock()
|
|
_wire_tx(db)
|
|
db.execute_raw = AsyncMock(return_value=1000)
|
|
after_last_run = SizeCleanupState(
|
|
allocated_bytes=32_030_720,
|
|
reclaimable_bytes=24_015_000,
|
|
bytes_per_row=1601.0,
|
|
inserted_total=20_000,
|
|
deleted_total=15_000,
|
|
)
|
|
saves = _wire_size_probes(
|
|
db,
|
|
[_storage(32_030_720, live=5000, dead=0, inserted=20_000, deleted=15_000)],
|
|
cutoff=datetime.now(timezone.utc),
|
|
saved_state=after_last_run,
|
|
)
|
|
client = MagicMock()
|
|
client.db = db
|
|
cleaner = _size_cleaner({"maximum_spend_logs_size": "16MB", "maximum_spend_logs_cleanup_target_size": "10MB"})
|
|
|
|
await cleaner.cleanup_spend_logs_by_size(client)
|
|
|
|
assert db.deletes.await_count == 0
|
|
assert saves[0].reclaimable_bytes == 24_015_000
|
|
|
|
|
|
_RUN_STATE = SizeCleanupState(
|
|
allocated_bytes=32_030_720,
|
|
reclaimable_bytes=24_015_000,
|
|
bytes_per_row=1601.0,
|
|
inserted_total=20_000,
|
|
deleted_total=15_000,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("name", "probe", "reclaimable_expected"),
|
|
[
|
|
(
|
|
"vacuum zeroed the dead rows, nothing else changed",
|
|
_storage(32_030_720, 5000, 0, 20_000, 15_000),
|
|
24_015_000,
|
|
),
|
|
("7000 new rows written into the free space", _storage(32_030_720, 12_000, 0, 27_000, 15_000), 12_808_000),
|
|
(
|
|
"7000 new rows that all grew the files",
|
|
_storage(32_030_720 + 7000 * 1601, 12_000, 0, 27_000, 15_000),
|
|
24_015_000,
|
|
),
|
|
("VACUUM FULL handed the free space back", _storage(8_015_720, 5000, 0, 20_000, 15_000), 0),
|
|
(
|
|
"statistics wiped while the files still hold data: the count is carried forward, not reseeded",
|
|
_storage(32_030_720, 0, 0, 0, 0),
|
|
24_015_000,
|
|
),
|
|
("the retention job deleted 2000 more rows", _storage(32_030_720, 3000, 2000, 20_000, 17_000), 27_217_000),
|
|
(
|
|
"statistics reset (a crash, then ANALYZE): the count is carried forward, not restarted from the dead rows",
|
|
_storage(32_030_720, 5000, 1000, 5000, 1000),
|
|
24_015_000,
|
|
),
|
|
(
|
|
"statistics reset on a file VACUUM FULL shrank meanwhile: the carried count loses the shrink",
|
|
_storage(8_015_720, 5000, 0, 5000, 0),
|
|
0,
|
|
),
|
|
(
|
|
"statistics reset with more dead rows than the carried count explains: the dead rows win",
|
|
_storage(32_030_720, 1000, 19_000, 1000, 19_000),
|
|
30_429_184,
|
|
),
|
|
],
|
|
)
|
|
def test_advance_size_state_tracks_free_space_through_vacuum_inserts_and_other_deletes(
|
|
name, probe, reclaimable_expected
|
|
):
|
|
advanced = advance_size_state(_RUN_STATE, StorageProbe(**probe))
|
|
|
|
assert advanced.reclaimable_bytes == reclaimable_expected, name
|
|
assert 0 <= probe["allocated_bytes"] - advanced.reclaimable_bytes <= probe["allocated_bytes"], name
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_size_only_run_leaves_the_age_rules_to_the_retention_job():
|
|
db = MagicMock()
|
|
_wire_tx(db)
|
|
db.execute_raw = AsyncMock(return_value=0)
|
|
_wire_size_probes(db, [_storage(allocated=10 * 1024**3, live=100, dead=0)], cutoff=datetime.now(timezone.utc))
|
|
client = MagicMock()
|
|
client.db = db
|
|
cleaner = _size_cleaner(
|
|
{"maximum_spend_logs_retention_period": "7d", "maximum_health_check_retention_period": "30d"}
|
|
)
|
|
|
|
await cleaner.cleanup_spend_logs_by_size(client)
|
|
|
|
tables = {call[0][0].split('"')[1] for call in db.deletes.call_args_list}
|
|
assert tables == {"LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(("cron", "size_pass_expected"), [(None, False), ("0 3 * * *", True)])
|
|
async def test_the_retention_run_carries_the_size_pass_only_when_one_cron_drives_both(cron, size_pass_expected):
|
|
db = MagicMock()
|
|
_wire_tx(db)
|
|
db.execute_raw = AsyncMock(return_value=0)
|
|
_wire_size_probes(db, [_storage(allocated=10 * 1024**3, live=100, dead=0)], cutoff=datetime.now(timezone.utc))
|
|
client = MagicMock()
|
|
client.db = db
|
|
cleaner = _size_cleaner({"maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_cron": cron})
|
|
|
|
await cleaner.cleanup_old_spend_logs(client)
|
|
|
|
measured = any("pg_total_relation_size" in call[0][0] for call in db.query_raw.call_args_list)
|
|
assert measured is size_pass_expected
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_partitioned_size_pass_drops_the_oldest_partitions_first_and_skips_row_deletes_once_under_target(
|
|
caplog,
|
|
):
|
|
db = MagicMock()
|
|
_wire_tx(db)
|
|
db.execute_raw = AsyncMock(return_value=1000)
|
|
saves = _wire_size_probes(
|
|
db,
|
|
[_storage(10 * 1024**3, live=1000, dead=0), _storage(256 * 1024**2, live=25, dead=0)],
|
|
cutoff=datetime.now(timezone.utc),
|
|
)
|
|
client = MagicMock()
|
|
client.db = db
|
|
verdicts = []
|
|
|
|
async def _drop(prisma_client, now, remaining_timeout_ms, still_over_target):
|
|
verdicts.append(await still_over_target())
|
|
return ("LiteLLM_SpendLogs_p20260601",)
|
|
|
|
partition_manager = MagicMock()
|
|
partition_manager.is_partitioned = AsyncMock(return_value=True)
|
|
partition_manager.ensure_partitions = AsyncMock()
|
|
partition_manager.drop_oldest_partitions_while = AsyncMock(side_effect=_drop)
|
|
cleaner = _size_cleaner(
|
|
{"use_spend_logs_partitioning": True, "maximum_spend_logs_cleanup_target_size": "512MB"},
|
|
partition_manager=partition_manager,
|
|
)
|
|
|
|
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
|
|
await cleaner.cleanup_spend_logs_by_size(client)
|
|
|
|
partition_manager.ensure_partitions.assert_awaited_once(), "a size-only setup still needs tomorrow's partition"
|
|
partition_manager.drop_oldest_partitions_while.assert_awaited_once()
|
|
assert verdicts == [False]
|
|
assert db.deletes.await_count == 0
|
|
summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage())
|
|
assert "partitions_dropped=1" in summary
|
|
assert "target_reached=true" in summary
|
|
assert (saves[-1].inserted_total, saves[-1].deleted_total, saves[-1].allocated_bytes) == (25, 0, 256 * 1024**2), (
|
|
"the saved counters must be the post-drop ones, or the next check reads the drop as a statistics reset"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_failed_remeasure_after_a_partition_drop_stops_the_pass_instead_of_deleting_on_old_numbers(caplog):
|
|
db = MagicMock()
|
|
_wire_tx(db)
|
|
db.execute_raw = AsyncMock(return_value=1000)
|
|
saves = _wire_size_probes(db, [_storage(10 * 1024**3, live=1000, dead=0)], cutoff=datetime.now(timezone.utc))
|
|
client = MagicMock()
|
|
client.db = db
|
|
|
|
async def _drop(prisma_client, now, remaining_timeout_ms, still_over_target):
|
|
db.query_raw = AsyncMock(side_effect=RuntimeError("connection reset"))
|
|
assert await still_over_target() is False
|
|
return ("LiteLLM_SpendLogs_p20260601",)
|
|
|
|
partition_manager = MagicMock()
|
|
partition_manager.is_partitioned = AsyncMock(return_value=True)
|
|
partition_manager.ensure_partitions = AsyncMock()
|
|
partition_manager.drop_oldest_partitions_while = AsyncMock(side_effect=_drop)
|
|
cleaner = _size_cleaner(
|
|
{"use_spend_logs_partitioning": True, "maximum_spend_logs_cleanup_target_size": "512MB"},
|
|
partition_manager=partition_manager,
|
|
)
|
|
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
await cleaner.cleanup_spend_logs_by_size(client)
|
|
|
|
assert db.deletes.await_count == 0, "the row step ran on the pre-drop measurement"
|
|
assert saves == [], "a count saved from pre-drop counters reads the drop as a statistics reset next check"
|
|
assert any("could not be re-measured after dropping partitions" in r.getMessage() for r in caplog.records)
|
|
|
|
|
|
def test_drop_oldest_partitions_while_only_drops_fully_elapsed_partitions_oldest_first():
|
|
from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import SpendLogsPartitionManager
|
|
|
|
manager = SpendLogsPartitionManager()
|
|
now = datetime(2026, 10, 17, 12, tzinfo=timezone.utc)
|
|
manager._list_partitions = AsyncMock(
|
|
return_value=[
|
|
("LiteLLM_SpendLogs_p20261017", datetime(2026, 10, 18)), # today's, still receiving writes
|
|
("LiteLLM_SpendLogs_p20261015", datetime(2026, 10, 16)),
|
|
("LiteLLM_SpendLogs_default", None),
|
|
("LiteLLM_SpendLogs_p20261014", datetime(2026, 10, 15)),
|
|
]
|
|
)
|
|
dropped_sql: list[str] = []
|
|
|
|
async def _ddl(prisma_client, statement, timeout_ms):
|
|
dropped_sql.append(statement)
|
|
|
|
manager._execute_bounded_ddl = _ddl
|
|
verdicts = iter([True, False])
|
|
|
|
async def still_over_target():
|
|
return next(verdicts)
|
|
|
|
dropped = asyncio.run(manager.drop_oldest_partitions_while(MagicMock(), now, lambda: 1000, still_over_target))
|
|
|
|
assert dropped == ("LiteLLM_SpendLogs_p20261014", "LiteLLM_SpendLogs_p20261015")
|
|
assert dropped_sql == [
|
|
'DROP TABLE IF EXISTS "LiteLLM_SpendLogs_p20261014"',
|
|
'DROP TABLE IF EXISTS "LiteLLM_SpendLogs_p20261015"',
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("settings", "retention_job", "size_job"),
|
|
[
|
|
({}, False, False),
|
|
({"maximum_spend_logs_retention_period": "7d"}, True, False),
|
|
({"maximum_spend_logs_size": "180GB"}, False, True),
|
|
({"maximum_spend_logs_size": "180GB", "maximum_spend_logs_retention_period": "7d"}, True, True),
|
|
({"maximum_spend_logs_size": "180GB", "maximum_spend_logs_cleanup_cron": "0 3 * * *"}, True, False),
|
|
],
|
|
)
|
|
def test_which_scheduler_jobs_the_cleanup_settings_ask_for(settings, retention_job, size_job):
|
|
from litellm.proxy.proxy_server import _wants_retention_job, _wants_size_check_job
|
|
|
|
assert _wants_retention_job(settings) is retention_job
|
|
assert _wants_size_check_job(settings) is size_job
|
|
|
|
|
|
def test_after_deleting_counts_the_rows_and_any_file_growth_as_reclaimable_and_to_state_keeps_the_counters():
|
|
storage = SpendLogStorage(allocated_bytes=32_030_720, reclaimable_bytes=0, bytes_per_row=1601.0, live_rows=20_000)
|
|
|
|
after = storage.after_deleting(15_000, 32_030_720 + 8192)
|
|
|
|
assert (after.reclaimable_bytes, after.live_rows, after.estimated_live_bytes) == (24_023_192, 5000, 8_015_720)
|
|
assert after.to_state(inserted_total=20_000, deleted_total=15_000) == SizeCleanupState(
|
|
allocated_bytes=32_038_912,
|
|
reclaimable_bytes=24_023_192,
|
|
bytes_per_row=1601.0,
|
|
inserted_total=20_000,
|
|
deleted_total=15_000,
|
|
)
|
|
assert StorageProbe(
|
|
allocated_bytes=8192, live_rows=0, dead_rows=0, inserted_total=0, deleted_total=0
|
|
).stats_are_stale
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_storage_reads_the_catalog_row_for_both_tables_and_a_failed_probe_is_none(caplog):
|
|
prisma_client = MagicMock()
|
|
_wire_tx(prisma_client.db)
|
|
prisma_client.db.query_raw = AsyncMock(return_value=[_storage(32_030_720, 20_000, 0)])
|
|
cleaner = SpendLogCleanup(general_settings={})
|
|
|
|
probe = await cleaner._probe_storage(prisma_client, _far_deadline())
|
|
|
|
assert probe is not None
|
|
assert (probe.allocated_bytes, probe.live_rows, probe.dead_rows, probe.inserted_total, probe.deleted_total) == (
|
|
32_030_720,
|
|
20_000,
|
|
0,
|
|
20_000,
|
|
0,
|
|
)
|
|
assert prisma_client.db.query_raw.call_args.args[1] == ["LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"]
|
|
|
|
prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("pg_stat_user_tables unavailable"))
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
assert await cleaner._probe_storage(prisma_client, _far_deadline()) is None
|
|
assert any("Could not measure spend-log storage" in record.getMessage() for record in caplog.records)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_size_state_round_trips_through_its_litellm_config_row_and_a_lost_save_is_only_a_warning(caplog):
|
|
saved = {}
|
|
prisma_client = MagicMock()
|
|
_wire_tx(prisma_client.db)
|
|
|
|
async def execute_raw(sql, *args):
|
|
saved[args[0]] = args[1]
|
|
return 1
|
|
|
|
prisma_client.db.execute_raw = execute_raw
|
|
prisma_client.db.query_raw = AsyncMock(
|
|
side_effect=lambda sql, param: [{"state_json": saved[param]}] if param in saved else []
|
|
)
|
|
cleaner = SpendLogCleanup(general_settings={})
|
|
|
|
assert await cleaner._load_size_state(prisma_client, _far_deadline()) is None
|
|
await cleaner._save_size_state(prisma_client, _far_deadline(), _RUN_STATE)
|
|
assert list(saved) == ["spend_logs_size_cleanup_state"]
|
|
loaded = await cleaner._load_size_state(prisma_client, _far_deadline())
|
|
assert loaded is not None
|
|
assert loaded.model_dump() == _RUN_STATE.model_dump()
|
|
|
|
prisma_client.db.execute_raw = AsyncMock(side_effect=RuntimeError("LiteLLM_Config is read only"))
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
await cleaner._save_size_state(prisma_client, _far_deadline(), _RUN_STATE)
|
|
assert any("Could not save the spend-log size cleanup state" in record.getMessage() for record in caplog.records)
|