mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(proxy): release completed max-parallel slots promptly
This commit is contained in:
Elif Naz Ozdamar
parent
f90b5cad8a
commit
07bed09119
2 changed files with 153 additions and 46 deletions
|
|
@ -522,6 +522,7 @@ class RequestRateLimiterStash:
|
|||
owner_litellm_call_id: str | None = None
|
||||
rate_limit_response: RateLimitResponse | None = None
|
||||
parallel_slot: ParallelSlotAcquisition | None = None
|
||||
parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False)
|
||||
reserved_tokens: int = 0
|
||||
reserved_model: str | None = None
|
||||
reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
|
||||
|
|
@ -1609,6 +1610,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
statuses.append(self._gauge_status(gauge, in_flight + 1, "OK"))
|
||||
return RateLimitResponse(overall_code="OK", statuses=statuses)
|
||||
|
||||
async def _release_stashed_parallel_slot(
|
||||
self,
|
||||
stash: RequestRateLimiterStash | None,
|
||||
parent_otel_span: Span | None,
|
||||
) -> None:
|
||||
if stash is None:
|
||||
return
|
||||
async with stash.parallel_slot_release_lock:
|
||||
acquisition: Final = stash.parallel_slot
|
||||
if acquisition is None:
|
||||
return
|
||||
await self._release_parallel_request_slots(acquisition, parent_otel_span)
|
||||
stash.parallel_slot = None # rebind-ok: marks this request's slot as released
|
||||
|
||||
async def _release_parallel_request_slots(
|
||||
self,
|
||||
acquisition: ParallelSlotAcquisition,
|
||||
|
|
@ -3368,13 +3383,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
stash.reservation_released = True
|
||||
acquisition: Final = stash.parallel_slot
|
||||
if acquisition is not None:
|
||||
await self._release_parallel_request_slots(
|
||||
acquisition=acquisition,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
self._handle_rate_limit_error(
|
||||
response=io_response,
|
||||
descriptors=descriptors,
|
||||
|
|
@ -3631,13 +3640,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
|
||||
if tpm_response["overall_code"] == "OVER_LIMIT":
|
||||
acquisition: Final = stash.parallel_slot
|
||||
if acquisition is not None:
|
||||
await self._release_parallel_request_slots(
|
||||
acquisition=acquisition,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
self._handle_rate_limit_error(
|
||||
response=tpm_response,
|
||||
descriptors=descriptors,
|
||||
|
|
@ -4450,13 +4453,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING")
|
||||
|
||||
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
|
||||
acquisition: Final = stash.parallel_slot if stash is not None else None
|
||||
if stash is not None and acquisition is not None:
|
||||
await self._release_parallel_request_slots(
|
||||
acquisition=acquisition,
|
||||
parent_otel_span=litellm_parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span)
|
||||
|
||||
pipeline_operations: Final = self._build_success_event_pipeline_operations(
|
||||
kwargs=kwargs,
|
||||
|
|
@ -4576,13 +4573,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = []
|
||||
|
||||
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
|
||||
acquisition: Final = stash.parallel_slot if stash is not None else None
|
||||
if stash is not None and acquisition is not None:
|
||||
await self._release_parallel_request_slots(
|
||||
acquisition=acquisition,
|
||||
parent_otel_span=litellm_parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span)
|
||||
|
||||
# Skip the reservation refund if async_post_call_failure_hook
|
||||
# already released it (proxy-level rejection that also bubbles up
|
||||
|
|
@ -4690,23 +4681,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
object's current max_parallel_requests configuration, which can
|
||||
change mid-request) decides whether there is anything to release.
|
||||
"""
|
||||
stash: Final = get_request_stash()
|
||||
if stash is None or stash.parallel_slot is None:
|
||||
return
|
||||
|
||||
await self._release_parallel_request_slots(
|
||||
acquisition=stash.parallel_slot,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
await self._release_stashed_parallel_slot(get_request_stash(), None)
|
||||
|
||||
async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response):
|
||||
"""
|
||||
Post-call hook to update rate limit headers in the response.
|
||||
Release completed-request slots and update rate limit headers in the response.
|
||||
"""
|
||||
try:
|
||||
stash: Final = get_request_stash()
|
||||
litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None
|
||||
slot_stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(data))
|
||||
await self._release_stashed_parallel_slot(slot_stash, user_api_key_dict.parent_otel_span)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error releasing parallel request slot in post-call hook: %s", e)
|
||||
|
||||
try:
|
||||
header_stash: Final = get_request_stash()
|
||||
litellm_proxy_rate_limit_response: Final = (
|
||||
header_stash.rate_limit_response if header_stash is not None else None
|
||||
)
|
||||
|
||||
if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response):
|
||||
additional_headers: Final = ensure_response_additional_headers(response)
|
||||
|
|
@ -4774,12 +4765,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
stash: Final = get_request_stash()
|
||||
if stash is None:
|
||||
return
|
||||
if stash.parallel_slot is not None:
|
||||
await self._release_parallel_request_slots(
|
||||
acquisition=stash.parallel_slot,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
|
||||
if stash.batch_enqueued_reservation is not None:
|
||||
await self.batch_enqueued_token_store.refund(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import logging
|
|||
import os
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
|
@ -32,6 +33,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
|||
)
|
||||
from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.utils import (
|
||||
EmbeddingResponse,
|
||||
ModelResponse,
|
||||
|
|
@ -4054,6 +4056,125 @@ async def _seed_max_parallel_requests_slots(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_responses_post_call_releases_parallel_slot() -> None:
|
||||
api_key = hash_token("sk-responses-post-call")
|
||||
local_cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache)
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
|
||||
data = {
|
||||
"model": "gpt-4o-mini",
|
||||
"input": "hello",
|
||||
"litellm_call_id": "responses-owner",
|
||||
}
|
||||
parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests"
|
||||
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data=data,
|
||||
call_type="aresponses",
|
||||
)
|
||||
assert handler._gauge_in_flight_from_cache_value(
|
||||
await local_cache.async_get_cache(key=parallel_key)
|
||||
) == 1
|
||||
|
||||
await handler.async_post_call_success_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_parallel_slot",
|
||||
created_at=0,
|
||||
model="gpt-4o-mini",
|
||||
object="response",
|
||||
output=[],
|
||||
status="completed",
|
||||
),
|
||||
)
|
||||
assert handler._gauge_in_flight_from_cache_value(
|
||||
await local_cache.async_get_cache(key=parallel_key)
|
||||
) == 0
|
||||
|
||||
await handler.async_log_success_event(
|
||||
kwargs={"litellm_call_id": data["litellm_call_id"]},
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
assert handler._gauge_in_flight_from_cache_value(
|
||||
await local_cache.async_get_cache(key=parallel_key)
|
||||
) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None:
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
api_key = hash_token("sk-concurrent-release")
|
||||
local_cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache)
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2)
|
||||
call_id = "concurrent-release-owner"
|
||||
parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests"
|
||||
release_started = asyncio.Event()
|
||||
allow_redis_failure = asyncio.Event()
|
||||
|
||||
async def failing_release(
|
||||
keys: Sequence[str], args: Sequence[object]
|
||||
) -> list[int]:
|
||||
release_started.set()
|
||||
await allow_redis_failure.wait()
|
||||
raise ConnectionError("redis unavailable")
|
||||
|
||||
release_script = AsyncMock(side_effect=failing_release)
|
||||
handler.parallel_release_script = release_script
|
||||
await local_cache.async_set_cache(key=parallel_key, value=2, local_only=True)
|
||||
stash = get_or_create_request_stash()
|
||||
stash.owner_litellm_call_id = call_id
|
||||
stash.parallel_slot = ParallelSlotAcquisition(
|
||||
slot_id="slot-concurrent-release",
|
||||
counter_keys=[parallel_key],
|
||||
)
|
||||
data = {"litellm_call_id": call_id}
|
||||
|
||||
post_call_task = asyncio.create_task(
|
||||
handler.async_post_call_success_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_concurrent_release",
|
||||
created_at=0,
|
||||
model="gpt-4o-mini",
|
||||
object="response",
|
||||
output=[],
|
||||
status="completed",
|
||||
),
|
||||
)
|
||||
)
|
||||
await asyncio.wait_for(release_started.wait(), timeout=5)
|
||||
logging_task = asyncio.create_task(
|
||||
handler.async_log_success_event(
|
||||
kwargs=data,
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
)
|
||||
allow_redis_failure.set()
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(post_call_task, logging_task),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
assert release_script.await_count == 1
|
||||
assert await local_cache.async_get_cache(key=parallel_key) == 1
|
||||
assert stash.parallel_slot is None
|
||||
|
||||
|
||||
async def _build_seeded_limiter():
|
||||
"""Build a v3 limiter whose api-key slot registry already holds the pre-call slot."""
|
||||
api_key = hash_token("sk-disconnect")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue