From 1fdb6e95a54e40611e957521813ab8100ea68fab Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Mon, 14 Sep 2026 08:20:04 -0400 Subject: [PATCH] fix(proxy): break reset_budget_job cyclic-import risk, tighten comments per repo convention Addresses CodeQL and greptile findings on the reset-race fix. CodeQL flagged reset_budget_job.py's module-level imports of RedisCache, the two RESET_BUDGET_SPEND_COUNTER_RESET_* constants, and evict_and_broadcast as cyclic-import risks: each target module can reach back to this one before finishing its own initialization. RedisCache is only used as a type hint, so it moves under TYPE_CHECKING with the annotations quoted; the constants and evict_and_broadcast are only used inside single functions, so they move to local imports there, per this repo's own stated exception for avoiding circular imports. Also fixes a CodeQL "empty except with no explanatory comment" on the cancel-path's CancelledError handler, and trims the docstrings this fix added or touched down to this repo's comment conventions (concise, non-obvious-why only). --- litellm/caching/redis_cache.py | 17 ++--- .../proxy/common_utils/reset_budget_job.py | 76 +++++++------------ .../spend_tracking/budget_reservation.py | 63 +++++---------- 3 files changed, 49 insertions(+), 107 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 61fef745991..19e904e60e9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1411,18 +1411,11 @@ class RedisCache(BaseCache): snapshot: float, ttl: int | None = None, ) -> float: - """Atomically reset ``key`` to ``new_base`` while preserving any amount - added since ``snapshot`` was read, so a reset racing a concurrent - ``async_increment`` cannot erase spend reserved after the reset boundary. - - ``snapshot`` is the value read once before the first attempt and held - fixed across retries by the caller, not re-read each attempt: replaying - this call after a transient failure stays correct no matter how many - increments landed in between, because the delta is always measured - against that same original baseline. The GET/compute/SET runs in a - single Lua call, atomic across racing callers and pods, mirroring - ``async_set_max``. Returns the resulting value. - """ + """Reset ``key`` to ``new_base`` plus whatever was added on top of + ``snapshot``, atomically (single Lua GET/compute/SET, mirrors + ``async_set_max``), so a concurrent ``async_increment`` survives a reset + retry instead of being overwritten by it. ``snapshot`` is read once by + the caller and held fixed across retries.""" _redis_client: Final = self.init_async_client() _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 25caa948247..f7edba6fd1d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -7,14 +7,13 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeVar +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeVar from typing_extensions import assert_never import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCache from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME, @@ -22,8 +21,6 @@ from litellm.constants import ( RESET_BUDGET_JOB_LOCK_TTL_SECONDS, RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN, RESET_BUDGET_JOB_NAME, - RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS, - RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( @@ -35,9 +32,6 @@ from litellm.proxy._types import ( LiteLLM_UserTable, LiteLLM_VerificationToken, ) -from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( - evict_and_broadcast, -) from litellm.proxy.common_utils.timezone_utils import ( BudgetResetSettings, compute_budget_reset_at, @@ -72,6 +66,9 @@ from litellm.repositories.verification_token_repository import ( ) from litellm.types.services import ServiceTypes +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + _RowT = TypeVar("_RowT") _LINKED_KEYS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"budget_duration": None, "spend": {"gt": 0}}) @@ -537,31 +534,17 @@ class ResetBudgetJob: @staticmethod async def _invalidate_spend_counter(counter_key: str, new_spend: float = 0.0) -> None: - """Overwrite a spend counter with the post-reset value (0, or the carried - overage when budget rollover is enabled) so a DB-row reset takes effect - immediately. + """Overwrite a spend counter with the post-reset value. Call AFTER the DB + write commits, or get_current_spend can read 0 from Redis while the DB + still holds the pre-reset value. - Call AFTER the DB write commits. Clearing Redis before the DB - commit opens a window where get_current_spend reads 0 from Redis - while the DB still holds the pre-reset value, allowing bypass. - - The reset is delta-preserving: a snapshot of the pre-reset value is read - once up front and held fixed across every retry, so a concurrent - async_increment landing during a retry (e.g. a request reserving spend - against the just-reset budget) is carried forward on top of new_spend - instead of being erased by a later attempt's write. See - async_reset_preserving_delta and _reset_redis_spend_counter. - - A reset that keeps failing after retrying falls back to deleting the key rather - than leaving the pre-reset (possibly far higher) value authoritative in Redis - until its TTL expires: a missing counter reads as cold and reseeds from the - DB on the next request (_ensure_spend_counter_initialized), which is always - closer to the truth than the stale value a failed reset would otherwise leave - behind. The same is true when the pre-reset snapshot itself cannot be read: - with no safe baseline to preserve increments against, deleting is the only - option that cannot silently erase a concurrent reservation. + Delta-preserving: holds one pre-reset snapshot fixed across retries so a + concurrent async_increment during a retry lands on top of new_spend + instead of being overwritten by it. Falls back to delete (reads as cold, + reseeds from DB) if every retry fails or the snapshot itself can't be read. """ try: + from litellm.constants import RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS from litellm.proxy.proxy_server import spend_counter_cache spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_spend, ttl=60) @@ -595,10 +578,9 @@ class ResetBudgetJob: verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e) @staticmethod - async def _snapshot_spend_counter(redis_cache: RedisCache, counter_key: str) -> float | None: - """Read the pre-reset baseline the delta-preserving reset holds fixed across - retries. A failed read leaves no safe baseline to reconcile against, so the - caller skips the atomic reset and falls straight back to delete.""" + async def _snapshot_spend_counter(redis_cache: "RedisCache", counter_key: str) -> float | None: + """Pre-reset baseline held fixed across retries. A failed read means no + safe baseline, so the caller falls straight back to delete.""" try: current = await redis_cache.async_get_cache(key=counter_key) return float(current) if current is not None else 0.0 @@ -613,8 +595,13 @@ class ResetBudgetJob: @staticmethod async def _reset_redis_spend_counter( - redis_cache: RedisCache, counter_key: str, new_spend: float, snapshot: float + redis_cache: "RedisCache", counter_key: str, new_spend: float, snapshot: float ) -> bool: + from litellm.constants import ( + RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS, + RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS, + ) + for attempt in range(RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS): try: await redis_cache.async_reset_preserving_delta( @@ -644,22 +631,13 @@ class ResetBudgetJob: @staticmethod async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None: - """Drop a stale management-cache entry, on every pod, so the next read - fetches from DB. - - Tags and end-users are not reseeded by ``SpendCounterReseed.from_db``; - for those, when the spend counter expires the budget check falls back - to ``cached_obj.spend``. Keys, orgs, and team memberships are reseeded - from the DB, but auth still may consult ``user_api_key_cache`` objects - whose ``.spend`` field can lag a cross-pod DB reset. - - Only one pod runs this job per tick (see ``_acquire_lease``), so a local-only - delete would leave every other pod's copy stale until its TTL: ``evict_and_broadcast`` - is the same LIT-3803 cross-pod eviction every other cache-mutating endpoint already - uses (e.g. auth_checks.delete_cache_team_object), broadcasting the delete over Redis - pub/sub so every pod's copy is dropped, not just the one that ran the reset. - """ + """Drop a stale management-cache entry on every pod (LIT-3803 cross-pod + eviction), so a pod other than the one that ran the reset doesn't keep + serving a cached object with the pre-reset spend.""" try: + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + ) from litellm.proxy.proxy_server import user_api_key_cache await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 93cba60ab21..9631eb216d6 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -54,12 +54,8 @@ class _BudgetCounter: _RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS: Final = 2 -"""Bounded retries of the per-reservation reconcile on the cancel path before -falling back to invalidating the shared aggregate counters. reconcile_budget_reservation -only ever adjusts this reservation's own recorded contribution to each counter (see -_set_reserved_entries_actual_cost), so it is safe and idempotent to retry; the counter -invalidation fallback below is not, since it deletes state every concurrent reservation -and recorded spend on that counter shares.""" +"""Retries of the per-reservation reconcile before falling back to the +counter-deleting invalidation, which a concurrent reservation also shares.""" _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { @@ -360,37 +356,19 @@ async def release_budget_reservation(budget_reservation: dict | None) -> None: async def release_budget_reservation_on_cancel( - budget_reservation: dict | None, + budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict ) -> None: - """Reconcile a still-open reservation when the request is cancelled mid-flight. + """Reconcile a still-open reservation when the request is cancelled mid-flight, + since neither the success nor failure hook runs in that case. Reconciles to the + request's input-token cost, not zero, since that portion was already billed to + the provider. asyncio.shield keeps this running through the surrounding + cancellation. - A client disconnect or timeout cancels the request task, which surfaces as - CancelledError / GeneratorExit rather than a normal exception, so neither the - success cost callback nor the failure hook runs and the pre-call reservation - is never reconciled. Left alone it pins the spend counter above real spend - and 429s subsequent requests until the counter's TTL expires. - - Reconcile to the request's input-token cost rather than refunding to zero: - by the time a request is cancelled in-flight the provider call was already - dispatched, so the input tokens were billed even if no chunk reached the - client. Refunding to zero would let a caller abort pre-token to dodge that - charge; the worst-case output portion of the reservation is still released. - - asyncio.shield keeps the reconcile running to completion even though the - surrounding task is being cancelled. The `finalized` guard makes this a no-op - when success/failure handling already reconciled, so calling it on every - cancellation path is safe. - - A reconcile failure here (e.g. a Redis timeout) is retried a bounded number of - times before falling back to invalidate_budget_reservation_counters: unlike that - fallback, the reconcile only ever adjusts this reservation's own contribution to - each counter, so it cannot clobber a concurrent request's reservation or recorded - spend on the same key/user/team counter the way deleting the aggregate can. Only - once every retry hits the same failure does this fall back to - invalidate_budget_reservation_counters, mirroring release_or_invalidate_budget_reservation's - handling of the same failure on the non-cancel release path: dropping the reserved counters - forces the next read to reseed from the DB instead of leaving the pre-charge stuck in Redis - with nothing left to correct it, since this is the terminal handler for a cancelled request. + A reconcile failure is retried before falling back to + invalidate_budget_reservation_counters, since unlike that fallback, reconcile + only ever adjusts this reservation's own contribution to each counter and can't + clobber a concurrent reservation sharing it. Mirrors + release_or_invalidate_budget_reservation's fallback on the non-cancel path. """ if not budget_reservation or budget_reservation.get("finalized") is True: return @@ -400,7 +378,7 @@ async def release_budget_reservation_on_cancel( reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost) ) except asyncio.CancelledError: - pass + pass # a second cancellation while shielded; the reconcile keeps running detached regardless except Exception: verbose_proxy_logger.exception( "Failed to reconcile budget reservation on cancel; retrying before invalidating reserved counters" @@ -423,16 +401,9 @@ async def _retry_reconcile_reservation_on_cancel( budget_reservation: dict, # mutable-ok: reconcile_budget_reservation stamps finalized on success incurred_cost: float, ) -> bool: - """Bounded retry of the reconcile that just failed once on the cancel path. - - Unlike invalidate_budget_reservation_counters, reconcile_budget_reservation only - ever adjusts this reservation's own recorded contribution to each counter (see - _set_reserved_entries_actual_cost's applied_adjustment bookkeeping), so retrying it - is safe and idempotent, and correct on any attempt that gets through: it does not - touch concurrent reservations or spend the counter also aggregates. Returns whether - a retry succeeded, so the caller falls back to the destructive invalidation only - after every retry hits the same (presumably persistent) failure. - """ + """Retry the reconcile that just failed once. Safe to retry since it only + ever adjusts this reservation's own contribution, unlike the destructive + invalidation the caller falls back to once every retry fails.""" for attempt in range(_RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS): try: await reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost)