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 227

Merge pull request #41113 from BerriAI/litellm_realtime_release_max_parallel_slot

Browse source 
fix(proxy): release max_parallel_requests slot when a realtime session ends without LLM callbacks
This commit is contained in:
Yassin Kortam 2026-09-14 14:56:11 -07:00 • committed by GitHub
commit 3e417fa6e0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 212 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/constants.py
View file

@ -317,6 +317,8 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
# SSL/TLS cipher configuration for faster handshakes
# Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones

5
litellm/litellm_core_utils/realtime_streaming.py
View file

@ -10,6 +10,7 @@ from typing_extensions import ReadOnly
import litellm
from litellm._logging import redact_internal_details_from_client_message, verbose_logger
from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.types.llms.openai import (
@ -35,9 +36,6 @@ else:
CLIENT_CONNECTION_CLASS = Any
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
@dataclass(frozen=True, slots=True)
class BackendClose:
code: int
@ -1153,6 +1151,7 @@ class RealTimeStreaming:
self._logging_worker.ensure_initialized_and_enqueue(
self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True)
)
self.logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True
@staticmethod
def _detect_beta_header(websocket: ScopedWebSocket) -> bool:

4
litellm/llms/bedrock/realtime/handler.py
View file

@ -20,6 +20,7 @@ from litellm.constants import (
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY,
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
)
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
@ -346,6 +347,7 @@ class BedrockRealtime(BaseAWSLLM):
prefer_async_handlers=True,
)
)
logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
if outcome.provider_failure is None:
return
@ -382,7 +384,7 @@ class BedrockRealtime(BaseAWSLLM):
)
bedrock_task: Final = asyncio.create_task(collect_logged_events())
await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_EXCEPTION)
await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_COMPLETED)
client_disconnected: Final = (
client_task.done() and not client_task.cancelled() and client_task.exception() is None
)

19
litellm/proxy/proxy_server.py
View file

@ -274,6 +274,8 @@ from litellm.constants import (
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
REALTIME_SESSION_FAILURE_LOGGED_KEY,
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG,
USER_SPEND_ALERTS_JOB_ID,
WEEKLY_SPEND_REPORT_JOB_ID,
@ -11899,6 +11901,13 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth
)
async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None:
release_like_http_disconnect: Final = (
proxy_logging_obj._arelease_max_parallel_requests_on_disconnect # pyright: ignore[reportPrivateUsage] # shared
)
await release_like_http_disconnect(user_api_key_dict)
async def _reject_realtime_session(
websocket: WebSocket,
user_api_key_dict: UserAPIKeyAuth,
@ -11918,6 +11927,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")
@ -12021,6 +12031,9 @@ async def realtime_websocket_endpoint(
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.
try:
@ -12050,12 +12063,10 @@ async def realtime_websocket_endpoint(
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)
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY):
await _release_realtime_max_parallel_slot(user_api_key_dict)
######################################################################

22
tests/test_litellm/litellm_core_utils/test_realtime_streaming.py
View file

@ -10,9 +10,9 @@ from websockets.exceptions import ConnectionClosed
from websockets.frames import Close
import litellm
from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.realtime_streaming import (
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
RealTimeStreaming,
client_sent_openai_beta_realtime_header,
)
@ -3399,6 +3399,26 @@ async def test_refused_session_does_not_stamp_the_reservation_ownership_marker()
assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details
@pytest.mark.asyncio
async def test_refused_session_stamps_the_failure_ownership_marker():
"""LIT-6463: the enqueued failure callback releases the key's max_parallel_requests
slot from the logging worker, so a refusal stamps REALTIME_SESSION_FAILURE_LOGGED_KEY.
The proxy endpoint reads it to leave the slot to that callback instead of racing it.
A session that relayed frames logs a success and must not carry the failure stamp."""
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
refused: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close))
session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode()
relayed: Final = _relay_session(
_client_ws_that_never_sends(), _backend_ws_closing_with(session_created, upstream_close)
)
await refused.run()
await relayed.run()
assert refused.logging.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY) is True
assert REALTIME_SESSION_FAILURE_LOGGED_KEY not in relayed.logging.model_call_details
@pytest.mark.asyncio
async def test_transformed_transcription_completion_never_sends_response_create():
from typing import Final

46
tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py
View file

@ -8,6 +8,7 @@ from unittest.mock import MagicMock
import pytest
import litellm
from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.realtime.handler import BedrockRealtime
from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig
@ -77,6 +78,7 @@ class UnavailableBedrockStream:
class FakeLogging:
def __init__(self, trace_id="trace-nova-sonic"):
self.litellm_trace_id = trace_id
self.model_call_details = {}
class DisconnectingClientWS:
@ -673,6 +675,31 @@ class TestBedrockRealtimeProviderFailurePropagation:
await spend_dispatch["coro"]
assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
@pytest.mark.asyncio
async def test_success_dispatch_stamps_the_ownership_marker_only_when_spend_was_logged(
self, stub_aws_sdk_client, spend_dispatch
):
stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)]
await BedrockRealtime().async_realtime(
model="amazon.nova-sonic-v1:0",
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
logging_obj=spend_dispatch["logging_obj"],
**self.AWS_PARAMS,
)
await spend_dispatch["coro"]
assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True
idle_logging = FakeLogging()
stub_aws_sdk_client["streams"] = [ScriptedBedrockStream([])]
await BedrockRealtime().async_realtime(
model="amazon.nova-sonic-v1:0",
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
logging_obj=idle_logging,
**self.AWS_PARAMS,
)
assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in idle_logging.model_call_details
@pytest.mark.asyncio
async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client):
stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver)
@ -684,6 +711,25 @@ class TestBedrockRealtimeProviderFailurePropagation:
assert stream.input_stream.closed
@pytest.mark.asyncio
async def test_client_disconnect_ends_the_session_while_bedrock_output_stays_open(self, stub_aws_sdk_client):
receiver = DrainedThenOpenBedrockReceiver([])
stream = ScriptedBedrockStream([], receiver_type=lambda _payloads: receiver)
stub_aws_sdk_client["streams"] = [stream]
await asyncio.wait_for(
BedrockRealtime().async_realtime(
model="amazon.nova-sonic-v1:0",
websocket=RealtimeClientWS(),
logging_obj=FakeLogging(),
**self.AWS_PARAMS,
),
timeout=1,
)
assert receiver.drained.is_set(), "the handler must have been waiting on the open provider stream"
assert stream.input_stream.closed
@pytest.mark.asyncio
async def test_session_updated_is_not_sent_before_bedrock_is_ready(self, stub_aws_models):
handler = BedrockRealtime()

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

@ -33,6 +33,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
@ -9999,6 +10000,7 @@ async def _lit6973_drive_realtime_session(
reservation: dict,
*,
backend_logged_success: bool,
backend_logged_failure: bool = False,
phase_one_exit: str | None = None,
websocket: MagicMock | None = None,
) -> MagicMock:
@ -10006,8 +10008,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
@ -10017,7 +10020,7 @@ async def _lit6973_drive_realtime_session(
logging object carries a real model_call_details dict so the stamp is
observable, and the reservation has empty entries so the real release touches
no counter store."""
from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY
from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
from litellm.proxy import proxy_server as ps
user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token")
@ -10029,6 +10032,8 @@ async def _lit6973_drive_realtime_session(
async def fake_llm_call() -> None:
if backend_logged_success:
logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
if backend_logged_failure:
logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True
from litellm.proxy._types import ProxyException
@ -10037,7 +10042,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)
)
@ -10155,6 +10166,114 @@ 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,
backend_logged_failure: bool = False,
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( # test-quality-ok: registers the real limiter the route's release reads
ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}
)
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,
backend_logged_failure=backend_logged_failure,
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(
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, 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."""
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
backend_logged_success=False, phase_one_exit=phase_one_exit
)
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_refused_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_failure_callback():
"""An upstream refusal before any frame enqueues the failure callback instead, and
the limiter's failure handler releases the slot from the logging worker just like
the success handler does. The route sees no success stamp, so it still settles the
budget reservation, but it must leave the slot to that callback or the two releases
race on the same acquisition."""
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
backend_logged_success=False, backend_logged_failure=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),