Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py: 19%
120 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Manages native Postgres range partitions for the LiteLLM_SpendLogs table.
4At high request volume, retention via batched DELETE leaves dead tuples that
5autovacuum cannot reclaim fast enough, so the table keeps growing on disk. When
6the table is range-partitioned on startTime, dropping old data becomes a
7DROP TABLE on a whole partition: an instant metadata operation that returns disk
8to the OS immediately, with no tombstones and no vacuum.
10This manager only acts when use_spend_logs_partitioning is enabled in
11general_settings AND the table is already partitioned (set up via the
12db_scripts/partition_spend_logs.sql runbook). Without both, the cleanup job
13keeps the batched-DELETE path, so existing deployments are untouched.
14"""
16import re
17from collections.abc import Callable
18from datetime import date, datetime, timedelta, timezone
19from typing import (
20 TYPE_CHECKING,
21 Final,
22 TypeAlias,
23 cast, # noqa: TID251 # db.tx is reached through untyped __getattr__ delegation
24)
26from litellm._logging import verbose_proxy_logger
27from litellm.constants import (
28 SPEND_LOG_PARTITION_INTERVAL,
29 SPEND_LOG_PARTITION_PRECREATE_AHEAD,
30)
32if TYPE_CHECKING: 32 ↛ 33line 32 didn't jump to line 33 because the condition on line 32 was never true
33 from prisma.client import TransactionManager
35 from litellm.proxy.utils import PrismaClient
37SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs"
39RemainingTimeoutMs: TypeAlias = Callable[[], "int | None"]
40"""
41The per-statement bound in milliseconds, or None once the caller's budget is
42spent.
44Injected rather than passed as a number so it is re-evaluated before EVERY
45statement: a value read once at entry would let a loop issue N statements each
46bounded by the budget that was left before the first of them, which is not a
47bound on the loop at all. The caller owns the policy; this module only asks how
48much time it may still use.
49"""
51PartitionInterval = str # "day" | "week" | "month"
53VALID_PARTITION_INTERVALS: Final = {"day", "week", "month"}
55_BOUND_UPPER_RE: Final = re.compile(r"TO \('([^']+)'\)")
58def period_start(day: date, interval: PartitionInterval) -> date:
59 """First day of the partition period that `day` falls into (UTC)."""
60 if interval == "day":
61 return day
62 if interval == "week":
63 return day - timedelta(days=day.weekday())
64 if interval == "month":
65 return day.replace(day=1)
66 raise ValueError(f"Unsupported partition interval: {interval}")
69def next_period_start(start: date, interval: PartitionInterval) -> date:
70 if interval == "day":
71 return start + timedelta(days=1)
72 if interval == "week":
73 return start + timedelta(days=7)
74 if interval == "month":
75 if start.month == 12:
76 return start.replace(year=start.year + 1, month=1)
77 return start.replace(month=start.month + 1)
78 raise ValueError(f"Unsupported partition interval: {interval}")
81def partition_name(start: date) -> str:
82 return f"{SPEND_LOGS_TABLE}_p{start.strftime('%Y%m%d')}"
85def upcoming_partitions(today: date, interval: PartitionInterval, ahead: int) -> list[tuple[str, date, date]]:
86 """
87 Specs (name, lower_inclusive, upper_exclusive) for the current period plus
88 the next `ahead` periods, so writes always have a partition to land in.
89 """
90 specs: Final[list[tuple[str, date, date]]] = []
91 start = period_start(today, interval)
92 for _ in range(ahead + 1):
93 upper = next_period_start(start, interval)
94 specs.append((partition_name(start), start, upper))
95 start = upper
96 return specs
99def parse_partition_upper_bound(bound_expr: str) -> datetime | None:
100 """
101 Upper bound of a Postgres partition from its `pg_get_expr(relpartbound)`
102 string, e.g. "FOR VALUES FROM ('2026-06-01 00:00:00') TO ('2026-06-02 00:00:00')".
103 Returns None for the DEFAULT partition or anything we cannot parse, so such
104 partitions are never selected for dropping.
105 """
106 if "DEFAULT" in bound_expr.upper():
107 return None
108 match: Final = _BOUND_UPPER_RE.search(bound_expr)
109 if match is None:
110 return None
111 try:
112 return datetime.fromisoformat(match.group(1))
113 except ValueError:
114 return None
117def select_partitions_to_drop(partitions: list[tuple[str, datetime | None]], cutoff: datetime) -> list[str]:
118 """
119 Names of partitions whose entire range is older than `cutoff` (upper bound
120 <= cutoff). `cutoff` and the bounds are UTC-naive. Partitions without a
121 parseable upper bound (e.g. DEFAULT) are kept.
122 """
123 return [name for name, upper in partitions if upper is not None and upper <= cutoff]
126_TX_COMMIT_SLACK: Final = timedelta(seconds=5)
129def _bounded_tx(prisma_client: "PrismaClient", timeout_ms: int) -> "TransactionManager":
130 """
131 Open an interactive transaction that outlives the statement bound it
132 carries. prisma's default 5s transaction timeout would close it mid
133 lock-wait, after which the engine answers the next call with a 422.
134 """
135 return cast( # cast-ok: PrismaWrapper delegates tx via __getattr__ (untyped)
136 "TransactionManager",
137 prisma_client.db.tx(timeout=timedelta(milliseconds=timeout_ms) + _TX_COMMIT_SLACK),
138 )
141class SpendLogsPartitionManager:
142 def __init__(
143 self,
144 interval: PartitionInterval = SPEND_LOG_PARTITION_INTERVAL,
145 precreate_ahead: int = SPEND_LOG_PARTITION_PRECREATE_AHEAD,
146 ):
147 if interval not in VALID_PARTITION_INTERVALS:
148 verbose_proxy_logger.warning(
149 "Invalid SPEND_LOG_PARTITION_INTERVAL %r, falling back to 'day'. Supported values: %s",
150 interval,
151 sorted(VALID_PARTITION_INTERVALS),
152 )
153 interval = "day"
154 self.interval = interval
155 self.precreate_ahead = precreate_ahead
157 async def is_partitioned(self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs) -> bool:
158 budget_ms: Final = remaining_timeout_ms()
159 if budget_ms is None:
160 return False
161 try:
162 async with _bounded_tx(prisma_client, budget_ms) as tx:
163 await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}")
164 rows: Final = await tx.query_raw(
165 """
166 SELECT EXISTS (
167 SELECT 1
168 FROM pg_partitioned_table pt
169 JOIN pg_class c ON c.oid = pt.partrelid
170 JOIN pg_namespace n ON n.oid = c.relnamespace
171 WHERE c.relname = $1
172 AND n.nspname = current_schema()
173 ) AS partitioned
174 """,
175 SPEND_LOGS_TABLE,
176 )
177 except Exception as e:
178 verbose_proxy_logger.warning(
179 "Could not determine if %s is partitioned, assuming it is not: %s",
180 SPEND_LOGS_TABLE,
181 e,
182 )
183 return False
184 return bool(rows and rows[0].get("partitioned"))
186 @staticmethod
187 async def _execute_bounded_ddl(prisma_client: "PrismaClient", statement: str, timeout_ms: int) -> None:
188 """
189 Run one DDL statement under a Postgres statement and lock timeout.
191 Partition DDL takes an ACCESS EXCLUSIVE lock, so an unbounded statement
192 queues behind any long-running reader for as long as that reader lives,
193 and the caller's run budget cannot cut it short. lock_timeout bounds the
194 wait for the lock and statement_timeout bounds the work itself, so a
195 partition this run cannot get is simply left for the next one.
196 """
197 async with _bounded_tx(prisma_client, timeout_ms) as tx:
198 await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
199 await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
200 await tx.execute_raw(statement)
202 async def ensure_partitions(
203 self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs
204 ) -> list[str]:
205 """
206 Ensure the current and upcoming partitions exist, returning the names
207 now present. CREATE TABLE IF NOT EXISTS is a no-op for partitions that
208 already exist, so this list is "ensured present", not "newly created".
209 """
210 ensured: Final[list[str]] = []
211 for name, lower, upper in upcoming_partitions(
212 datetime.now(timezone.utc).date(), self.interval, self.precreate_ahead
213 ):
214 budget_ms = remaining_timeout_ms()
215 if budget_ms is None:
216 verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run")
217 break
218 try:
219 await self._execute_bounded_ddl(
220 prisma_client,
221 f'CREATE TABLE IF NOT EXISTS "{name}" '
222 f'PARTITION OF "{SPEND_LOGS_TABLE}" '
223 f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')",
224 budget_ms,
225 )
226 ensured.append(name)
227 except Exception as e:
228 verbose_proxy_logger.warning("Failed to ensure spend-log partition %s: %s", name, e)
229 return ensured
231 async def _list_partitions(
232 self, prisma_client: "PrismaClient", timeout_ms: int
233 ) -> list[tuple[str, datetime | None]]:
234 async with _bounded_tx(prisma_client, timeout_ms) as tx:
235 await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
236 rows: Final = await tx.query_raw(
237 """
238 SELECT c.relname AS name,
239 pg_get_expr(c.relpartbound, c.oid) AS bound
240 FROM pg_inherits i
241 JOIN pg_class c ON c.oid = i.inhrelid
242 JOIN pg_class p ON p.oid = i.inhparent
243 JOIN pg_namespace n ON n.oid = p.relnamespace
244 WHERE p.relname = $1
245 AND n.nspname = current_schema()
246 """,
247 SPEND_LOGS_TABLE,
248 )
249 return [(row["name"], parse_partition_upper_bound(row.get("bound") or "")) for row in rows]
251 async def drop_partitions_older_than(
252 self, prisma_client: "PrismaClient", cutoff: datetime, remaining_timeout_ms: RemainingTimeoutMs
253 ) -> list[str]:
254 """DROP every partition whose whole range is older than `cutoff`."""
255 list_budget_ms: Final = remaining_timeout_ms()
256 if list_budget_ms is None:
257 return []
258 cutoff_naive: Final = cutoff.astimezone(timezone.utc).replace(tzinfo=None)
259 partitions: Final = await self._list_partitions(prisma_client, list_budget_ms)
260 to_drop: Final = select_partitions_to_drop(partitions, cutoff_naive)
261 dropped: Final[list[str]] = []
262 for name in to_drop:
263 budget_ms = remaining_timeout_ms()
264 if budget_ms is None:
265 verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run")
266 break
267 try:
268 await self._execute_bounded_ddl(prisma_client, f'DROP TABLE IF EXISTS "{name}"', budget_ms)
269 dropped.append(name)
270 except Exception as e:
271 verbose_proxy_logger.warning("Failed to drop spend-log partition %s: %s", name, e)
272 return dropped