scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

fix(proxy): break reset_budget_job cyclic-import risk, tighten comments per repo convention

Browse source 
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).
This commit is contained in:
Abhyuday 2026-09-14 08:20:04 -04:00
parent bdd409430e
commit 1fdb6e95a5
3 changed files with 49 additions and 107 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/caching/redis_cache.py
View file

@ -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)

76
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -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)

63
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -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)