scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-20 05:04:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 155

fix(proxy): keep the legacy limiter's WebSocket slot until the socket closes

Browse source
This commit is contained in:
mateo-berri 2026-10-15 19:44:50 -07:00
parent 10d6748a90
commit 6ccd11643f
3 changed files with 193 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

77
litellm/proxy/hooks/parallel_request_limiter.py
View file

@ -1,6 +1,8 @@
import asyncio
import sys
from collections.abc import Callable
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn
@ -23,7 +25,7 @@ from litellm.proxy.auth.budget_throttle import throttled_limit
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.types.utils import Usage
from litellm.types.utils import CallTypes, Usage
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -48,6 +50,17 @@ def _response_total_tokens(response_obj: object) -> int:
return response_usage.total_tokens if isinstance(response_usage, Usage) else 0
@dataclass(frozen=True, slots=True)
class _WebSocketConnectionSlot:
request_count_key: str | None
holds_global_slot: bool
_websocket_connection_slot: Final[ContextVar[_WebSocketConnectionSlot | None]] = ContextVar(
"litellm_legacy_rate_limiter_websocket_connection_slot", default=None
)
class CacheObject(TypedDict):
current_global_requests: dict | None
request_count_api_key: dict | None
@ -63,7 +76,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
self.internal_usage_cache = internal_usage_cache
self._clock: Final = clock
def print_verbose(self, print_statement) -> None:
def print_verbose(self, print_statement: object) -> None:
try:
verbose_proxy_logger.debug(print_statement)
if litellm.set_verbose:
@ -499,6 +512,41 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
) # don't block execution for cache updates
)
if call_type == CallTypes.aresponses_websocket.value:
_websocket_connection_slot.set(
_WebSocketConnectionSlot(
request_count_key=(f"{api_key}::{precise_minute}::request_count" if api_key is not None else None),
holds_global_slot=global_max_parallel_requests is not None,
)
)
@with_service_target("rate_limits")
async def async_release_websocket_connection_slot(self) -> None:
connection: Final = _websocket_connection_slot.get()
if connection is None:
return
_websocket_connection_slot.set(None)
if connection.holds_global_slot:
await self.internal_usage_cache.async_increment_cache(
key="global_max_parallel_requests",
value=-1,
local_only=True,
litellm_parent_otel_span=None,
)
if connection.request_count_key is None:
return
current: Final = await self.internal_usage_cache.async_get_cache(
key=connection.request_count_key,
litellm_parent_otel_span=None,
)
if current is None:
return
await self.internal_usage_cache.async_set_cache(
connection.request_count_key,
{**current, "current_requests": max(current["current_requests"] - 1, 0)},
ttl=60,
litellm_parent_otel_span=None,
)
@with_service_target("rate_limits")
async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None:
@ -520,6 +568,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
"user_api_key_model_max_budget", None
)
user_api_key_end_user_id: Final = kwargs.get("user")
requests_released: Final = 0 if _websocket_connection_slot.get() is not None else 1
user_api_key_metadata: Final = kwargs["litellm_params"]["metadata"].get("user_api_key_metadata", {}) or {}
user_api_key_team_metadata = kwargs["litellm_params"]["metadata"].get("user_api_key_team_metadata", None)
@ -540,7 +589,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
# decrement
await self.internal_usage_cache.async_increment_cache(
key=_key,
value=-1,
value=-requests_released,
local_only=True,
litellm_parent_otel_span=litellm_parent_otel_span,
)
@ -562,13 +611,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_requests": requests_released,
"current_tpm": 0,
"current_rpm": 0,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -603,13 +652,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_requests": requests_released,
"current_tpm": 0,
"current_rpm": 0,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -629,13 +678,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_requests": requests_released,
"current_tpm": total_tokens,
"current_rpm": 1,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -655,13 +704,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_requests": requests_released,
"current_tpm": total_tokens,
"current_rpm": 1,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -681,13 +730,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_requests": requests_released,
"current_tpm": total_tokens,
"current_rpm": 1,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -712,7 +761,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
global_max_parallel_requests: Final = _metadata.get("global_max_parallel_requests", None)
user_api_key: Final = _metadata.get("user_api_key", None)
self.print_verbose(f"user_api_key: [set={user_api_key is not None}]")
if user_api_key is None:
if user_api_key is None or _websocket_connection_slot.get() is not None:
return
## decrement call count if call failed

3
litellm/proxy/utils.py
View file

@ -4405,6 +4405,9 @@ class ProxyLogging:
double-decrement under the limiter's in-memory fallback.
"""
limiter: Final = self.get_proxy_hook("parallel_request_limiter")
if isinstance(limiter, PROXY_MaxParallelRequestsHandler):
await limiter.async_release_websocket_connection_slot()
return
if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3):
return
await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)

130
tests/unit/proxy/hooks/test_parallel_request_limiter.py
View file

@ -2,8 +2,9 @@
Unit Tests for the max parallel request limiter v1 for the proxy
"""
import asyncio
import itertools
from collections.abc import Callable, Iterator, Mapping
from collections.abc import Awaitable, Callable, Iterator, Mapping
from datetime import datetime
from typing import Final
@ -12,11 +13,12 @@ import pytest
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.parallel_request_limiter import (
PROXY_MaxParallelRequestsHandler,
)
from litellm.proxy.utils import InternalUsageCache, hash_token
from litellm.types.utils import EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage
from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
from litellm.types.utils import CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage
FROZEN_INSTANT: Final = datetime(2026, 1, 31, 23, 59, 30)
LAST_MICROSECOND_OF_JANUARY: Final = datetime(2026, 1, 31, 23, 59, 59, 999999)
@ -323,3 +325,125 @@ async def test_retry_after_is_the_seconds_until_the_next_minute_of_the_injected_
)
assert exc_info.value.headers["retry-after"] == "14.5"
async def _log_websocket_turn_success(
handler: PROXY_MaxParallelRequestsHandler, metadata: Mapping[str, object]
) -> None:
await handler.async_log_success_event(
kwargs={"litellm_params": {"metadata": dict(metadata)}},
response_obj=None,
start_time=FROZEN_INSTANT,
end_time=FROZEN_INSTANT,
)
async def _log_websocket_turn_failure(
handler: PROXY_MaxParallelRequestsHandler, metadata: Mapping[str, object]
) -> None:
await handler.async_log_failure_event(
kwargs={"litellm_params": {"metadata": dict(metadata)}, "exception": Exception("upstream turn failed")},
response_obj=None,
start_time=FROZEN_INSTANT,
end_time=FROZEN_INSTANT,
)
def _proxy_logging_with(handler: PROXY_MaxParallelRequestsHandler) -> ProxyLogging:
proxy_logging: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache())
proxy_logging.proxy_hook_mapping["parallel_request_limiter"] = handler
return proxy_logging
async def _requests_admitted(
handler: PROXY_MaxParallelRequestsHandler, session: UserAPIKeyAuth, metadata: Mapping[str, object]
) -> int:
for admitted in range(3):
try:
await handler.async_pre_call_hook(
user_api_key_dict=session,
cache=DualCache(),
data={"metadata": dict(metadata)},
call_type=CallTypes.acompletion.value,
)
except ProxyRateLimitError:
return admitted
await asyncio.sleep(0)
return 3
@pytest.mark.parametrize("log_turn", [_log_websocket_turn_success, _log_websocket_turn_failure])
@pytest.mark.parametrize(
"max_parallel_requests,limit_metadata,admitted_while_open",
[(1, {}, 0), (None, {"global_max_parallel_requests": 2}, 1)],
ids=["key_limit", "global_limit"],
)
@pytest.mark.asyncio
async def test_a_websocket_connection_holds_its_parallel_slot_through_its_turns_until_it_closes(
log_turn: Callable[[PROXY_MaxParallelRequestsHandler, Mapping[str, object]], Awaitable[None]],
max_parallel_requests: int | None,
limit_metadata: Mapping[str, object],
admitted_while_open: int,
):
handler: Final = PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock
)
api_key: Final = hash_token("sk-websocket-slot")
session: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=max_parallel_requests)
turn_logged: Final = asyncio.Event()
socket_closed: Final = asyncio.Event()
async def websocket_connection() -> None:
await handler.async_pre_call_hook(
user_api_key_dict=session,
cache=DualCache(),
data={"metadata": dict(limit_metadata)},
call_type=CallTypes.aresponses_websocket.value,
)
await asyncio.sleep(0)
await log_turn(handler, {"user_api_key": api_key, "user_api_key_model_max_budget": {}, **limit_metadata})
turn_logged.set()
await socket_closed.wait()
await _proxy_logging_with(handler).arelease_max_parallel_requests_on_disconnect(session)
connection: Final = asyncio.create_task(websocket_connection())
await turn_logged.wait()
assert await _requests_admitted(handler, session, limit_metadata) == admitted_while_open
socket_closed.set()
await connection
assert await _requests_admitted(handler, session, limit_metadata) == 1
@pytest.mark.asyncio
async def test_a_websocket_turn_logged_in_a_later_minute_leaves_no_slot_held_after_the_socket_closes():
handler: Final = PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), clock=_clock_rolling_over_after_first_read()
)
api_key: Final = hash_token("sk-websocket-rollover")
session: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
async def websocket_connection() -> None:
await handler.async_pre_call_hook(
user_api_key_dict=session, cache=DualCache(), data={}, call_type=CallTypes.aresponses_websocket.value
)
await asyncio.sleep(0)
await _log_websocket_turn_success(handler, {"user_api_key": api_key, "user_api_key_model_max_budget": {}})
await handler.async_release_websocket_connection_slot()
await asyncio.create_task(websocket_connection())
assert await _requests_admitted(handler, session, {}) == 1
@pytest.mark.asyncio
async def test_a_disconnect_release_outside_a_websocket_connection_leaves_the_request_slot_held():
handler: Final = PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock
)
session: Final = UserAPIKeyAuth(api_key=hash_token("sk-http-slot"), max_parallel_requests=1)
assert await _requests_admitted(handler, session, {}) == 1
await _proxy_logging_with(handler).arelease_max_parallel_requests_on_disconnect(session)
assert await _requests_admitted(handler, session, {}) == 0