diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 3ea8b874ef7..0cda6fb46a5 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -852,9 +852,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) - ################## Auto-Router Benchmarks Rollup ################## - await self.auto_router_session_queue.flush(prisma_client=prisma_client) - async def _commit_spend_updates_to_db_with_redis( self, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 621dea00e90..b7729de42f7 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -15,6 +15,9 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import ( SpendLogsPartitionManager, ) +from litellm.proxy.spend_tracking.auto_router_benchmarks import ( + AUTO_ROUTER_SESSION_RETENTION_DAYS, +) from litellm.proxy.utils import PrismaClient @@ -187,7 +190,7 @@ class SpendLogCleanup: ) async def _delete_old_auto_router_sessions(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int: - """Expire auto-router rollups on the spend-log cutoff. + """Expire auto-router rollups past the given cutoff. Keyed on last activity rather than session start, so a conversation still running when the cutoff passes is not pruned out from under itself. @@ -200,22 +203,28 @@ class SpendLogCleanup: time_column="last_turn_at", ) + def _auto_router_session_cutoff(self, retention_seconds: float | None) -> datetime: + """Rows the benchmarks endpoint can no longer read are collected at the read + horizon even with no retention configured; a shorter configured retention wins.""" + now: Final = datetime.now(timezone.utc) + horizon: Final = now - timedelta(days=AUTO_ROUTER_SESSION_RETENTION_DAYS) + return horizon if retention_seconds is None else max(horizon, now - timedelta(seconds=retention_seconds)) + async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None: """ Main cleanup function. Deletes old spend logs in batches. If pod_lock_manager is available, ensures only one pod runs cleanup. If no pod_lock_manager, runs cleanup without distributed locking. + + Spend logs and their tool index only expire when + ``maximum_spend_logs_retention_period`` is configured. The auto-router rollup is + collected on every run regardless; see ``_auto_router_session_cutoff``. """ lock_acquired = False try: verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now()) - if not self._should_delete_spend_logs(): - return - - if self.retention_seconds is None: - verbose_proxy_logger.error("Retention seconds is None, cannot proceed with cleanup") - return + retention_seconds: Final = self.retention_seconds if self._should_delete_spend_logs() else None # If we have a pod lock manager, try to acquire the lock if self.pod_lock_manager and self.pod_lock_manager.redis_cache: @@ -233,33 +242,38 @@ class SpendLogCleanup: verbose_proxy_logger.info("Another pod is already running cleanup") return - cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds)) - verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat()) + if retention_seconds is not None: + cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds)) + verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat()) - if self.general_settings.get( - "use_spend_logs_partitioning", False - ) and await self.partition_manager.is_partitioned(prisma_client): - await self.partition_manager.ensure_partitions(prisma_client) - dropped: Final = await self.partition_manager.drop_partitions_older_than(prisma_client, cutoff_date) - verbose_proxy_logger.info( - "Dropped %d expired spend-log partitions: %s", - len(dropped), - dropped, - ) - # DROP only reclaims whole expired partitions. Expired rows can - # still sit in the DEFAULT partition (backfill, coverage gaps) - # or in a partition that spans the cutoff, so retention must - # also delete those stragglers row-wise. - total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) - verbose_proxy_logger.info("Deleted %s expired logs not covered by dropped partitions", total_deleted) - else: - total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) - verbose_proxy_logger.info("Deleted %s logs", total_deleted) + if self.general_settings.get( + "use_spend_logs_partitioning", False + ) and await self.partition_manager.is_partitioned(prisma_client): + await self.partition_manager.ensure_partitions(prisma_client) + dropped: Final = await self.partition_manager.drop_partitions_older_than(prisma_client, cutoff_date) + verbose_proxy_logger.info( + "Dropped %d expired spend-log partitions: %s", + len(dropped), + dropped, + ) + # DROP only reclaims whole expired partitions. Expired rows can + # still sit in the DEFAULT partition (backfill, coverage gaps) + # or in a partition that spans the cutoff, so retention must + # also delete those stragglers row-wise. + total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) + verbose_proxy_logger.info( + "Deleted %s expired logs not covered by dropped partitions", total_deleted + ) + else: + total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) + verbose_proxy_logger.info("Deleted %s logs", total_deleted) - index_deleted: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date) - verbose_proxy_logger.info("Deleted %s expired tool index rows", index_deleted) + index_deleted: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date) + verbose_proxy_logger.info("Deleted %s expired tool index rows", index_deleted) - sessions_deleted: Final = await self._delete_old_auto_router_sessions(prisma_client, cutoff_date) + sessions_deleted: Final = await self._delete_old_auto_router_sessions( + prisma_client, self._auto_router_session_cutoff(retention_seconds) + ) verbose_proxy_logger.info("Deleted %s expired auto-router session rollups", sessions_deleted) except Exception as e: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8bd6f682d20..86b6e6775e3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5932,7 +5932,9 @@ class ProxyConfig: """ Reschedule the spend log cleanup job based on current general_settings. This is called when maximum_spend_logs_retention_period is updated dynamically. - If the retention period is None, the job will be removed. + The job is always rescheduled: spend-log pruning inside it still requires the + retention setting, but the auto-router rollup is garbage collected past its + read horizon regardless. """ global scheduler, general_settings, prisma_client if scheduler is None: @@ -5945,53 +5947,50 @@ class ProxyConfig: except Exception: pass # Job might not exist, which is fine - # Schedule new job if retention period is set (not None) - retention_period: Final = general_settings.get("maximum_spend_logs_retention_period") - if retention_period is not None: - from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( - SpendLogCleanup, + from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( + SpendLogCleanup, + ) + + spend_log_cleanup: Final = SpendLogCleanup() + cleanup_cron: Final = general_settings.get("maximum_spend_logs_cleanup_cron") + + if cleanup_cron: + from apscheduler.triggers.cron import CronTrigger + + try: + cron_trigger: Final = CronTrigger.from_crontab(cleanup_cron) + scheduler.add_job( + spend_log_cleanup.cleanup_old_spend_logs, + cron_trigger, + args=[prisma_client], + id="spend_log_cleanup_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info("Spend log cleanup rescheduled with cron: %s", cleanup_cron) + except ValueError: + verbose_proxy_logger.error("Invalid maximum_spend_logs_cleanup_cron value: %s", cleanup_cron) + else: + # Interval-based scheduling (existing behavior) + from litellm.litellm_core_utils.duration_parser import ( + duration_in_seconds, ) - spend_log_cleanup: Final = SpendLogCleanup() - cleanup_cron: Final = general_settings.get("maximum_spend_logs_cleanup_cron") - - if cleanup_cron: - from apscheduler.triggers.cron import CronTrigger - - try: - cron_trigger: Final = CronTrigger.from_crontab(cleanup_cron) - scheduler.add_job( - spend_log_cleanup.cleanup_old_spend_logs, - cron_trigger, - args=[prisma_client], - id="spend_log_cleanup_job", - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - verbose_proxy_logger.info("Spend log cleanup rescheduled with cron: %s", cleanup_cron) - except ValueError: - verbose_proxy_logger.error("Invalid maximum_spend_logs_cleanup_cron value: %s", cleanup_cron) - else: - # Interval-based scheduling (existing behavior) - from litellm.litellm_core_utils.duration_parser import ( - duration_in_seconds, + retention_interval: Final = general_settings.get("maximum_spend_logs_retention_interval", "1d") + try: + interval_seconds: Final = duration_in_seconds(retention_interval) + scheduler.add_job( + spend_log_cleanup.cleanup_old_spend_logs, + "interval", + seconds=interval_seconds + random.randint(0, 60), + args=[prisma_client], + id="spend_log_cleanup_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - - retention_interval: Final = general_settings.get("maximum_spend_logs_retention_interval", "1d") - try: - interval_seconds: Final = duration_in_seconds(retention_interval) - scheduler.add_job( - spend_log_cleanup.cleanup_old_spend_logs, - "interval", - seconds=interval_seconds + random.randint(0, 60), - args=[prisma_client], - id="spend_log_cleanup_job", - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - verbose_proxy_logger.info("Spend log cleanup rescheduled with interval: %s", retention_interval) - except ValueError: - verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") + verbose_proxy_logger.info("Spend log cleanup rescheduled with interval: %s", retention_interval) + except ValueError: + verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") async def _update_general_settings(self, db_general_settings: Json | None): """ @@ -8204,6 +8203,19 @@ class ProxyStartupEvent: f"({tag_spend_update_interval / batch_writing_interval:.1f}x main job interval)" ) + ### AUTO-ROUTER BENCHMARKS ROLLUP (separate scheduler job) ### + from litellm.proxy.utils import update_auto_router_sessions + + scheduler.add_job( + update_auto_router_sessions, + "interval", + seconds=batch_writing_interval, + args=(prisma_client, proxy_logging_obj), + id="update_auto_router_sessions_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + ### MONITOR SPEND LOGS QUEUE (queue-size-based job) ### if general_settings.get("disable_spend_logs", False) is False: from litellm.proxy.utils import _monitor_spend_logs_queue @@ -8320,42 +8332,43 @@ class ProxyStartupEvent: await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) ### SPEND LOG CLEANUP ### - if general_settings.get("maximum_spend_logs_retention_period") is not None: - spend_log_cleanup: Final = SpendLogCleanup() - cleanup_cron: Final = general_settings.get("maximum_spend_logs_cleanup_cron") + ## Always scheduled: spend-log pruning inside it still requires the retention + ## setting, but the auto-router rollup is collected past its read horizon regardless + spend_log_cleanup: Final = SpendLogCleanup() + cleanup_cron: Final = general_settings.get("maximum_spend_logs_cleanup_cron") - if cleanup_cron: - from apscheduler.triggers.cron import CronTrigger + if cleanup_cron: + from apscheduler.triggers.cron import CronTrigger - try: - cron_trigger: Final = CronTrigger.from_crontab(cleanup_cron) - scheduler.add_job( - spend_log_cleanup.cleanup_old_spend_logs, - cron_trigger, - args=[prisma_client], - id="spend_log_cleanup_job", - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - verbose_proxy_logger.info("Spend log cleanup scheduled with cron: %s", cleanup_cron) - except ValueError: - verbose_proxy_logger.error("Invalid maximum_spend_logs_cleanup_cron value: %s", cleanup_cron) - else: - # Interval-based scheduling (existing behavior) - retention_interval: Final = general_settings.get("maximum_spend_logs_retention_interval", "1d") - try: - interval_seconds: Final = duration_in_seconds(retention_interval) - scheduler.add_job( - spend_log_cleanup.cleanup_old_spend_logs, - "interval", - seconds=interval_seconds + random.randint(0, 60), - args=[prisma_client], - id="spend_log_cleanup_job", - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - except ValueError: - verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") + try: + cron_trigger: Final = CronTrigger.from_crontab(cleanup_cron) + scheduler.add_job( + spend_log_cleanup.cleanup_old_spend_logs, + cron_trigger, + args=[prisma_client], + id="spend_log_cleanup_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info("Spend log cleanup scheduled with cron: %s", cleanup_cron) + except ValueError: + verbose_proxy_logger.error("Invalid maximum_spend_logs_cleanup_cron value: %s", cleanup_cron) + else: + # Interval-based scheduling (existing behavior) + retention_interval: Final = general_settings.get("maximum_spend_logs_retention_interval", "1d") + try: + interval_seconds: Final = duration_in_seconds(retention_interval) + scheduler.add_job( + spend_log_cleanup.cleanup_old_spend_logs, + "interval", + seconds=interval_seconds + random.randint(0, 60), + args=[prisma_client], + id="spend_log_cleanup_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + except ValueError: + verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") ### CHECK BATCH COST ### if llm_router is not None and PROXY_BATCH_POLLING_ENABLED: try: diff --git a/litellm/proxy/spend_tracking/auto_router_benchmarks.py b/litellm/proxy/spend_tracking/auto_router_benchmarks.py index aeee35db0d6..d8f478756e0 100644 --- a/litellm/proxy/spend_tracking/auto_router_benchmarks.py +++ b/litellm/proxy/spend_tracking/auto_router_benchmarks.py @@ -20,6 +20,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient MAX_WINDOW_DAYS: Final = 30 +AUTO_ROUTER_SESSION_RETENTION_DAYS: Final = MAX_WINDOW_DAYS + 1 @dataclass(frozen=True, slots=True) @@ -228,13 +229,17 @@ def summarize(counters: _Counters) -> AutoRouterBenchmark: ) -def clamp_window(start_date: date, end_date: date) -> tuple[datetime, datetime]: - """The half-open UTC interval to read, clamped to ``MAX_WINDOW_DAYS``; ``end_date`` is - inclusive to the caller, so the upper bound is the start of the following day.""" - span_start: Final = max(start_date, end_date - timedelta(days=MAX_WINDOW_DAYS - 1)) +def clamp_window(start_date: date, end_date: date, today: date) -> tuple[datetime, datetime]: + """The half-open UTC interval to read, clamped into the most recent ``MAX_WINDOW_DAYS`` + ending ``today``; ``end_date`` is inclusive to the caller, so the upper bound is the + start of the following day. This recency clamp is what makes garbage collecting rows + past ``AUTO_ROUTER_SESSION_RETENTION_DAYS`` safe: a pruned row is one no window can + read. A window entirely before the horizon degenerates to an empty interval.""" + span_end: Final = min(end_date, today) + span_start: Final = max(start_date, today - timedelta(days=MAX_WINDOW_DAYS - 1)) return ( datetime.combine(span_start, time.min, tzinfo=timezone.utc), - datetime.combine(end_date + timedelta(days=1), time.min, tzinfo=timezone.utc), + datetime.combine(span_end + timedelta(days=1), time.min, tzinfo=timezone.utc), ) @@ -270,7 +275,7 @@ async def fetch_benchmarks( ) -> AutoRouterBenchmarksResponse: """Benchmarks for the window actually read; sessions are attributed to the window they started in, and the response echoes the clamped dates rather than the requested ones.""" - window_start, window_end = clamp_window(start_date, end_date) + window_start, window_end = clamp_window(start_date, end_date, today=datetime.now(timezone.utc).date()) rows: Final = await prisma_client.db.query_raw( _AGGREGATE_SQL, window_start.isoformat(), diff --git a/litellm/proxy/spend_tracking/auto_router_sessions.py b/litellm/proxy/spend_tracking/auto_router_sessions.py index 4771472571b..dd26767c695 100644 --- a/litellm/proxy/spend_tracking/auto_router_sessions.py +++ b/litellm/proxy/spend_tracking/auto_router_sessions.py @@ -182,10 +182,12 @@ ON CONFLICT (api_key, session_id, model_group) DO UPDATE SET class AutoRouterSessionQueue(BaseUpdateQueue): - """Stages turns in memory; writes an interval as one ordered batch on the spend flush. + """Stages turns in memory; a dedicated scheduler job writes an interval as one batch, + kept off the spend commit so a busy interval never delays budget enforcement. - Batched because a busy interval drains up to ``MAX_IN_MEMORY_QUEUE_FLUSH_COUNT`` turns. - Ordered because a session's classification depends on the turn before it. Never raises. + Sorted by session key so every pod locks rows in the same order (no cross-pod + deadlock), with a session's turns in the time order its classification depends on. + Never raises. """ async def flush(self, prisma_client: PrismaClient) -> None: @@ -194,7 +196,7 @@ class AutoRouterSessionQueue(BaseUpdateQueue): return try: async with prisma_client.db.batch_() as batcher: - for turn in sorted(staged, key=lambda staged_turn: staged_turn.started_at): + for turn in sorted(staged, key=lambda t: (t.api_key, t.session_id, t.model_group, t.started_at)): batcher.execute_raw(_UPSERT_SQL, *bind(turn)) except Exception as e: # noqa: BLE001 # a dashboard rollup must never fail spend tracking verbose_proxy_logger.warning("auto_router_sessions: dropped %d turns (%s)", len(staged), e) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 99e566c2da1..74a173d6b55 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5606,6 +5606,16 @@ async def update_daily_tag_spend( verbose_proxy_logger.error("Error updating daily tag spend: %s", e) +async def update_auto_router_sessions( + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, +): + """Separate scheduler job for the auto-router benchmarks rollup, kept off the + update_spend job so rollup upserts never extend the wall time of key, team and org + budget commits. The queue drops rather than blocks when full; the flush never raises.""" + await proxy_logging_obj.db_spend_update_writer.auto_router_session_queue.flush(prisma_client=prisma_client) + + async def update_spend_logs_job( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, diff --git a/tests/test_litellm/proxy/spend_tracking/test_auto_router_benchmarks.py b/tests/test_litellm/proxy/spend_tracking/test_auto_router_benchmarks.py index a04d65a56bf..725b5483b26 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_auto_router_benchmarks.py +++ b/tests/test_litellm/proxy/spend_tracking/test_auto_router_benchmarks.py @@ -273,17 +273,33 @@ class TestTotals: class TestWindow: + TODAY = dt.date(2026, 8, 3) + def test_end_date_is_inclusive(self): - start, end = clamp_window(dt.date(2026, 8, 3), dt.date(2026, 8, 3)) + start, end = clamp_window(dt.date(2026, 8, 3), dt.date(2026, 8, 3), today=self.TODAY) assert start == dt.datetime(2026, 8, 3, tzinfo=dt.timezone.utc) assert end == dt.datetime(2026, 8, 4, tzinfo=dt.timezone.utc) def test_a_wider_request_is_clamped_to_the_cap_measured_in_dates_spanned(self): - start, end = clamp_window(dt.date(2020, 1, 1), dt.date(2026, 8, 3)) + start, end = clamp_window(dt.date(2020, 1, 1), dt.date(2026, 8, 3), today=self.TODAY) assert (end.date() - start.date()).days == MAX_WINDOW_DAYS assert start == dt.datetime(2026, 7, 5, tzinfo=dt.timezone.utc) def test_a_window_inside_the_cap_is_left_alone(self): - start, end = clamp_window(dt.date(2026, 8, 1), dt.date(2026, 8, 3)) + start, end = clamp_window(dt.date(2026, 8, 1), dt.date(2026, 8, 3), today=self.TODAY) assert start == dt.datetime(2026, 8, 1, tzinfo=dt.timezone.utc) assert end == dt.datetime(2026, 8, 4, tzinfo=dt.timezone.utc) + + def test_the_window_cannot_start_before_the_recency_horizon(self): + """Retention leans on this: a row older than the horizon is one no window can + read, so pruning it is garbage collection rather than data loss.""" + start, _ = clamp_window(dt.date(2026, 6, 1), dt.date(2026, 7, 10), today=self.TODAY) + assert start == dt.datetime(2026, 7, 5, tzinfo=dt.timezone.utc) + + def test_a_future_end_date_is_clamped_to_today(self): + _, end = clamp_window(dt.date(2026, 8, 1), dt.date(2026, 9, 9), today=self.TODAY) + assert end == dt.datetime(2026, 8, 4, tzinfo=dt.timezone.utc) + + def test_a_window_entirely_before_the_horizon_reads_nothing(self): + start, end = clamp_window(dt.date(2026, 1, 1), dt.date(2026, 2, 1), today=self.TODAY) + assert end <= start diff --git a/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py b/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py index 82e7d0f9c53..50c5d3102d9 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py +++ b/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py @@ -6,12 +6,16 @@ here, is deciding whether a request is an auto-routed turn at all and what it co """ import datetime as dt +from dataclasses import replace +from types import SimpleNamespace import pytest from litellm.proxy.spend_tracking.auto_router_sessions import ( CACHE_TTL_1H_SECONDS, CACHE_TTL_5M_SECONDS, + AutoRouterSessionQueue, + TurnFacts, build_turn_facts, ttl_seconds, ) @@ -139,3 +143,50 @@ class TestCacheEvidence: }, } assert ttl_seconds(usage) is None + + +class TestFlushOrdering: + @pytest.mark.asyncio + async def test_turns_apply_in_session_key_order_and_in_time_order_within_a_session(self): + """Key order means every pod locks rollup rows in the same sequence (no cross-pod + deadlock); time order within a session is what classification depends on.""" + recorded: list[tuple[object, ...]] = [] + + class _DB: + def batch_(self): + return self + + async def __aenter__(self): + return SimpleNamespace(execute_raw=lambda sql, *params: recorded.append(params)) + + async def __aexit__(self, *exc: object) -> bool: + return False + + base = TurnFacts( + api_key="k", + session_id="a", + model_group="g", + router_kind="complexity", + baseline_model=None, + model="m", + started_at=0.0, + total_tokens=0, + spend=0.0, + baseline_spend=0.0, + cache_hit=False, + cache_creation_tokens=0, + cached_prefix_tokens=0, + ttl_seconds=None, + ) + queue = AutoRouterSessionQueue() + for turn in ( + replace(base, session_id="b", started_at=3.0), + replace(base, started_at=4.0), + replace(base, started_at=2.0), + replace(base, session_id="b", started_at=1.0), + ): + queue.update_queue.put_nowait(turn) + + await queue.flush(SimpleNamespace(db=_DB())) + + assert [(params[1], params[6]) for params in recorded] == [("a", 2.0), ("a", 4.0), ("b", 1.0), ("b", 3.0)] diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index e663e8b2de6..a360598930f 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -193,11 +193,16 @@ async def test_cleanup_old_spend_logs_batch_deletion(): tool_index_sql = mock_db.execute_raw.call_args_list[3][0][0] assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in tool_index_sql - # The auto-router rollup expires on the same cutoff, keyed on last activity so a - # conversation still running when the cutoff passes is not pruned mid-session - session_sql = mock_db.execute_raw.call_args_list[4][0][0] + # The auto-router rollup expires on the same cutoff when retention is shorter than + # its read horizon, keyed on last activity so a conversation still running when the + # cutoff passes is not pruned mid-session + from datetime import datetime, timedelta, timezone + + session_sql, session_cutoff = mock_db.execute_raw.call_args_list[4][0][:2] assert 'DELETE FROM "LiteLLM_AutoRouterSession"' in session_sql assert '"last_turn_at" <' in session_sql + expected_session_cutoff = datetime.now(timezone.utc) - timedelta(days=7) + assert abs((session_cutoff - expected_session_cutoff).total_seconds()) < 60 # The LiteLLM_DailyToolSpend rollup must outlive spend-log retention: it is # the only copy of tool spend history once its per-request sources expire, @@ -349,22 +354,32 @@ async def test_cleanup_uses_delete_when_not_partitioned(): @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 + With no retention period set, spend logs are untouched but the auto-router rollup is + still collected at its read horizon: rows past it are unreadable by the benchmarks + endpoint, so the table stays bounded without any configuration. """ + from datetime import datetime, timedelta, timezone + + from litellm.proxy.spend_tracking.auto_router_benchmarks import AUTO_ROUTER_SESSION_RETENTION_DAYS + mock_prisma_client = MagicMock() - mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.execute_raw = AsyncMock(return_value=0) cleaner = SpendLogCleanup(general_settings={}) # no retention await cleaner.cleanup_old_spend_logs(mock_prisma_client) - mock_prisma_client.db.execute_raw.assert_not_called() + assert mock_prisma_client.db.execute_raw.await_count == 1 + session_sql, cutoff = mock_prisma_client.db.execute_raw.call_args_list[0][0][:2] + assert 'DELETE FROM "LiteLLM_AutoRouterSession"' in session_sql + expected_cutoff = datetime.now(timezone.utc) - timedelta(days=AUTO_ROUTER_SESSION_RETENTION_DAYS) + assert abs((cutoff - expected_cutoff).total_seconds()) < 60 @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. + When another pod holds the cleanup lock, nothing is deleted (including the always-on + rollup collection) and the lock is not released by this pod. """ mock_prisma_client = MagicMock() mock_prisma_client.db.execute_raw = AsyncMock() @@ -372,17 +387,17 @@ async def test_lock_not_released_when_not_acquired(): 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.acquire_lock = AsyncMock(return_value=False) 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.acquire_lock.assert_called_once() mock_pod_lock_manager.release_lock.assert_not_called() + mock_prisma_client.db.execute_raw.assert_not_called() @pytest.mark.asyncio