scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(proxy): leave the realtime max_parallel slot to the success callback when one is enqueued

Browse source 
Releasing the slot unconditionally from the route raced the limiter's own
success handler on the logging worker: both could read the same stashed
acquisition before either cleared it, and under the integer in-memory
fallback that double-decrements the counter. The route now releases only on
exits without a success callback (pre-call rejection, pre-call cancellation,
and Phase 2 exits without the success stamp), matching the HTTP disconnect
path's ownership rule.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-14 20:23:23 +00:00
parent 81ffc3125f
commit 1b31be1a9c
2 changed files with 131 additions and 86 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

122
litellm/proxy/proxy_server.py
View file

@ -11893,6 +11893,10 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth
)
async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None:
await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses
async def _reject_realtime_session(
websocket: WebSocket,
user_api_key_dict: UserAPIKeyAuth,
@ -11912,6 +11916,7 @@ async def _reject_realtime_session(
await websocket.close(code=code, reason=reason)
finally:
await _release_realtime_budget_reservation(user_api_key_dict)
await _release_realtime_max_parallel_slot(user_api_key_dict)
@app.websocket("/openai/v1/realtime")
@ -11991,68 +11996,69 @@ async def realtime_websocket_endpoint(
# Errors here (e.g. guardrail block) are sent back to the client as an
# error event before closing, so the caller knows what happened.
try:
try:
(
data,
litellm_logging_obj,
) = await base_llm_response_processor.common_processing_pre_call_logic(
request=request,
general_settings=general_settings,
user_api_key_dict=user_api_key_dict,
version=version,
proxy_logging_obj=proxy_logging_obj,
proxy_config=proxy_config,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=route_model,
route_type="_arealtime",
)
except Exception as e:
verbose_proxy_logger.exception("Realtime pre-call error")
await _reject_realtime_session(
websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e)
)
return
(
data,
litellm_logging_obj,
) = await base_llm_response_processor.common_processing_pre_call_logic(
request=request,
general_settings=general_settings,
user_api_key_dict=user_api_key_dict,
version=version,
proxy_logging_obj=proxy_logging_obj,
proxy_config=proxy_config,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=route_model,
route_type="_arealtime",
)
except Exception as e:
verbose_proxy_logger.exception("Realtime pre-call error")
await _reject_realtime_session(
websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e)
)
return
except BaseException:
await _release_realtime_max_parallel_slot(user_api_key_dict)
raise
# Phase 2: route to upstream LLM.
# Phase 2: route to upstream LLM.
try:
data["user_api_key_dict"] = user_api_key_dict
llm_call: Final = await route_request(
data=data,
route_type="_arealtime",
llm_router=llm_router,
user_model=user_model,
)
await llm_call
except websockets.exceptions.InvalidStatusCode as e:
verbose_proxy_logger.exception("Invalid status code")
await websocket.close(code=e.status_code, reason="Invalid status code")
except Exception as e:
verbose_proxy_logger.exception("Internal server error")
redacted_error: Final = _redact_string(str(e))
try:
data["user_api_key_dict"] = user_api_key_dict
llm_call: Final = await route_request(
data=data,
route_type="_arealtime",
llm_router=llm_router,
user_model=user_model,
await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error"))
except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below
verbose_proxy_logger.debug("Could not send realtime error event to client; closing anyway")
try:
await websocket.close(
code=1011,
reason=websocket_close_reason(redacted_error, fallback="Internal server error"),
)
await llm_call
except websockets.exceptions.InvalidStatusCode as e:
verbose_proxy_logger.exception("Invalid status code")
await websocket.close(code=e.status_code, reason="Invalid status code")
except Exception as e:
verbose_proxy_logger.exception("Internal server error")
redacted_error: Final = _redact_string(str(e))
try:
await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error"))
except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below
verbose_proxy_logger.debug("Could not send realtime error event to client; closing anyway")
try:
await websocket.close(
code=1011,
reason=websocket_close_reason(redacted_error, fallback="Internal server error"),
)
except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error
verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone")
finally:
from litellm.litellm_core_utils.realtime_streaming import (
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
)
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY):
await _release_realtime_budget_reservation(user_api_key_dict)
except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error
verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone")
finally:
await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses
from litellm.litellm_core_utils.realtime_streaming import (
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
)
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY):
await _release_realtime_budget_reservation(user_api_key_dict)
await _release_realtime_max_parallel_slot(user_api_key_dict)
######################################################################

95
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -31,6 +31,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash
from litellm.proxy.proxy_server import app, initialize
from litellm.utils import _invalidate_model_cost_lowercase_map
@ -10075,6 +10076,50 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c
assert reservation["finalized"] is False
_LIT6463_COUNTER_KEY: Final = "{api_key:hashed-token}:max_parallel_requests"
async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
*,
backend_logged_success: bool,
phase_one_exit: str | None = None,
) -> tuple[DualCache, RequestRateLimiterStash]:
"""Run the realtime endpoint with a real v3 limiter registered and the request's
stash already holding slot-1 of a two-slot counter, the state pre-call leaves
behind. Returns the limiter's cache and the stash so the test can read what the
endpoint did to the slot."""
from litellm.proxy import proxy_server as ps
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
_request_stash,
)
from litellm.proxy.utils import InternalUsageCache
dual_cache: Final = DualCache()
await dual_cache.async_set_cache(
key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True
)
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache))
stash: Final = RequestRateLimiterStash(
parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}
)
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
stash_token: Final = _request_stash.set(stash)
try:
hooks: Final = patch.dict(ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}) # test-quality-ok: registers a real limiter on the module-global hook map the route reads; assertion observes its counter
expected_exit: Final = (
pytest.raises(asyncio.CancelledError) if phase_one_exit == "pre_call_cancelled" else contextlib.nullcontext()
)
with hooks, expected_exit:
await _lit6973_drive_realtime_session(
reservation, backend_logged_success=backend_logged_success, phase_one_exit=phase_one_exit
)
finally:
_request_stash.reset(stash_token)
return dual_cache, stash
@pytest.mark.asyncio
@pytest.mark.parametrize("phase_one_exit", [None, "pre_call", "pre_call_cancelled"])
async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_parallel_slot(
@ -10086,39 +10131,33 @@ async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_pa
rejects the session, or the task is cancelled while still in pre-call) has to be
released by the route itself, or the slot stays occupied until its TTL and the key's
next session is refused with a 429."""
from litellm.caching.caching import DualCache
from litellm.proxy import proxy_server as ps
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
RequestRateLimiterStash,
_PROXY_MaxParallelRequestsHandler_v3,
_request_stash,
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
backend_logged_success=False, phase_one_exit=phase_one_exit
)
from litellm.proxy.utils import InternalUsageCache
counter_key: Final = "{api_key:hashed-token}:max_parallel_requests"
dual_cache: Final = DualCache()
await dual_cache.async_set_cache(key=counter_key, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True)
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache))
stash: Final = RequestRateLimiterStash(parallel_slot={"slot_id": "slot-1", "counter_keys": [counter_key]})
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
stash_token: Final = _request_stash.set(stash)
try:
hooks: Final = patch.dict(ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}) # test-quality-ok: registers a real limiter on the module-global hook map the route reads; assertion observes its counter
expected_exit: Final = (
pytest.raises(asyncio.CancelledError) if phase_one_exit == "pre_call_cancelled" else contextlib.nullcontext()
)
with hooks, expected_exit:
await _lit6973_drive_realtime_session(
reservation, backend_logged_success=False, phase_one_exit=phase_one_exit
)
finally:
_request_stash.reset(stash_token)
assert await dual_cache.async_get_cache(key=counter_key, local_only=True) == {"slot-2": 2.0}
assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {"slot-2": 2.0}
assert stash.parallel_slot is None
@pytest.mark.asyncio
async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_callback():
"""A session that enqueued its success callback hands the slot to the limiter's
own success handler, which runs on the logging worker. If the route also released
it, the two releases would race on the same stashed acquisition and, under the
limiter's integer in-memory fallback, double-decrement the counter so the key
admits more sessions than max_parallel_requests allows. With the success stamp
present the route leaves the slot and the stash alone."""
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
backend_logged_success=True
)
assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {
"slot-1": 1.0,
"slot-2": 2.0,
}
assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}
@pytest.mark.asyncio
async def test_release_or_invalidate_falls_back_to_invalidating_the_counters():
"""If releasing the reservation itself fails (e.g. the counter store is down),