Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/common_utils/reset_budget_job.py: 52%
626 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import asyncio
2import json
3import math
4import time
5from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
6from dataclasses import dataclass, field
7from datetime import datetime, timedelta, timezone
8from enum import Enum
9from types import MappingProxyType
10from typing import Final, Generic, Literal, Protocol, TypeVar
12from typing_extensions import assert_never
14import litellm
15from litellm._logging import verbose_proxy_logger
16from litellm.caching.dual_cache import DualCache
17from litellm.constants import (
18 GLOBAL_PROXY_SPEND_CACHE_KEY,
19 LITELLM_PROXY_BUDGET_NAME,
20 RESET_BUDGET_JOB_BATCH_SIZE,
21 RESET_BUDGET_JOB_LOCK_TTL_SECONDS,
22 RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN,
23 RESET_BUDGET_JOB_NAME,
24)
25from litellm.litellm_core_utils.duration_parser import duration_in_seconds
26from litellm.proxy._types import (
27 DB_RETRY_SAFE_ERROR_TYPES,
28 LiteLLM_BudgetTableFull,
29 Litellm_EntityType,
30 LiteLLM_TeamTable,
31 LiteLLM_UserTable,
32 LiteLLM_VerificationToken,
33)
34from litellm.proxy.common_utils.timezone_utils import (
35 BudgetResetSettings,
36 compute_budget_reset_at,
37 get_budget_reset_settings,
38)
39from litellm.proxy.common_utils.user_api_key_cache import (
40 end_user_cache_key,
41 model_access_group_cache_key,
42 model_access_group_spend_counter_key,
43 project_cache_key,
44 project_spend_counter_key,
45 tag_cache_key,
46)
47from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
48from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
49from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
50from litellm.proxy.utils import PrismaClient, ProxyLogging
51from litellm.repositories.organization_repository import OrganizationRepository
52from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable
53from litellm.repositories.project_repository import ProjectRepository
54from litellm.repositories.table_repositories import (
55 EndUserRepository,
56 ModelAccessGroupBudgetRepository,
57 TagRepository,
58 TeamMembershipRepository,
59)
60from litellm.repositories.team_repository import TeamRepository
61from litellm.repositories.unit_of_work import (
62 LinkedSpendResetWrites,
63 budget_cascade_unit_of_work,
64 spend_reset_unit_of_work,
65)
66from litellm.repositories.verification_token_repository import (
67 VerificationTokenRepository,
68)
69from litellm.types.services import ServiceTypes
71_RowT = TypeVar("_RowT")
74@dataclass(frozen=True, slots=True)
75class _RowReset(Generic[_RowT]):
76 row: _RowT
77 spend_decrement: float
80_LINKED_KEYS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"budget_duration": None, "spend": {"gt": 0}})
81_SPENT_ROWS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"spend": {"gt": 0}})
84class _BudgetLinkedRow(Protocol):
85 @property
86 def spend(self) -> float | None: ... 86 ↛ exitline 86 didn't return from function 'spend' because
88 @property
89 def budget_id(self) -> str | None: ... 89 ↛ exitline 89 didn't return from function 'budget_id' because
92class _TeamMembershipRow(_BudgetLinkedRow, Protocol):
93 @property
94 def user_id(self) -> str: ... 94 ↛ exitline 94 didn't return from function 'user_id' because
96 @property
97 def team_id(self) -> str: ... 97 ↛ exitline 97 didn't return from function 'team_id' because
100class _KeyRow(_BudgetLinkedRow, Protocol):
101 @property
102 def token(self) -> str: ... 102 ↛ exitline 102 didn't return from function 'token' because
105class _OrgRow(_BudgetLinkedRow, Protocol):
106 @property
107 def organization_id(self) -> str: ... 107 ↛ exitline 107 didn't return from function 'organization_id' because
110class _TagRow(_BudgetLinkedRow, Protocol):
111 @property
112 def tag_name(self) -> str: ... 112 ↛ exitline 112 didn't return from function 'tag_name' because
115class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol):
116 @property
117 def access_group_name(self) -> str: ... 117 ↛ exitline 117 didn't return from function 'access_group_name' because
120class _ProjectRow(_BudgetLinkedRow, Protocol):
121 @property
122 def project_id(self) -> str: ... 122 ↛ exitline 122 didn't return from function 'project_id' because
125class _EndUserRow(_BudgetLinkedRow, Protocol):
126 @property
127 def user_id(self) -> str: ... 127 ↛ exitline 127 didn't return from function 'user_id' because
130def _rollover_enabled() -> bool:
131 return litellm.budget_rollover is True
134def _rollover_cap(max_budget: float | None) -> float | None:
135 if max_budget is None or not math.isfinite(max_budget):
136 return None
137 return max_budget
140def _carried_spend(spend: float | None, cap: float | None) -> float:
141 if cap is None:
142 return 0.0
143 return max(0.0, (spend or 0.0) - cap)
146def _row_carried_spend(row: _BudgetLinkedRow, caps: Mapping[str, float]) -> float:
147 if not caps:
148 return 0.0
149 return _carried_spend(row.spend, caps.get(row.budget_id) if row.budget_id is not None else None)
152def _team_membership_counter_key(row: _TeamMembershipRow) -> str:
153 return f"spend:team_member:{row.user_id}:{row.team_id}"
156def _team_membership_cache_keys(row: _TeamMembershipRow) -> tuple[str, ...]:
157 return (f"{row.team_id}_{row.user_id}",)
160def _key_counter_key(row: _KeyRow) -> str:
161 return f"spend:key:{row.token}"
164def _key_cache_keys(row: _KeyRow) -> tuple[str, ...]:
165 return (row.token,)
168def _org_counter_key(row: _OrgRow) -> str:
169 return f"spend:org:{row.organization_id}"
172def _org_cache_keys(row: _OrgRow) -> tuple[str, ...]:
173 return (
174 f"org_id:{row.organization_id}",
175 f"org_id:{row.organization_id}:with_budget",
176 )
179def _tag_counter_key(row: _TagRow) -> str:
180 return f"spend:tag:{row.tag_name}"
183def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]:
184 return (tag_cache_key(row.tag_name),)
187def _model_access_group_counter_key(row: _ModelAccessGroupRow) -> str:
188 return model_access_group_spend_counter_key(row.access_group_name)
191def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]:
192 return (model_access_group_cache_key(row.access_group_name),)
195def _project_counter_key(row: _ProjectRow) -> str:
196 return project_spend_counter_key(row.project_id)
199def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]:
200 return (project_cache_key(row.project_id),)
203def _enduser_counter_key(row: _EndUserRow) -> str:
204 return f"spend:end_user:{row.user_id}"
207def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]:
208 return (end_user_cache_key(row.user_id),)
211def _budget_link_where(
212 budget_ids: Sequence[str],
213 extra: Mapping[str, object] = MappingProxyType({}),
214) -> dict[str, object]:
215 return {"budget_id": {"in": list(budget_ids)}, **extra}
218def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]:
219 """Customers whose cached spend a committed reset of these tiers invalidated.
221 Mirrors ``_queue_enduser_resets`` without its ``spend > 0`` filter, which
222 post-commit would match nobody.
223 """
224 linked: Final = _budget_link_where(budget_ids)
225 default_budget_id: Final = litellm.max_end_user_budget_id
226 if default_budget_id is None or default_budget_id not in budget_ids:
227 return linked
228 return {"OR": [linked, {"budget_id": None}]} # mutable-ok: prisma where filter must be a dict
231def _queue_budget_linked_resets(
232 writes: LinkedSpendResetWrites,
233 cascade: "_BudgetCascade",
234 extra: Mapping[str, object] = MappingProxyType({}),
235) -> None:
236 """Reset one linked table's spend for every expiring tier: tiers with a
237 rollover cap keep spend beyond the cap (decrement preserves writes racing
238 the reset), everything else is zeroed as before. Zero the under-cap rows
239 BEFORE decrementing the over-cap ones: the statements run sequentially in
240 one transaction, so the reverse order lets the zero re-match a row the
241 decrement just moved into the (0, cap] range and erase its carried spend."""
242 for budget_id, cap in cascade.rollover_caps.items():
243 writes.queue_spend_zero(where={"budget_id": budget_id, **extra, "spend": {"gt": 0, "lte": cap}})
244 writes.queue_spend_decrement(where={"budget_id": budget_id, **extra, "spend": {"gt": cap}}, amount=cap)
245 plain_ids: Final = tuple(bid for bid in cascade.budget_ids if bid not in cascade.rollover_caps)
246 if plain_ids:
247 writes.queue_spend_zero(where=_budget_link_where(plain_ids, extra))
250def _queue_enduser_resets(writes: LinkedSpendResetWrites, cascade: "_BudgetCascade") -> None:
251 """End users reset on the budget link like every other gated table, plus a
252 NULL-budget_id branch: rows created implicitly persist no link and ride the
253 default tier (litellm.max_end_user_budget_id).
255 Matching on the link rather than enumerating user ids keeps a statement's
256 bind count proportional to the expiring tiers instead of the customer
257 population, which past ~32,700 dependents exceeds PostgreSQL's per-statement
258 bind ceiling and wedges the cascade permanently (#40564).
259 """
260 _queue_budget_linked_resets(writes, cascade, extra=_SPENT_ROWS_WHERE)
261 default_budget_id: Final = litellm.max_end_user_budget_id
262 if default_budget_id is None or default_budget_id not in cascade.budget_ids:
263 return
264 cap: Final = cascade.rollover_caps.get(default_budget_id)
265 if cap is None:
266 writes.queue_spend_zero(where={"budget_id": None, **_SPENT_ROWS_WHERE})
267 return
268 writes.queue_spend_zero(where={"budget_id": None, "spend": {"gt": 0, "lte": cap}})
269 writes.queue_spend_decrement(where={"budget_id": None, "spend": {"gt": cap}}, amount=cap)
272@dataclass(frozen=True, slots=True)
273class _BudgetCascade:
274 """Everything one budget-tier reset touches, resolved before any write."""
276 budgets: tuple[LiteLLM_BudgetTableFull, ...] = ()
277 budget_ids: tuple[str, ...] = ()
278 budget_resets: tuple[tuple[str, datetime], ...] = ()
279 counter_resets: tuple[tuple[str, float], ...] = ()
280 cache_keys: tuple[str, ...] = ()
281 rollover_caps: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({}))
284@dataclass(frozen=True, slots=True)
285class _EndUserWalk:
286 """Where the customer walk stands. ``cursor`` is None once it is done, and
287 ``truncated`` says a failed page read cut it short of the tail."""
289 cursor: str | None = ""
290 invalidated: int = 0
291 truncated: bool = False
294_ENDUSER_WALK_DONE: Final = _EndUserWalk(cursor=None)
297@dataclass(frozen=True, slots=True)
298class _BudgetCascadeCommitted:
299 cascade: _BudgetCascade
300 advanced: int
301 endusers: _EndUserWalk
304@dataclass(frozen=True, slots=True)
305class _BudgetCascadeFailed:
306 cascade: _BudgetCascade
307 error: Exception
310_EMPTY_CASCADE: Final = _BudgetCascade()
312_InvalidatedCache = Literal["spend counter", "user_api_key_cache"]
315@dataclass(frozen=True, slots=True)
316class _ChunkOutcome:
317 """One chunk of a reset phase: rows read, and rows whose new budget_reset_at
318 cleared the due cutoff. Anything else is still due and would come straight
319 back on the next fetch, so it is not progress."""
321 fetched: int
322 advanced: int
325_NO_PROGRESS: Final = _ChunkOutcome(fetched=0, advanced=0)
328def _as_utc(moment: datetime) -> datetime:
329 return moment if moment.tzinfo is not None else moment.replace(tzinfo=timezone.utc)
332def _count_advanced(reset_ats: Iterable[object], cutoff: datetime) -> int:
333 """How many rows the write actually moved past the due cutoff.
335 A budget_duration of "0s" (or one the parser cannot read) resolves to the
336 current time, so the row is written and stays due. Counting it as progress
337 would re-read the same chunk until the per-run cap on every tick.
338 """
339 utc_cutoff: Final = _as_utc(cutoff)
340 return sum(1 for reset_at in reset_ats if isinstance(reset_at, datetime) and _as_utc(reset_at) > utc_cutoff)
343def _phase_is_drained(outcome: _ChunkOutcome) -> bool:
344 """A short chunk means the due rows ran out. A full chunk that advanced
345 nothing would be re-read unchanged forever, so it ends the phase too and
346 those rows wait for the next tick."""
347 return outcome.fetched < RESET_BUDGET_JOB_BATCH_SIZE or outcome.advanced == 0
350async def _run_phase_in_chunks(process_chunk: Callable[[], Awaitable[_ChunkOutcome]]) -> None:
351 """Drive one reset phase a chunk at a time, capped so a single run cannot
352 spin unbounded: leftovers are picked up by the next tick."""
353 for _ in range(RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN): 353 ↛ exitline 353 didn't return from function '_run_phase_in_chunks' because the loop on line 353 didn't complete
354 if _phase_is_drained(await process_chunk()): 354 ↛ 353line 354 didn't jump to line 353 because the condition on line 354 was always true
355 return
358@dataclass(frozen=True, slots=True)
359class _LazyJson:
360 """Serialize only if a log record is actually emitted.
362 ``logger.debug("... %s", json.dumps(rows))`` evaluates the dump before the
363 logger decides to drop the record, so a chunk of rows is serialized on the
364 event loop on every tick at any log level. Passing this instead defers the
365 work to the formatter.
366 """
368 value: object
370 def __str__(self) -> str:
371 return json.dumps(self.value, indent=4, default=str)
374class _Lease(Enum):
375 """Whether this pod may sweep, and whether it owes a lock release."""
377 LEADER = "leader"
378 UNGUARDED = "unguarded"
379 FOLLOWER = "follower"
382async def _write_key_windows(prisma_client: PrismaClient, row_id: str, payload: str) -> None:
383 await VerificationTokenRepository(prisma_client).table.update(
384 where={"token": row_id},
385 data={"budget_limits": payload},
386 )
389async def _write_team_windows(prisma_client: PrismaClient, row_id: str, payload: str) -> None:
390 await TeamRepository(prisma_client).table.update(
391 where={"team_id": row_id},
392 data={"budget_limits": payload},
393 )
396@dataclass(frozen=True, slots=True)
397class _WindowSource:
398 """A table whose rows carry their own per-window budget limits."""
400 table: str
401 id_column: str
402 entity_type: Litellm_EntityType
403 counter_prefix: str
404 log_subject: str
405 retry_subject: str
406 write: Callable[[PrismaClient, str, str], Awaitable[None]]
408 def page_query(self) -> str:
409 """One keyset page, ordered by the primary key so the cursor never repeats a row.
411 prisma-client-python cannot null-filter a ``Json?`` column (no DbNull /
412 JsonNull sentinel, RobertCraigie/prisma-client-py#714), so the read stays
413 raw SQL; the table and column names are module constants, never input.
414 Writes still go through the ORM.
415 """
416 return (
417 f'SELECT {self.id_column}, budget_limits FROM "{self.table}" '
418 f"WHERE budget_limits IS NOT NULL AND {self.id_column} > $1 "
419 f"ORDER BY {self.id_column} LIMIT $2"
420 )
423_WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = (
424 _WindowSource(
425 table="LiteLLM_VerificationToken",
426 id_column="token",
427 entity_type=Litellm_EntityType.KEY,
428 counter_prefix="spend:key",
429 log_subject="keys",
430 retry_subject="key",
431 write=_write_key_windows,
432 ),
433 _WindowSource(
434 table="LiteLLM_TeamTable",
435 id_column="team_id",
436 entity_type=Litellm_EntityType.TEAM,
437 counter_prefix="spend:team",
438 log_subject="teams",
439 retry_subject="team",
440 write=_write_team_windows,
441 ),
442)
445def _budget_cascade_event_metadata(
446 cascade: _BudgetCascade, endusers: _EndUserWalk = _ENDUSER_WALK_DONE
447) -> dict[str, object]:
448 return {
449 "num_budgets_found": len(cascade.budgets),
450 "num_endusers_found": endusers.invalidated,
451 }
454class ResetBudgetJob:
455 """
456 Resets the budget for all the keys, users, and teams that need it
457 """
459 def __init__(
460 self,
461 proxy_logging_obj: ProxyLogging,
462 prisma_client: PrismaClient,
463 reset_settings: BudgetResetSettings | None = None,
464 pod_lock_manager: PodLockManager | None = None,
465 ):
466 self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
467 self.prisma_client: PrismaClient = prisma_client
468 self.reset_settings: BudgetResetSettings = reset_settings or get_budget_reset_settings()
469 self.pod_lock_manager: PodLockManager | None = pod_lock_manager
471 @property
472 def _new_batch(self) -> Callable[[], PrismaBatch]:
473 new_batch: Final[Callable[[], PrismaBatch]] = self.prisma_client.db.batch_
474 return new_batch
476 async def _lease_is_held(self, lock_manager: PodLockManager) -> bool:
477 """True only when the lease is readable and someone holds it.
479 An unreadable lock reports as unheld so the caller sweeps rather than
480 skipping; being wrong here costs a duplicate sweep, and the alternative
481 strands every expired budget at its cap.
482 """
483 if lock_manager.redis_cache is None:
484 return False
485 try:
486 lock_key: Final = lock_manager.get_redis_lock_key(RESET_BUDGET_JOB_NAME)
487 return bool(await lock_manager.redis_cache.async_get_cache(lock_key))
488 except Exception as exc: # noqa: BLE001 # an unreadable lease must not strand the sweep
489 verbose_proxy_logger.warning("Reset budget job: could not read the reset lease: %s", exc)
490 return False
492 async def _acquire_lease(self) -> _Lease:
493 """Elect one sweeper per tick.
495 Every pod schedules this job, and each one otherwise re-reads the whole
496 due population and writes it back at the same calendar boundary, so a
497 fleet multiplies one sweep's Postgres load by its replica count. A
498 deployment with no Redis-backed lock manager runs unguarded, as it
499 always has.
500 """
501 lock_manager: Final = self.pod_lock_manager
502 if lock_manager is None or lock_manager.redis_cache is None: 502 ↛ 505line 502 didn't jump to line 505 because the condition on line 502 was always true
503 return _Lease.UNGUARDED
505 if await lock_manager.acquire_lock(
506 cronjob_id=RESET_BUDGET_JOB_NAME,
507 ttl=RESET_BUDGET_JOB_LOCK_TTL_SECONDS,
508 ):
509 return _Lease.LEADER
511 if await self._lease_is_held(lock_manager):
512 verbose_proxy_logger.debug("Reset budget job: another pod holds the reset lease, skipping this tick")
513 return _Lease.FOLLOWER
515 # acquire_lock reports contention and an unreachable Redis identically, so
516 # treating a failed acquire as contention would skip the sweep on every pod
517 # at once for as long as Redis is down. Sweeping unguarded costs duplicate
518 # work; not sweeping leaves every expired budget pinned at its cap.
519 verbose_proxy_logger.warning(
520 "Reset budget job: could not take the reset lease and no other pod holds it, "
521 "sweeping unguarded rather than skipping the tick"
522 )
523 return _Lease.UNGUARDED
525 async def reset_budget(
526 self,
527 ):
528 """
529 Gets all the non-expired keys for a db, which need spend to be reset
531 Resets their spend
533 Updates db
535 Runs on one pod per tick where a Redis lease is available.
536 """
537 if self.prisma_client is None: 537 ↛ 538line 537 didn't jump to line 538 because the condition on line 537 was never true
538 return
540 lease: Final = await self._acquire_lease()
541 if lease is _Lease.FOLLOWER: 541 ↛ 542line 541 didn't jump to line 542 because the condition on line 541 was never true
542 return
544 try:
545 await self.reset_budget_for_litellm_keys()
546 await self.reset_budget_for_litellm_users()
547 await self.reset_budget_for_litellm_teams()
548 await self.reset_budget_for_litellm_budget_table()
549 await self.reset_budget_windows()
550 finally:
551 if lease is _Lease.LEADER and self.pod_lock_manager is not None: 551 ↛ 552line 551 didn't jump to line 552 because the condition on line 551 was never true
552 await self.pod_lock_manager.release_lock(cronjob_id=RESET_BUDGET_JOB_NAME)
554 async def _with_db_retry(self, operation: Callable[[], Awaitable[_RowT]], *, reason: str) -> _RowT:
555 """Reconnect and retry once on a transport error, so a dropped connection
556 costs one retry instead of the whole tick.
557 """
558 return await call_with_db_reconnect_retry(self.prisma_client, operation, reason=reason)
560 async def _with_db_write_retry(self, operation: Callable[[], Awaitable[_RowT]], *, reason: str) -> _RowT:
561 """Same, for writes: only replay when the statements provably never
562 reached the database. A reset zeroes spend unconditionally, so replaying
563 an ambiguous commit would erase spend accrued since it landed.
564 """
565 return await call_with_db_reconnect_retry(
566 self.prisma_client,
567 operation,
568 reason=reason,
569 retry_safe_error_types=DB_RETRY_SAFE_ERROR_TYPES,
570 )
572 @staticmethod
573 async def _invalidate_spend_counter(counter_key: str) -> None:
574 """Drop a spend counter so the next read reseeds from the committed DB
575 row, the only value that includes increments that raced the reset.
577 Call AFTER the DB write commits. Clearing Redis before the DB
578 commit opens a window where get_current_spend reads 0 from Redis
579 while the DB still holds the pre-reset value, allowing bypass.
580 """
581 try:
582 from litellm.proxy.proxy_server import spend_counter_cache
584 spend_counter_cache.in_memory_cache.delete_cache(key=counter_key)
585 if spend_counter_cache.redis_cache is not None:
586 try:
587 await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key)
588 except Exception as redis_err:
589 verbose_proxy_logger.warning(
590 "Failed to reset spend counter %s in Redis: %s. "
591 "Budget may be over-enforced until counter expires.",
592 counter_key,
593 redis_err,
594 )
595 except Exception as e:
596 verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e)
598 @staticmethod
599 async def _invalidate_global_proxy_spend_cache() -> None:
600 """Drop the cached global-proxy spend accumulator after the proxy
601 budget aggregate row is reset, so the next auth-time load reads the
602 zeroed row instead of a stale (potentially never-expiring) counter.
603 """
604 await ResetBudgetJob._invalidate_user_api_key_cache_entry(GLOBAL_PROXY_SPEND_CACHE_KEY)
606 @staticmethod
607 async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None:
608 """Drop a stale management-cache entry so the next read fetches from DB.
610 Tags and end-users are not reseeded by ``SpendCounterReseed.from_db``;
611 for those, when the spend counter expires the budget check falls back
612 to ``cached_obj.spend``. Keys, orgs, and team memberships are reseeded
613 from the DB, but auth still may consult ``user_api_key_cache`` objects
614 whose ``.spend`` field can lag a cross-pod DB reset. Deleting the cache
615 entry forces the next auth-time fetch to reload the zeroed row from
616 Postgres.
617 """
618 try:
619 from litellm.proxy.proxy_server import user_api_key_cache
621 await user_api_key_cache.async_delete_cache(key=cache_key)
622 except Exception as e:
623 verbose_proxy_logger.warning(
624 "Failed to invalidate user_api_key_cache entry %s: %s",
625 cache_key,
626 e,
627 )
629 @staticmethod
630 async def _invalidate_caches(counter_keys: Sequence[str], cache_keys: Sequence[str]) -> None:
631 """Batch twin of ``_invalidate_spend_counter`` and
632 ``_invalidate_user_api_key_cache_entry``, after the commit like both:
633 one round trip per chunk where a tier's dependents are unbounded."""
634 await ResetBudgetJob._invalidate_cache("spend counter", counter_keys)
635 await ResetBudgetJob._invalidate_cache("user_api_key_cache", cache_keys)
637 @staticmethod
638 async def _invalidate_cache(cache: _InvalidatedCache, keys: Sequence[str]) -> None:
639 """One cache's share of a batch, awaited separately so either failing
640 still leaves the other invalidated."""
641 if not keys: 641 ↛ 643line 641 didn't jump to line 643 because the condition on line 641 was always true
642 return
643 try:
644 from litellm.proxy.proxy_server import spend_counter_cache, user_api_key_cache
646 match cache:
647 case "spend counter":
648 await spend_counter_cache.async_delete_cache_keys(keys)
649 case "user_api_key_cache":
650 await user_api_key_cache.async_delete_cache_keys(keys)
651 case _:
652 assert_never(cache)
653 except Exception as e:
654 verbose_proxy_logger.warning(
655 "Failed to invalidate %d %s entries: %s. Budgets may be over-enforced until they expire.",
656 len(keys),
657 cache,
658 e,
659 )
661 async def _fetch_linked_rows(
662 self,
663 table: SpendLinkedTable[_RowT],
664 where: Mapping[str, object],
665 log_subject: str,
666 ) -> tuple[_RowT, ...]:
667 """Read the rows the cascade will zero, so their counters can be
668 invalidated once the transaction commits."""
669 try:
670 return tuple(
671 await self._with_db_retry(
672 lambda: table.find_many(where=where),
673 reason=f"reset_budget_read_{log_subject.replace(' ', '_')}_failure",
674 )
675 )
676 except Exception as e:
677 verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e)
678 return ()
680 async def _invalidate_enduser_caches(self, budget_ids: Sequence[str]) -> _EndUserWalk:
681 """Drop the cached spend of every customer the committed tier reset zeroed.
683 Paged like ``_reset_windows_for``, and capless for its reason too: the
684 customers on one tier are unbounded, and a cap cannot keep its position
685 across pod elections, so it would restart at the first customer forever.
686 """
687 if not budget_ids: 687 ↛ 689line 687 didn't jump to line 689 because the condition on line 687 was always true
688 return _ENDUSER_WALK_DONE
689 where: Final = _enduser_invalidation_where(budget_ids)
690 walk = _EndUserWalk()
691 while walk.cursor is not None:
692 walk = await self._invalidate_enduser_page(where=where, cursor=walk.cursor, reached=walk.invalidated)
693 return walk
695 async def _invalidate_enduser_page(self, where: Mapping[str, object], cursor: str, reached: int) -> _EndUserWalk:
696 """Invalidate one page of customers and say where the walk goes next."""
697 try:
698 rows: Final = await self._fetch_enduser_page(where=where, cursor=cursor)
699 except Exception as e:
700 verbose_proxy_logger.warning(
701 "Failed to fetch end users for cache invalidation after %s customers (cursor %r): %s. "
702 "The customers past that page keep their cached spend until it expires.",
703 reached,
704 cursor,
705 e,
706 )
707 return _EndUserWalk(cursor=None, invalidated=reached, truncated=True)
708 if not rows:
709 return _EndUserWalk(cursor=None, invalidated=reached)
710 await self._invalidate_caches(
711 counter_keys=tuple(_enduser_counter_key(row) for row in rows),
712 cache_keys=tuple(key for row in rows for key in _enduser_cache_keys(row)),
713 )
714 walked: Final = reached + len(rows)
715 if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE:
716 return _EndUserWalk(cursor=None, invalidated=walked)
717 return _EndUserWalk(cursor=rows[-1].user_id, invalidated=walked)
719 async def _fetch_enduser_page(self, where: Mapping[str, object], cursor: str) -> tuple[_EndUserRow, ...]:
720 """One keyset page of customers, ordered by primary key so the cursor never repeats a row."""
721 return tuple(
722 await self._with_db_retry(
723 lambda: EndUserRepository(self.prisma_client).table.find_many(
724 where={**where, "user_id": {"gt": cursor}}, # mutable-ok: prisma where filter must be a dict
725 order={"user_id": "asc"}, # mutable-ok: prisma order filter must be a dict
726 take=RESET_BUDGET_JOB_BATCH_SIZE,
727 ),
728 reason="reset_budget_read_endusers_failure",
729 )
730 )
732 async def _collect_budget_cascade(self, budgets_to_reset: Sequence[LiteLLM_BudgetTableFull]) -> _BudgetCascade:
733 """Resolve every row the expiring budget tiers gate, before any write.
735 Keys carrying their own budget_duration are left out: they run on their
736 own schedule via reset_budget_for_litellm_keys(), so sweeping them here
737 would reset them twice.
738 """
739 budget_ids: Final = tuple(b.budget_id for b in budgets_to_reset if b.budget_id is not None)
740 if not budget_ids: 740 ↛ 743line 740 didn't jump to line 743 because the condition on line 740 was always true
741 return _EMPTY_CASCADE
743 team_memberships: Final[tuple[_TeamMembershipRow, ...]] = await self._fetch_linked_rows(
744 table=TeamMembershipRepository(self.prisma_client).table,
745 where=_budget_link_where(budget_ids),
746 log_subject="team memberships",
747 )
748 keys: Final[tuple[_KeyRow, ...]] = await self._fetch_linked_rows(
749 table=VerificationTokenRepository(self.prisma_client).table,
750 where=_budget_link_where(budget_ids, _LINKED_KEYS_WHERE),
751 log_subject="keys",
752 )
753 orgs: Final[tuple[_OrgRow, ...]] = await self._fetch_linked_rows(
754 table=OrganizationRepository(self.prisma_client).table,
755 where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
756 log_subject="orgs",
757 )
758 tags: Final[tuple[_TagRow, ...]] = await self._fetch_linked_rows(
759 table=TagRepository(self.prisma_client).table,
760 where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
761 log_subject="tags",
762 )
763 model_access_groups: Final[tuple[_ModelAccessGroupRow, ...]] = await self._fetch_linked_rows(
764 table=ModelAccessGroupBudgetRepository(self.prisma_client).table,
765 where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
766 log_subject="model access groups",
767 )
768 projects: Final[tuple[_ProjectRow, ...]] = await self._fetch_linked_rows(
769 table=ProjectRepository(self.prisma_client).table,
770 where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
771 log_subject="projects",
772 )
773 rollover_caps: Final[Mapping[str, float]] = MappingProxyType(
774 { # mutable-ok: MappingProxyType wraps a one-shot dict comprehension
775 b.budget_id: cap
776 for b in budgets_to_reset
777 if b.budget_id is not None and (cap := _rollover_cap(b.max_budget)) is not None
778 }
779 if _rollover_enabled()
780 else {} # mutable-ok: empty sentinel immediately frozen by MappingProxyType
781 )
782 return _BudgetCascade(
783 budgets=tuple(budgets_to_reset),
784 budget_ids=budget_ids,
785 budget_resets=tuple(
786 (
787 b.budget_id,
788 compute_budget_reset_at(budget_duration=b.budget_duration, settings=self.reset_settings),
789 )
790 for b in budgets_to_reset
791 if b.budget_id is not None and b.budget_duration is not None
792 ),
793 counter_resets=(
794 *(
795 (_team_membership_counter_key(row), _row_carried_spend(row, rollover_caps))
796 for row in team_memberships
797 ),
798 *((_key_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in keys),
799 *((_org_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in orgs),
800 *((_tag_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in tags),
801 *(
802 (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps))
803 for row in model_access_groups
804 ),
805 *((_project_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in projects),
806 ),
807 rollover_caps=rollover_caps,
808 cache_keys=(
809 *(key for row in team_memberships for key in _team_membership_cache_keys(row)),
810 *(key for row in keys for key in _key_cache_keys(row)),
811 *(key for row in orgs for key in _org_cache_keys(row)),
812 *(key for row in tags for key in _tag_cache_keys(row)),
813 *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)),
814 *(key for row in projects for key in _project_cache_keys(row)),
815 ),
816 )
818 async def _commit_budget_cascade(self, cascade: _BudgetCascade) -> None:
819 """Zero the gated spend and advance ``budget_reset_at`` in one transaction.
821 Advancing the window on its own hides the tier from every later tick
822 while its dependents stay pinned at the cap for the whole window;
823 batching both means a mid-cascade failure persists nothing and the rows
824 stay due for the next run.
825 """
826 if not cascade.budget_ids: 826 ↛ 829line 826 didn't jump to line 829 because the condition on line 826 was always true
827 return
829 await self._with_db_write_retry(
830 lambda: self._commit_budget_cascade_once(cascade),
831 reason="reset_budget_write_budget_cascade_failure",
832 )
834 async def _commit_budget_cascade_once(self, cascade: _BudgetCascade) -> None:
835 async with budget_cascade_unit_of_work(self._new_batch) as uow:
836 _queue_budget_linked_resets(uow.team_memberships, cascade)
837 _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE)
838 _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
839 _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE)
840 _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE)
841 _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE)
842 _queue_enduser_resets(uow.endusers, cascade)
843 for budget_id, budget_reset_at in cascade.budget_resets:
844 uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at)
846 async def _invalidate_budget_cascade_caches(self, cascade: _BudgetCascade) -> None:
847 await self._invalidate_caches(
848 counter_keys=tuple(counter_key for counter_key, _ in cascade.counter_resets),
849 cache_keys=cascade.cache_keys,
850 )
852 async def _reset_expired_budget_cascade(self) -> _BudgetCascadeCommitted | _BudgetCascadeFailed:
853 now: Final = datetime.now(timezone.utc)
854 try:
855 budgets_to_reset: Final[Sequence[LiteLLM_BudgetTableFull] | None] = await self._with_db_retry(
856 lambda: self.prisma_client.get_data(
857 table_name="budget",
858 query_type="find_all",
859 reset_at=now,
860 limit=RESET_BUDGET_JOB_BATCH_SIZE,
861 ),
862 reason="reset_budget_read_budgets_failure",
863 )
864 cascade: Final = await self._collect_budget_cascade(budgets_to_reset or ())
865 except Exception as e:
866 return _BudgetCascadeFailed(cascade=_EMPTY_CASCADE, error=e)
868 try:
869 await self._commit_budget_cascade(cascade)
870 except Exception as e:
871 return _BudgetCascadeFailed(cascade=cascade, error=e)
873 await self._invalidate_budget_cascade_caches(cascade)
874 return _BudgetCascadeCommitted(
875 cascade=cascade,
876 advanced=_count_advanced(
877 (reset_at for _, reset_at in cascade.budget_resets),
878 cutoff=datetime.now(timezone.utc),
879 ),
880 endusers=await self._invalidate_enduser_caches(cascade.budget_ids),
881 )
883 async def reset_budget_for_litellm_budget_table(self) -> None:
884 """
885 Resets the spend a budget tier gates (end users, team members, keys,
886 orgs, tags, model access groups) and advances the tier's
887 budget_reset_at, atomically.
889 Caches are invalidated only after the transaction commits, so a failed
890 run cannot leave a zeroed counter in front of an un-reset DB row.
891 """
892 await _run_phase_in_chunks(self._reset_budget_for_litellm_budget_table_chunk)
894 async def _reset_budget_for_litellm_budget_table_chunk(self) -> _ChunkOutcome:
895 start_time: Final = time.time()
896 outcome: Final = await self._reset_expired_budget_cascade()
897 end_time: Final = time.time()
899 match outcome:
900 case _BudgetCascadeCommitted() as committed: 900 ↛ 917line 900 didn't jump to line 917 because the pattern on line 900 always matched
901 asyncio.create_task(
902 self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
903 service=ServiceTypes.RESET_BUDGET_JOB,
904 duration=end_time - start_time,
905 call_type="reset_budget_budget_table",
906 start_time=start_time,
907 end_time=end_time,
908 event_metadata={
909 **_budget_cascade_event_metadata(committed.cascade, committed.endusers),
910 "num_endusers_updated": committed.endusers.invalidated,
911 "num_endusers_failed": 0,
912 "enduser_invalidation_truncated": committed.endusers.truncated,
913 },
914 )
915 )
916 return _ChunkOutcome(fetched=len(committed.cascade.budgets), advanced=committed.advanced)
917 case _BudgetCascadeFailed(cascade=cascade, error=error):
918 verbose_proxy_logger.exception(
919 "Failed to reset the budget table cascade (team member, enduser, org, tag and model access "
920 "group spend, plus budget_reset_at); nothing was committed and the budgets stay due for the "
921 "next run: %s",
922 error,
923 exc_info=error,
924 )
925 asyncio.create_task(
926 self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
927 service=ServiceTypes.RESET_BUDGET_JOB,
928 duration=end_time - start_time,
929 error=error,
930 call_type="reset_budget_endusers",
931 start_time=start_time,
932 end_time=end_time,
933 event_metadata=_budget_cascade_event_metadata(cascade),
934 )
935 )
936 return _NO_PROGRESS
937 case _:
938 assert_never(outcome)
940 async def _write_key_reset_updates(self, updated_keys: Sequence[_RowReset[LiteLLM_VerificationToken]]) -> None:
941 """
942 Write per-row {spend, budget_reset_at} updates for keys.
944 Avoids the batched full-model update path, which trips
945 prisma.errors.DataError on any row carrying object_permission_id or
946 budget_limits (see #27730). Both fields are rejected by Prisma's
947 update input type for LiteLLM_VerificationToken, and the failure
948 aborts the entire batch — silently leaving spend over the cap and
949 budget_reset_at unchanged forever.
950 """
951 await self._with_db_write_retry(
952 lambda: self._write_key_reset_updates_once(updated_keys),
953 reason="reset_budget_write_keys_failure",
954 )
956 async def _write_key_reset_updates_once(self, updated_keys: Sequence[_RowReset[LiteLLM_VerificationToken]]) -> None:
957 async with spend_reset_unit_of_work(self._new_batch) as uow:
958 for k in updated_keys:
959 if k.row.token is None:
960 continue
961 uow.keys.queue_spend_reset(
962 token=k.row.token,
963 budget_reset_at=k.row.budget_reset_at,
964 spend_decrement=k.spend_decrement,
965 )
967 async def _write_user_reset_updates(self, updated_users: Sequence[_RowReset[LiteLLM_UserTable]]) -> None:
968 """
969 Write per-row {spend, budget_reset_at} updates for users.
971 Mirrors _write_key_reset_updates — avoids the full-model update path
972 that trips Prisma's DataError on rows carrying unrecognised fields
973 (see #27730).
974 """
975 await self._with_db_write_retry(
976 lambda: self._write_user_reset_updates_once(updated_users),
977 reason="reset_budget_write_users_failure",
978 )
980 async def _write_user_reset_updates_once(self, updated_users: Sequence[_RowReset[LiteLLM_UserTable]]) -> None:
981 async with spend_reset_unit_of_work(self._new_batch) as uow:
982 for u in updated_users:
983 uow.users.queue_spend_reset(
984 user_id=u.row.user_id,
985 budget_reset_at=u.row.budget_reset_at,
986 spend_decrement=u.spend_decrement,
987 )
989 async def _write_team_reset_updates(self, updated_teams: Sequence[_RowReset[LiteLLM_TeamTable]]) -> None:
990 """
991 Write per-row {spend, budget_reset_at} updates for teams.
993 Mirrors _write_key_reset_updates — avoids the full-model update path
994 that trips Prisma's DataError on rows carrying unrecognised fields
995 (see #27730).
996 """
997 await self._with_db_write_retry(
998 lambda: self._write_team_reset_updates_once(updated_teams),
999 reason="reset_budget_write_teams_failure",
1000 )
1002 async def _write_team_reset_updates_once(self, updated_teams: Sequence[_RowReset[LiteLLM_TeamTable]]) -> None:
1003 async with spend_reset_unit_of_work(self._new_batch) as uow:
1004 for t in updated_teams:
1005 uow.teams.queue_spend_reset(
1006 team_id=t.row.team_id,
1007 budget_reset_at=t.row.budget_reset_at,
1008 spend_decrement=t.spend_decrement,
1009 )
1011 def _emit_phase_failure(
1012 self,
1013 call_type: str,
1014 error: Exception,
1015 start_time: float,
1016 end_time: float,
1017 event_metadata: dict[str, object],
1018 ) -> None:
1019 """Report rows that could not be reset without failing the chunk: the
1020 rows that did reset are already committed, and raising here would cost
1021 the phase every remaining chunk this tick.
1022 """
1023 verbose_proxy_logger.error("%s: %s", call_type, error)
1024 asyncio.create_task(
1025 self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
1026 service=ServiceTypes.RESET_BUDGET_JOB,
1027 duration=end_time - start_time,
1028 error=error,
1029 call_type=call_type,
1030 start_time=start_time,
1031 end_time=end_time,
1032 event_metadata=event_metadata,
1033 )
1034 )
1036 async def reset_budget_for_litellm_keys(self) -> None:
1037 """
1038 Resets the budget for all the litellm keys
1040 Catches Exceptions and logs them
1041 """
1042 await _run_phase_in_chunks(self._reset_budget_for_litellm_keys_chunk)
1044 async def _reset_budget_for_litellm_keys_chunk(self) -> _ChunkOutcome:
1045 now: Final = datetime.utcnow()
1046 start_time: Final = time.time()
1047 keys_to_reset: list[LiteLLM_VerificationToken] | None = None
1048 try:
1049 keys_to_reset = await self._with_db_retry(
1050 lambda: self.prisma_client.get_data(
1051 table_name="key",
1052 query_type="find_all",
1053 expires=now,
1054 reset_at=now,
1055 limit=RESET_BUDGET_JOB_BATCH_SIZE,
1056 ),
1057 reason="reset_budget_read_keys_failure",
1058 )
1059 verbose_proxy_logger.debug("Keys to reset %s", _LazyJson(keys_to_reset))
1060 updated_keys: Final[list[_RowReset[LiteLLM_VerificationToken]]] = []
1061 failed_keys: Final = []
1062 if keys_to_reset is not None and len(keys_to_reset) > 0: 1062 ↛ 1063line 1062 didn't jump to line 1063 because the condition on line 1062 was never true
1063 for key in keys_to_reset:
1064 try:
1065 pre_reset_spend = float(key.spend or 0.0)
1066 updated_key = await ResetBudgetJob._reset_budget_for_key(
1067 key=key,
1068 current_time=now,
1069 reset_settings=self.reset_settings,
1070 )
1071 if updated_key is not None:
1072 updated_keys.append(
1073 _RowReset(
1074 row=updated_key,
1075 spend_decrement=pre_reset_spend - float(updated_key.spend or 0.0),
1076 )
1077 )
1078 else:
1079 failed_keys.append({"key": key, "error": "Returned None without exception"})
1080 except Exception as e:
1081 failed_keys.append({"key": key, "error": str(e)})
1082 verbose_proxy_logger.exception("Failed to reset budget for key: %s", key)
1084 verbose_proxy_logger.debug("Updated keys %s", _LazyJson(updated_keys))
1086 if updated_keys:
1087 await self._write_key_reset_updates(updated_keys=updated_keys)
1088 for k in updated_keys:
1089 token = getattr(k.row, "token", None)
1090 if token:
1091 await self._invalidate_spend_counter(f"spend:key:{token}")
1093 end_time = time.time()
1094 outcome: Final = _ChunkOutcome(
1095 fetched=len(keys_to_reset) if keys_to_reset else 0,
1096 advanced=_count_advanced(
1097 (k.row.budget_reset_at for k in updated_keys),
1098 cutoff=datetime.now(timezone.utc),
1099 ),
1100 )
1101 if len(failed_keys) > 0: 1101 ↛ 1102line 1101 didn't jump to line 1102 because the condition on line 1101 was never true
1102 self._emit_phase_failure(
1103 call_type="reset_budget_keys",
1104 error=Exception(f"Failed to reset {len(failed_keys)} keys: {json.dumps(failed_keys, default=str)}"),
1105 start_time=start_time,
1106 end_time=end_time,
1107 event_metadata={
1108 "num_keys_found": len(keys_to_reset) if keys_to_reset else 0,
1109 },
1110 )
1111 return outcome
1113 asyncio.create_task(
1114 self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
1115 service=ServiceTypes.RESET_BUDGET_JOB,
1116 duration=end_time - start_time,
1117 call_type="reset_budget_keys",
1118 start_time=start_time,
1119 end_time=end_time,
1120 event_metadata={
1121 "num_keys_found": len(keys_to_reset) if keys_to_reset else 0,
1122 "num_keys_updated": len(updated_keys),
1123 "num_keys_failed": len(failed_keys),
1124 },
1125 )
1126 )
1127 except Exception as e:
1128 end_time = time.time()
1129 asyncio.create_task(
1130 self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
1131 service=ServiceTypes.RESET_BUDGET_JOB,
1132 duration=end_time - start_time,
1133 error=e,
1134 call_type="reset_budget_keys",
1135 start_time=start_time,
1136 end_time=end_time,
1137 event_metadata={
1138 "num_keys_found": len(keys_to_reset) if keys_to_reset else 0,
1139 },
1140 )
1141 )
1142 verbose_proxy_logger.exception("Failed to reset budget for keys: %s", e)
1143 return _NO_PROGRESS
1144 else:
1145 return outcome
1147 async def reset_budget_for_litellm_users(self) -> None:
1148 """
1149 Resets the budget for all LiteLLM Internal Users if their budget has expired
1150 """
1151 await _run_phase_in_chunks(self._reset_budget_for_litellm_users_chunk)
1153 async def _reset_budget_for_litellm_users_chunk(self) -> _ChunkOutcome:
1154 now: Final = datetime.utcnow()
1155 start_time: Final = time.time()
1156 users_to_reset: list[LiteLLM_UserTable] | None = None
1157 try:
1158 users_to_reset = await self._with_db_retry(
1159 lambda: self.prisma_client.get_data(
1160 table_name="user",
1161 query_type="find_all",
1162 reset_at=now,
1163 limit=RESET_BUDGET_JOB_BATCH_SIZE,
1164 ),
1165 reason="reset_budget_read_users_failure",
1166 )
1167 updated_users: Final[list[_RowReset[LiteLLM_UserTable]]] = []
1168 failed_users: Final = []
1169 if users_to_reset is not None and len(users_to_reset) > 0:
1170 for user in users_to_reset:
1171 try:
1172 pre_reset_spend = float(user.spend or 0.0)
1173 updated_user = await ResetBudgetJob._reset_budget_for_user(
1174 user=user,
1175 current_time=now,
1176 reset_settings=self.reset_settings,
1177 )
1178 if updated_user is not None: 1178 ↛ 1186line 1178 didn't jump to line 1186 because the condition on line 1178 was always true
1179 updated_users.append(
1180 _RowReset(
1181 row=updated_user,
1182 spend_decrement=pre_reset_spend - float(updated_user.spend or 0.0),
1183 )
1184 )
1185 else:
1186 failed_users.append(
1187 {
1188 "user": user,
1189 "error": "Returned None without exception",
1190 }
1191 )
1192 except Exception as e:
1193 failed_users.append({"user": user, "error": str(e)})
1194 verbose_proxy_logger.exception("Failed to reset budget for user: %s", user)
1196 verbose_proxy_logger.debug("Updated users %s", _LazyJson(updated_users))
1197 if updated_users: 1197 ↛ 1206line 1197 didn't jump to line 1206 because the condition on line 1197 was always true
1198 await self._write_user_reset_updates(updated_users=updated_users)
1199 for u in updated_users:
1200 user_id = getattr(u.row, "user_id", None)
1201 if user_id:
1202 await self._invalidate_spend_counter(f"spend:user:{user_id}")
1203 if user_id == LITELLM_PROXY_BUDGET_NAME:
1204 await self._invalidate_global_proxy_spend_cache()
1206 end_time = time.time()
1207 outcome: Final = _ChunkOutcome(
1208 fetched=len(users_to_reset) if users_to_reset else 0,
1209 advanced=_count_advanced(
1210 (u.row.budget_reset_at for u in updated_users),
1211 cutoff=datetime.now(timezone.utc),
1212 ),
1213 )
1214 if len(failed_users) > 0: 1214 ↛ 1215line 1214 didn't jump to line 1215 because the condition on line 1214 was never true
1215 self._emit_phase_failure(
1216 call_type="reset_budget_users",
1217 error=Exception(
1218 f"Failed to reset {len(failed_users)} users: {json.dumps(failed_users, default=str)}"
1219 ),
1220 start_time=start_time,
1221 end_time=end_time,
1222 event_metadata={
1223 "num_users_found": len(users_to_reset) if users_to_reset else 0,
1224 },
1225 )
1226 return outcome
1228 asyncio.create_task(
1229 self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
1230 service=ServiceTypes.RESET_BUDGET_JOB,
1231 duration=end_time - start_time,
1232 call_type="reset_budget_users",
1233 start_time=start_time,
1234 end_time=end_time,
1235 event_metadata={
1236 "num_users_found": len(users_to_reset) if users_to_reset else 0,
1237 "num_users_updated": len(updated_users),
1238 "num_users_failed": len(failed_users),
1239 },
1240 )
1241 )
1242 except Exception as e:
1243 end_time = time.time()
1244 asyncio.create_task(
1245 self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
1246 service=ServiceTypes.RESET_BUDGET_JOB,
1247 duration=end_time - start_time,
1248 error=e,
1249 call_type="reset_budget_users",
1250 start_time=start_time,
1251 end_time=end_time,
1252 event_metadata={
1253 "num_users_found": len(users_to_reset) if users_to_reset else 0,
1254 },
1255 )
1256 )
1257 verbose_proxy_logger.exception("Failed to reset budget for users: %s", e)
1258 return _NO_PROGRESS
1259 else:
1260 return outcome
1262 async def reset_budget_for_litellm_teams(self) -> None:
1263 """
1264 Resets the budget for all LiteLLM Internal Teams if their budget has expired
1265 """
1266 await _run_phase_in_chunks(self._reset_budget_for_litellm_teams_chunk)
1268 async def _reset_budget_for_litellm_teams_chunk(self) -> _ChunkOutcome:
1269 now: Final = datetime.utcnow()
1270 start_time: Final = time.time()
1271 teams_to_reset: list[LiteLLM_TeamTable] | None = None
1272 try:
1273 teams_to_reset = await self._with_db_retry(
1274 lambda: self.prisma_client.get_data(
1275 table_name="team",
1276 query_type="find_all",
1277 reset_at=now,
1278 limit=RESET_BUDGET_JOB_BATCH_SIZE,
1279 ),
1280 reason="reset_budget_read_teams_failure",
1281 )
1282 updated_teams: Final[list[_RowReset[LiteLLM_TeamTable]]] = []
1283 failed_teams: Final = []
1284 if teams_to_reset is not None and len(teams_to_reset) > 0: 1284 ↛ 1285line 1284 didn't jump to line 1285 because the condition on line 1284 was never true
1285 for team in teams_to_reset:
1286 try:
1287 pre_reset_spend = float(team.spend or 0.0)
1288 updated_team = await ResetBudgetJob._reset_budget_for_team(
1289 team=team,
1290 current_time=now,
1291 reset_settings=self.reset_settings,
1292 )
1293 if updated_team is not None:
1294 updated_teams.append(
1295 _RowReset(
1296 row=updated_team,
1297 spend_decrement=pre_reset_spend - float(updated_team.spend or 0.0),
1298 )
1299 )
1300 else:
1301 failed_teams.append(
1302 {
1303 "team": team,
1304 "error": "Returned None without exception",
1305 }
1306 )
1307 except Exception as e:
1308 failed_teams.append({"team": team, "error": str(e)})
1309 verbose_proxy_logger.exception("Failed to reset budget for team: %s", team)
1311 verbose_proxy_logger.debug("Updated teams %s", _LazyJson(updated_teams))
1312 if updated_teams:
1313 await self._write_team_reset_updates(updated_teams=updated_teams)
1314 for t in updated_teams:
1315 team_id = getattr(t.row, "team_id", None)
1316 if team_id:
1317 await self._invalidate_spend_counter(f"spend:team:{team_id}")
1319 end_time = time.time()
1320 outcome: Final = _ChunkOutcome(
1321 fetched=len(teams_to_reset) if teams_to_reset else 0,
1322 advanced=_count_advanced(
1323 (t.row.budget_reset_at for t in updated_teams),
1324 cutoff=datetime.now(timezone.utc),
1325 ),
1326 )
1327 if len(failed_teams) > 0: 1327 ↛ 1328line 1327 didn't jump to line 1328 because the condition on line 1327 was never true
1328 self._emit_phase_failure(
1329 call_type="reset_budget_teams",
1330 error=Exception(
1331 f"Failed to reset {len(failed_teams)} teams: {json.dumps(failed_teams, default=str)}"
1332 ),
1333 start_time=start_time,
1334 end_time=end_time,
1335 event_metadata={
1336 "num_teams_found": len(teams_to_reset) if teams_to_reset else 0,
1337 },
1338 )
1339 return outcome
1341 asyncio.create_task(
1342 self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
1343 service=ServiceTypes.RESET_BUDGET_JOB,
1344 duration=end_time - start_time,
1345 call_type="reset_budget_teams",
1346 start_time=start_time,
1347 end_time=end_time,
1348 event_metadata={
1349 "num_teams_found": len(teams_to_reset) if teams_to_reset else 0,
1350 "num_teams_updated": len(updated_teams),
1351 "num_teams_failed": len(failed_teams),
1352 },
1353 )
1354 )
1355 except Exception as e:
1356 end_time = time.time()
1357 asyncio.create_task(
1358 self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
1359 service=ServiceTypes.RESET_BUDGET_JOB,
1360 duration=end_time - start_time,
1361 error=e,
1362 call_type="reset_budget_teams",
1363 start_time=start_time,
1364 end_time=end_time,
1365 event_metadata={
1366 "num_teams_found": len(teams_to_reset) if teams_to_reset else 0,
1367 },
1368 )
1369 )
1370 verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e)
1371 return _NO_PROGRESS
1372 else:
1373 return outcome
1375 @staticmethod
1376 async def _reset_expired_window(
1377 window: dict,
1378 counter_key: str,
1379 spend_counter_cache: DualCache,
1380 now: datetime,
1381 reset_settings: BudgetResetSettings,
1382 prisma_client: PrismaClient,
1383 entity_type: Litellm_EntityType,
1384 entity_id: str,
1385 ) -> bool:
1386 """Reset a single budget window if expired. Returns True if the window was reset."""
1387 reset_at_str: Final = window.get("reset_at")
1388 if not reset_at_str: 1388 ↛ 1389line 1388 didn't jump to line 1389 because the condition on line 1388 was never true
1389 return False
1390 reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None)
1391 if reset_at > now: 1391 ↛ 1393line 1391 didn't jump to line 1393 because the condition on line 1391 was always true
1392 return False
1393 new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_cache)
1394 spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_value)
1395 if spend_counter_cache.redis_cache is not None:
1396 try:
1397 await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_value)
1398 except Exception as redis_err:
1399 verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err)
1400 budget_duration: Final = window["budget_duration"]
1401 next_reset_at: Final = compute_budget_reset_at(budget_duration=budget_duration, settings=reset_settings)
1402 window["reset_at"] = next_reset_at.isoformat()
1403 await ResetBudgetJob._roll_window_spend_row(
1404 prisma_client=prisma_client,
1405 entity_type=entity_type,
1406 entity_id=entity_id,
1407 budget_duration=budget_duration,
1408 next_reset_at=next_reset_at,
1409 )
1410 return True
1412 @staticmethod
1413 async def _roll_window_spend_row(
1414 prisma_client: PrismaClient,
1415 entity_type: Litellm_EntityType,
1416 entity_id: str,
1417 budget_duration: str,
1418 next_reset_at: datetime,
1419 ) -> None:
1420 """Move this window's LiteLLM_BudgetWindowSpend row onto the window
1421 that just started, so the maintained total the read path uses starts
1422 from zero alongside the counter.
1424 Best effort: the row is an optimization over aggregating
1425 LiteLLM_SpendLogs, so a failure here must not stop the remaining
1426 windows from having their counters reset.
1427 """
1428 try:
1429 window_start: Final = next_reset_at - timedelta(seconds=duration_in_seconds(budget_duration))
1430 except Exception as e: # noqa: BLE001 # duration_in_seconds raises bare exceptions on bad input
1431 verbose_proxy_logger.warning("Unparseable budget_duration %s: %s", budget_duration, e)
1432 return
1433 try:
1434 await roll_window_spend_row(
1435 prisma_client=prisma_client,
1436 entity_type=entity_type.value,
1437 entity_id=entity_id,
1438 window_duration=budget_duration,
1439 new_window_start=window_start,
1440 )
1441 except Exception as e: # noqa: BLE001 # the row is best effort; counter resets must still land
1442 verbose_proxy_logger.warning(
1443 "Failed to roll budget window spend row for %s=%s window=%s: %s",
1444 entity_type.value,
1445 entity_id,
1446 budget_duration,
1447 e,
1448 )
1450 @staticmethod
1451 async def _window_carried_spend(
1452 window: Mapping[str, object], counter_key: str, spend_counter_cache: DualCache
1453 ) -> float:
1454 """Per-window spend lives only in the counter, so the carried overage is
1455 read from it before the reset overwrites it."""
1456 if not _rollover_enabled():
1457 return 0.0
1458 window_max: Final = window.get("max_budget")
1459 cap: Final = _rollover_cap(window_max) if isinstance(window_max, (int, float)) else None
1460 if cap is None:
1461 return 0.0
1462 try:
1463 current: Final = await spend_counter_cache.async_get_cache(key=counter_key)
1464 except Exception as e: # noqa: BLE001 # an unreadable counter falls back to a plain zero reset
1465 verbose_proxy_logger.warning("Failed to read spend counter %s for rollover: %s", counter_key, e)
1466 return 0.0
1467 if not isinstance(current, (int, float)):
1468 return 0.0
1469 return _carried_spend(float(current), cap)
1471 async def reset_budget_windows(self) -> None:
1472 """
1473 For keys and teams with budget_limits, reset any individual windows where
1474 reset_at <= now. Only the expired windows are reset; other windows are untouched.
1475 """
1477 from litellm.proxy.proxy_server import spend_counter_cache
1479 now: Final = datetime.utcnow()
1480 for source in _WINDOW_SOURCES:
1481 try:
1482 await self._reset_windows_for(source=source, now=now, spend_counter_cache=spend_counter_cache)
1483 except Exception as e:
1484 verbose_proxy_logger.exception("Failed to reset budget windows for %s: %s", source.log_subject, e)
1486 async def _reset_windows_for(
1487 self,
1488 source: _WindowSource,
1489 now: datetime,
1490 spend_counter_cache: DualCache,
1491 ) -> None:
1492 """Walk one table's windowed rows a page at a time, to the end.
1494 Paging is what bounds the memory: the previous form pulled every row
1495 carrying budget_limits into one result set on every tick, which grows
1496 with the deployment's key count and is paid on the event loop.
1498 The walk deliberately has no per-run page cap. A cap has to remember
1499 where it stopped, and that position cannot live in the process: the
1500 lease is released after each sweep, so the next tick can elect a
1501 different pod whose own position is unset. It would restart at the first
1502 row and never reach the tail, pinning those windows at their cap for
1503 good. The cursor strictly advances, so the walk terminates on its own
1504 without needing a bound.
1505 """
1506 cursor = ""
1507 while True:
1508 next_cursor = await self._reset_window_page(
1509 source=source,
1510 cursor=cursor,
1511 now=now,
1512 spend_counter_cache=spend_counter_cache,
1513 )
1514 if next_cursor is None: 1514 ↛ 1516line 1514 didn't jump to line 1516 because the condition on line 1514 was always true
1515 return
1516 cursor = next_cursor
1518 async def _reset_window_page(
1519 self,
1520 source: _WindowSource,
1521 cursor: str,
1522 now: datetime,
1523 spend_counter_cache: DualCache,
1524 ) -> str | None:
1525 """Reset one page of windows; return the next cursor, or None when drained."""
1526 rows: Final = await self._with_db_retry(
1527 lambda: self.prisma_client.db.query_raw(source.page_query(), cursor, RESET_BUDGET_JOB_BATCH_SIZE),
1528 reason=f"reset_budget_read_{source.retry_subject}_windows_failure",
1529 )
1530 for row in rows:
1531 raw = row["budget_limits"]
1532 if not raw:
1533 continue
1534 row_id: str = row[source.id_column]
1535 windows: list[dict[str, object]] = raw if isinstance(raw, list) else json.loads(raw)
1536 changed = False
1537 for window in windows:
1538 counter_key = f"{source.counter_prefix}:{row_id}:window:{window['budget_duration']}"
1539 if await ResetBudgetJob._reset_expired_window( 1539 ↛ 1549line 1539 didn't jump to line 1549 because the condition on line 1539 was never true
1540 window,
1541 counter_key,
1542 spend_counter_cache,
1543 now,
1544 self.reset_settings,
1545 prisma_client=self.prisma_client,
1546 entity_type=source.entity_type,
1547 entity_id=row_id,
1548 ):
1549 changed = True
1550 if changed: 1550 ↛ 1551line 1550 didn't jump to line 1551 because the condition on line 1550 was never true
1551 await self._with_db_write_retry(
1552 lambda: source.write(self.prisma_client, row_id, json.dumps(windows)),
1553 reason=f"reset_budget_write_{source.retry_subject}_windows_failure",
1554 )
1556 if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: 1556 ↛ 1558line 1556 didn't jump to line 1558 because the condition on line 1556 was always true
1557 return None
1558 return rows[-1][source.id_column]
1560 @staticmethod
1561 async def _reset_budget_common(
1562 item: LiteLLM_TeamTable | LiteLLM_UserTable | LiteLLM_VerificationToken,
1563 current_time: datetime,
1564 item_type: Literal["key", "team", "user"],
1565 reset_settings: BudgetResetSettings,
1566 ):
1567 """
1568 In-place, updates spend=0, and sets budget_reset_at to current_time + budget_duration
1570 Common logic for resetting budget for a team, user, or key.
1572 Spend-counter invalidation happens in the caller, AFTER the DB write
1573 commits. Zeroing the counter here would open a bypass window when the
1574 DB write fails: get_current_spend reads 0 from Redis while the DB
1575 still holds the pre-reset value, admitting requests past the cap.
1576 """
1577 try:
1578 item.spend = _carried_spend(item.spend, _rollover_cap(item.max_budget)) if _rollover_enabled() else 0.0
1579 if hasattr(item, "budget_duration") and item.budget_duration is not None: 1579 ↛ 1583line 1579 didn't jump to line 1583 because the condition on line 1579 was always true
1580 item.budget_reset_at = compute_budget_reset_at(
1581 budget_duration=item.budget_duration, settings=reset_settings
1582 )
1583 return item
1584 except Exception as e:
1585 verbose_proxy_logger.exception("Error resetting budget for %s: %s. Item: %s", item_type, e, item)
1586 raise e
1588 @staticmethod
1589 async def _reset_budget_for_team(
1590 team: LiteLLM_TeamTable,
1591 current_time: datetime,
1592 reset_settings: BudgetResetSettings,
1593 ) -> LiteLLM_TeamTable | None:
1594 await ResetBudgetJob._reset_budget_common(
1595 item=team,
1596 current_time=current_time,
1597 item_type="team",
1598 reset_settings=reset_settings,
1599 )
1600 return team
1602 @staticmethod
1603 async def _reset_budget_for_user(
1604 user: LiteLLM_UserTable,
1605 current_time: datetime,
1606 reset_settings: BudgetResetSettings,
1607 ) -> LiteLLM_UserTable | None:
1608 await ResetBudgetJob._reset_budget_common(
1609 item=user,
1610 current_time=current_time,
1611 item_type="user",
1612 reset_settings=reset_settings,
1613 )
1614 return user
1616 @staticmethod
1617 async def _reset_budget_for_key(
1618 key: LiteLLM_VerificationToken,
1619 current_time: datetime,
1620 reset_settings: BudgetResetSettings,
1621 ) -> LiteLLM_VerificationToken | None:
1622 await ResetBudgetJob._reset_budget_common(
1623 item=key,
1624 current_time=current_time,
1625 item_type="key",
1626 reset_settings=reset_settings,
1627 )
1628 return key