mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(proxy): release realtime max_parallel_requests slot when the task is cancelled during pre-call
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
ebcd9bcb18
commit
81ffc3125f
2 changed files with 77 additions and 65 deletions
|
|
@ -11912,7 +11912,6 @@ async def _reject_realtime_session(
|
|||
await websocket.close(code=code, reason=reason)
|
||||
finally:
|
||||
await _release_realtime_budget_reservation(user_api_key_dict)
|
||||
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
|
||||
|
||||
|
||||
@app.websocket("/openai/v1/realtime")
|
||||
|
|
@ -11992,65 +11991,67 @@ 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:
|
||||
(
|
||||
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
|
||||
|
||||
# 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:
|
||||
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"),
|
||||
(
|
||||
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: # 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,
|
||||
)
|
||||
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
|
||||
|
||||
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY):
|
||||
await _release_realtime_budget_reservation(user_api_key_dict)
|
||||
# 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:
|
||||
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)
|
||||
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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -9919,8 +9919,9 @@ async def _lit6973_drive_realtime_session(
|
|||
|
||||
phase_one_exit picks a rejection before the relay: "model_access" makes the
|
||||
key/model check raise ProxyException, "pre_call" makes pre-call processing
|
||||
(rate limits, guardrails) raise. Neither reaches route_request, so no success
|
||||
log can own the reservation and the endpoint has to release it on that exit.
|
||||
(rate limits, guardrails) raise, "pre_call_cancelled" cancels the task inside
|
||||
pre-call processing. None reaches route_request, so no success log can own the
|
||||
reservation and the endpoint has to release it on that exit.
|
||||
|
||||
route_request resolves normally in both cases: the relay owns the session
|
||||
once route_request returns. A successful session enqueues its success cost
|
||||
|
|
@ -9950,7 +9951,13 @@ async def _lit6973_drive_realtime_session(
|
|||
if phase_one_exit == "model_access"
|
||||
else None
|
||||
)
|
||||
pre_call_error: Final = Exception("Rate limit exceeded") if phase_one_exit == "pre_call" else None
|
||||
pre_call_error: Final = (
|
||||
asyncio.CancelledError()
|
||||
if phase_one_exit == "pre_call_cancelled"
|
||||
else Exception("Rate limit exceeded")
|
||||
if phase_one_exit == "pre_call"
|
||||
else None
|
||||
)
|
||||
pre_call: Final = AsyncMock(
|
||||
side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)
|
||||
)
|
||||
|
|
@ -10069,15 +10076,16 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase_one_exit", [None, "pre_call"])
|
||||
@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(
|
||||
phase_one_exit: str | None,
|
||||
):
|
||||
"""The rate limiter acquires the key's max_parallel_requests slot in pre-call and
|
||||
only frees it from the LLM success/failure callbacks. A realtime session that ends
|
||||
without either callback (Bedrock closes without usage events, or a later pre-call
|
||||
hook rejects the session) 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."""
|
||||
without either callback (Bedrock closes without usage events, a later pre-call hook
|
||||
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 (
|
||||
|
|
@ -10097,7 +10105,10 @@ async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_pa
|
|||
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
|
||||
with hooks:
|
||||
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
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue