scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(proxy): refund batch TPD reservation on failure and report active window reset time

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-13 10:56:17 +00:00
parent c47120cbf7
commit a41b719920
4 changed files with 191 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
litellm/proxy/hooks/batch_rate_limiter.py
View file

@ -18,12 +18,13 @@ Quick summary:
"""
import json
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Callable, Iterable, Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, TypeAlias
from fastapi import HTTPException
from pydantic import BaseModel, Field, TypeAdapter
from pydantic import BaseModel, Field, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -56,6 +57,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROJECT_ITPM_DESCRIPTOR_KEY,
PROJECT_OTPM_DESCRIPTOR_KEY,
ReservationAwareIncrementOperation,
get_or_create_request_stash,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
@ -93,6 +95,7 @@ else:
_BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object])
_WINDOW_START_ADAPTER: Final[TypeAdapter[int | float | str | None]] = TypeAdapter(int | float | str | None)
IncrementAmounts: TypeAlias = dict[Literal["requests", "tokens"], int]
@ -129,6 +132,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
self,
internal_usage_cache: InternalUsageCache,
parallel_request_limiter: ParallelRequestLimiter,
time_provider: Callable[[], datetime] | None = None,
):
"""
Initialize the batch rate limiter.
@ -139,9 +143,11 @@ class _PROXY_BatchRateLimiter(CustomLogger):
Args:
internal_usage_cache: Cache for storing rate limit data (auto-injected)
parallel_request_limiter: Existing rate limiter to integrate with (needs custom injection)
time_provider: Clock used for rate limit reset times (defaults to ``datetime.now``)
"""
self.internal_usage_cache = internal_usage_cache
self.parallel_request_limiter = parallel_request_limiter
self._time_provider: Final = time_provider or datetime.now
self._warned_unsupported_model_skip = False
def _get_file_bound_batch_model(self, data: dict) -> str | None:
@ -618,9 +624,14 @@ class _PROXY_BatchRateLimiter(CustomLogger):
batch_usage: BatchFileUsage,
limit_type: str,
requested_model: str | None = None,
window_start: int | None = None,
) -> NoReturn:
"""Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded."""
from datetime import datetime
"""Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded.
``window_start`` is the active counter window's start (unix seconds) when
known, so the reset time reflects that window's actual end rather than a
full window from now.
"""
# Find the descriptor for this status. Matching on (key, value) is
# required, not key alone: a batch can carry several project ITPM/OTPM
@ -644,11 +655,12 @@ class _PROXY_BatchRateLimiter(CustomLogger):
descriptors[descriptor_index] if descriptors else {"key": "", "value": "", "rate_limit": None}
)
now: Final = datetime.now().timestamp()
now: Final = self._time_provider().timestamp()
window_size: Final = (descriptor.get("rate_limit") or {}).get(
"window_size"
) or self.parallel_request_limiter.window_size
reset_time: Final = now + window_size
reset_time: Final = now + window_size if window_start is None else window_start + window_size
retry_after: Final = max(0, int(reset_time - now))
reset_time_formatted: Final = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC")
remaining_display: Final = max(0, status["limit_remaining"])
@ -694,7 +706,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
raise ProxyRateLimitError(
detail=detail,
headers={
"retry-after": str(window_size),
"retry-after": str(retry_after),
"rate_limit_type": limit_type,
"reset_at": reset_time_formatted,
},
@ -752,6 +764,8 @@ class _PROXY_BatchRateLimiter(CustomLogger):
parent_otel_span=user_api_key_dict.parent_otel_span,
)
stash: Final = get_or_create_request_stash()
stash.batch_tpd_refund_ops = ()
if rate_limit_response["overall_code"] == "OVER_LIMIT":
requested_model: Final = data.get("model") if data else None
for status in rate_limit_response["statuses"]:
@ -762,8 +776,70 @@ class _PROXY_BatchRateLimiter(CustomLogger):
batch_usage,
status["rate_limit_type"],
requested_model=requested_model,
window_start=await self._read_tpd_window_start(
status=status, parent_otel_span=user_api_key_dict.parent_otel_span
),
)
stash.batch_tpd_refund_ops = self._build_tpd_refund_ops(
descriptors=descriptors,
tokens=batch_usage.total_tokens,
reservation_windows=rate_limit_response.get("reservation_windows", frozenset()),
)
async def _read_tpd_window_start(self, status: "RateLimitStatus", parent_otel_span: "Span | None") -> int | None:
descriptor_key: Final = status.get("descriptor_key") or ""
if not descriptor_key.endswith(BATCH_TPD_DESCRIPTOR_SUFFIX):
return None
try:
window_start: Final = _WINDOW_START_ADAPTER.validate_python(
await self.parallel_request_limiter.internal_usage_cache.async_get_cache(
key=f"{{{descriptor_key}:{status.get('descriptor_value') or ''}}}:window",
litellm_parent_otel_span=parent_otel_span,
),
strict=True,
)
return None if window_start is None else int(float(window_start))
except (ValidationError, ValueError):
return None
def _build_tpd_refund_ops(
self,
descriptors: Sequence["RateLimitDescriptor"],
tokens: int,
reservation_windows: frozenset[tuple[str, str, Literal["redis", "local"]]],
) -> tuple[ReservationAwareIncrementOperation, ...]:
"""Refund operations for the daily token counters this batch charged.
The v3 limiter's failure hook applies them when the submission fails
after the counters were incremented. Each operation carries the window
identity the charge landed in, so the refund is skipped once that
window has rolled over.
"""
if tokens <= 0 or not reservation_windows:
return ()
tpd_descriptors_by_counter: Final[Mapping[str, RateLimitDescriptor]] = MappingProxyType(
{
self.parallel_request_limiter.create_rate_limit_keys(
descriptor["key"], descriptor["value"], "tokens"
): descriptor
for descriptor in descriptors
if descriptor["key"].endswith(BATCH_TPD_DESCRIPTOR_SUFFIX)
}
)
return tuple(
ReservationAwareIncrementOperation(
key=counter_key,
increment_value=-tokens,
ttl=BATCH_TPD_WINDOW_SECONDS,
window_key=f"{{{descriptor['key']}:{descriptor['value']}}}:window",
expected_window_start=window_start,
reservation_backend=backend,
)
for counter_key, window_start, backend in sorted(reservation_windows)
if (descriptor := tpd_descriptors_by_counter.get(counter_key)) is not None
)
async def count_input_file_usage(
self,
file_id: str,

15
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -390,6 +390,8 @@ CacheCounterValue: TypeAlias = int | float | str | bytes
CacheCounterValues: TypeAlias = Sequence[CacheCounterValue | None]
ReservationWindowIdentity: TypeAlias = tuple[str, str, Literal["redis", "local"]]
ParallelGaugeCacheValue: TypeAlias = dict[str, object] | int | float | str | bytes
@ -536,6 +538,7 @@ class RequestRateLimiterStash:
default_factory=frozenset
)
batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None
batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = ()
reservation_released: bool = False
@ -677,6 +680,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._batch_rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=self.internal_usage_cache,
parallel_request_limiter=self,
time_provider=self._time_provider,
)
except Exception as e:
verbose_proxy_logger.debug("Could not load batch rate limiter: %s", e)
@ -1817,6 +1821,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
applied: Final[list[list[AtomicCounterMeta]]] = []
statuses: Final[list[RateLimitStatus]] = []
reservation_windows: Final[set[ReservationWindowIdentity]] = set() # mutable-ok: filled by the group loop
raw: list[CacheCounterValue]
for _idx, (keys, args, meta) in enumerate(descriptor_groups):
@ -1854,11 +1859,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return response
applied.append(meta)
statuses.extend(response["statuses"])
reservation_windows.update(response.get("reservation_windows", frozenset()))
return RateLimitResponse(
overall_code="OK",
statuses=statuses,
reservation_windows=frozenset(),
reservation_windows=frozenset(reservation_windows),
)
async def _refund_applied_descriptor_groups(
@ -4788,6 +4794,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
stash.batch_enqueued_reservation = None
if stash.batch_tpd_refund_ops:
await self.async_increment_reservation_aware_tokens(
pipeline_operations=stash.batch_tpd_refund_ops,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
stash.batch_tpd_refund_ops = ()
if stash.reservation_released:
return
reserved_tokens: Final = stash.reserved_tokens

98
tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py
View file

@ -6,6 +6,8 @@ batch under a per-minute RPM/TPM budget. Scopes that configure `tpd_limit`
are charged against a 24h token window instead of their minute counters.
"""
from datetime import datetime
import pytest
from fastapi import HTTPException
@ -19,9 +21,17 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
from litellm.proxy.utils import InternalUsageCache, hash_token
def _make_limiters():
class _Clock:
def __init__(self, start: datetime):
self.now = start
def __call__(self) -> datetime:
return self.now
def _make_limiters(clock: _Clock | None = None):
internal_usage_cache = InternalUsageCache(dual_cache=DualCache())
rate_limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache)
rate_limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache, time_provider=clock)
batch_limiter = rate_limiter._get_batch_rate_limiter()
assert batch_limiter is not None
return internal_usage_cache, rate_limiter, batch_limiter
@ -51,8 +61,10 @@ async def test_batch_over_rpm_and_tpm_but_under_tpd_is_accepted():
@pytest.mark.asyncio
async def test_cumulative_batch_tokens_over_tpd_returns_429_with_daily_reset():
_internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters()
async def test_cumulative_batch_tokens_over_tpd_returns_429_with_remaining_daily_window():
window_start = datetime(2026, 9, 13, 8, 0, 0)
clock = _Clock(window_start)
_internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters(clock)
user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("tpd-key-2"), rpm_limit=1, tpd_limit=1000)
await batch_limiter._check_and_increment_batch_counters(
@ -60,6 +72,7 @@ async def test_cumulative_batch_tokens_over_tpd_returns_429_with_daily_reset():
data={},
batch_usage=BatchFileUsage(total_tokens=600, request_count=6),
)
clock.now = datetime(2026, 9, 13, 11, 0, 0)
with pytest.raises(HTTPException) as exc:
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
@ -70,7 +83,82 @@ async def test_cumulative_batch_tokens_over_tpd_returns_429_with_daily_reset():
assert exc.value.status_code == 429
assert "api_key_tpd" in str(exc.value.detail)
assert "600 tokens but only 400 tokens remaining out of 1000 TPD limit" in str(exc.value.detail)
assert exc.value.headers["retry-after"] == str(BATCH_TPD_WINDOW_SECONDS)
assert exc.value.headers["retry-after"] == str(BATCH_TPD_WINDOW_SECONDS - 3 * 3600)
assert exc.value.headers["reset_at"] == "2026-09-14 08:00:00 UTC"
@pytest.mark.asyncio
async def test_failed_batch_submission_refunds_tpd_tokens():
internal_usage_cache, rate_limiter, batch_limiter = _make_limiters()
api_key = hash_token("tpd-refund-key")
user_api_key_dict = UserAPIKeyAuth(api_key=api_key, rpm_limit=1, tpd_limit=1000)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=600, request_count=6),
)
await rate_limiter.async_post_call_failure_hook(
request_data={},
original_exception=RuntimeError("provider rejected the file"),
user_api_key_dict=user_api_key_dict,
)
assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 0
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=1000, request_count=10),
)
assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 1000
@pytest.mark.asyncio
async def test_tpd_refund_applies_once_and_only_to_daily_counters():
internal_usage_cache, rate_limiter, batch_limiter = _make_limiters()
team_key = UserAPIKeyAuth(
api_key=hash_token("tpd-refund-team-key"),
rpm_limit=100,
tpm_limit=10_000,
team_id="team-r",
team_tpd_limit=5000,
)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=team_key,
data={},
batch_usage=BatchFileUsage(total_tokens=800, request_count=8),
)
await rate_limiter.async_post_call_failure_hook(
request_data={}, original_exception=RuntimeError("boom"), user_api_key_dict=team_key
)
await rate_limiter.async_post_call_failure_hook(
request_data={}, original_exception=RuntimeError("boom"), user_api_key_dict=team_key
)
assert await _counter(internal_usage_cache, rate_limiter, "team_tpd", "team-r", "tokens") == 0
assert await _counter(internal_usage_cache, rate_limiter, "api_key", team_key.api_key, "tokens") == 800
assert await _counter(internal_usage_cache, rate_limiter, "api_key", team_key.api_key, "requests") == 8
@pytest.mark.asyncio
async def test_rejected_batch_leaves_nothing_to_refund():
internal_usage_cache, rate_limiter, batch_limiter = _make_limiters()
api_key = hash_token("tpd-rejected-key")
user_api_key_dict = UserAPIKeyAuth(api_key=api_key, tpd_limit=100)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9)
)
with pytest.raises(HTTPException):
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict, data={}, batch_usage=BatchFileUsage(total_tokens=20, request_count=2)
)
await rate_limiter.async_post_call_failure_hook(
request_data={}, original_exception=RuntimeError("429 bubbled up"), user_api_key_dict=user_api_key_dict
)
assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 90
@pytest.mark.asyncio

1
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -10484,6 +10484,7 @@ export interface paths {
* - max_budget: *Optional[float]* - Max budget for org
* - tpm_limit: *Optional[int]* - Max tpm limit for org
* - rpm_limit: *Optional[int]* - Max rpm limit for org
* - tpd_limit: *Optional[int]* - Max tokens per day stored on the org budget. Batch submissions enforce tpd_limit at the key, team and end user scopes only.
* - model_rpm_limit: *Optional[Dict[str, int]]* - The RPM (Requests Per Minute) limit per model for this organization.
* - model_tpm_limit: *Optional[Dict[str, int]]* - The TPM (Tokens Per Minute) limit per model for this organization.
* - max_parallel_requests: *Optional[int]* - [Not Implemented Yet] Max parallel requests for org