[Core][Frontend] Bound duplex output delivery and reject cancelled audio (#7644)

Signed-off-by: Nolen Liang <nliang@nvidia.com>
Signed-off-by: bcsdhjew <nliang@nvidia.com>
This commit is contained in:
bcsdhjew authored and GitHub committed 2026-10-04 14:06:04 -07:00
1 parent 3f20252151
commit f781f7bb08
26 files changed
+1151 -128

No files matched your search

+55 -9
View File
@@ -220,6 +220,7 @@ vllm_omni/
│ │ DUPLEX_REALTIME_CAPABILITIES, duplex_response_format
│ ├── events.py re-export shim (DuplexEvent = RealtimeEvent) + the runner's
│ │ epoch-filter sets DOMAIN_TERMINAL_EVENTS / MODEL_OUTPUT_EVENTS
│ ├── delivery.py DuplexOutputBuffer (bounded output shared with the session handle)
│ ├── realtime_events.py RealtimeProjectionState: internal event -> typed events
│ ├── messages.py queue envelopes (Open/Close/Resume/Touch/Command/Result/Event), DuplexSessionError
│ ├── config.py DuplexSessionConfig, DuplexCapabilities, ResponseCreateOptions
@@ -284,7 +285,8 @@ command handle.submit(DuplexCommand)
-> engine.submit_command_async -> DuplexSessionCommandMessage on the request queue [one-way]
-> DuplexSessionManager.dispatch: unknown_session / input_backpressure checks, then runner mailbox
output DuplexOrchestrator._intercept_stage_output -> runner.on_stage_output -> mailbox -> typed events
-> output_sink (DuplexSessionEventMessage) -> DuplexOmni._route_engine_message -> handle.events()
-> session's bounded output buffer -> handle.events(); closure and errors without a live session
use output_sink (DuplexSessionEventMessage) -> DuplexOmni._route_engine_message
detach DuplexOmni.detach_session(expected_lease_generation) -> touch(DETACH): engine-owned disconnect
grace, refused for a lease newer than the one the caller held; expiry -> SessionExpired
resume DuplexOmni.resume_session(expected_lease_generation) -> lease CAS keyed by the control id (a replay
@@ -303,7 +305,9 @@ reap DuplexSessionManager.reaper_loop: idle TTL / disconnect grace expiry, c
### Concurrency and ordering model of the runner
Everything below runs on the orchestrator asyncio loop; there is no lock.
The runner's session state is owned by the orchestrator asyncio loop and
needs no lock. The output buffer shared with the caller thread uses a lock
for queue updates and the final validity check before delivery.
- **Single writer per session.** `DuplexEngineSession` is mutated only by its
runner. Inputs reach the runner through one mailbox in this order: client
@@ -332,13 +336,15 @@ Everything below runs on the orchestrator asyncio loop; there is no lock.
on one loop in program order, nothing can be emitted for an old epoch after
its terminal event; `runner.emit()` still drops a terminal carrying a stale
epoch and stamps `epoch` on every event for clients that filter
defensively. This is a statement about order, not about latency: the engine
output queue and each `DuplexSessionHandle` outbox are unbounded and FIFO,
so a cancellation is delivered behind whatever audio was already emitted for
the response it cancels. A client that must stop quickly cancels its own
playback on `response.done` / `audio.cancelled` rather than waiting for the
stream to drain. Bounding those buffers and letting an accepted invalidation
skip undelivered media is left to the follow-up RFC.
defensively. In this experimental follow-up, public output goes directly to
one bounded buffer shared by the engine and `DuplexSessionHandle`. Accepted
cancellation removes queued `AudioDelta` events for that response through
the cancelled epoch. Text increments, audio/transcript completion markers,
other responses and normal completion retain FIFO order. The consumer also
checks an event already taken from the buffer: WebSocket delivery repeats
that check under the connection lock, before sequencing and journaling.
Already sequenced replay entries are unchanged; already delivered audio
still requires the client to stop its own playback.
- **Backpressure before the mailbox.** `DuplexSessionManager.dispatch`
checks `max_pending_input_bytes_per_session` and reserves a pending turn for
`Commit` before the put; a rejected command is answered with
@@ -349,6 +355,46 @@ Everything below runs on the orchestrator asyncio loop; there is no lock.
`ResponseDone.stage_metrics` and `metadata.vllm_omni.stage_metrics` on the
wire keep their meaning.
### Experimental output limits
`max_pending_output_bytes_per_session` defaults to 2 MiB and
`max_pending_output_events_per_session` to 512. The byte limit counts compact
Realtime JSON, including base64 audio, not decoded PCM or Python object
overhead. Audio is already base64-encoded by the producer; counting its length
avoids serializing or decoding that large field again. Neither default has
been tuned from a performance study.
Errors and response endings have an additional 64 KiB / eight-event reserve
and retain FIFO order. One final session-closure notification has an independent
slot, so an exhausted reserve cannot prevent closure or break delivery to
other sessions. If that reserve also fills, an error or response-ending
notification may be omitted; the final session-closure notification still
has its independent slot. Oversized close details (over 4 KiB with generated identity)
are replaced by a compact reason, preserving terminal type and identity.
Late writes after closure are ignored; a locally ended iterator is not reopened
by a late terminal, and a failed open closes the abandoned handle. Events other
than errors are not forwarded once their session's buffer has been removed.
One event held by the consumer is outside the queued budget.
The overflow policy fails the active response without committing its partial
text to history and closes only the affected session. If overflow occurs while
a completion batch is being delivered, its not-yet-queued `ResponseDone` is
converted to failure with its identity and output preserved; any history entry
for that response is removed. An ending already accepted by the buffer is not
replaced or followed by a second ending. After overflow, newly
emitted ordinary events are suppressed, including text/audio completion,
content-part and output-item completion markers. Previously queued non-audio
events retain FIFO order; error and failed response-ending notifications may
use the reserve before the independent closure notification. This differs from
accepted cancellation, which removes only matching audio. See the
[public error and closure fields](../serving/realtime_duplex_api.md#output-limits-and-slow-consumers).
Response-only recovery is not implemented. These limits do not bound the runner's raw stage
output mailbox, replay journal, or client playback queues. The shared buffer
requires the current same-process, separate-thread engine arrangement; it is
not a cross-process delivery protocol. Direct handle iteration and WebSocket
delivery share buffer validity, but end-to-end client behavior still needs
independent real-model verification.
## Orchestrator seams
`OrchestratorBase` keeps the generic machinery (request handler skeleton,
+37 -1
View File
@@ -250,6 +250,40 @@ call; transport failures raise `DuplexConnectionError`; a session that ends
while you are awaiting it raises `DuplexSessionClosedError`. All three derive
from `DuplexClientError`.
### Output limits and slow consumers
Deployment settings under `duplex_session` bound each session's pending
output: `max_pending_output_bytes_per_session` defaults to 2 MiB and
`max_pending_output_events_per_session` to 512 events. Bytes count compact
Realtime JSON, including base64 audio, not decoded audio. The limits apply
to both WebSocket delivery and direct Python handles, not to replay journals
or client playback queues.
Exceeding a limit fails the active response without committing its partial
text to history, then closes only that session. If overflow occurs while a
completion is being queued, its not-yet-queued `response.done` becomes failed
with the same identity and output, and its response history entry is removed.
An ending already accepted by the buffer is not replaced or followed by a
second ending. Open a new session to continue; the closed session cannot be
resumed. The notifications use these fields:
| Event | Field values |
| --- | --- |
| `error` | `error.type: "rate_limit_error"`, `error.code: "output_backpressure"` |
| `response.done` (active response or not-yet-queued ending) | `response.status: "failed"`, `response.status_details: {"type": "failed", "reason": "output_backpressure"}` |
| `session.closed` | `reason: "output_backpressure"` |
Queued audio for the affected response is removed. Previously queued non-audio
events retain their order, but newly emitted ordinary events, including
audio/text, content-part and output-item completion markers, are suppressed.
Clients must therefore handle closure without waiting for every completion
marker. Errors and response endings have an extra 64 KiB / eight-event reserve;
if that also fills, either notification may be omitted. Closure has its own
independent slot. A closure notification exceeding 4 KiB is compacted to
`reason: "close_details_exceed_output_limit"` with empty details, preserving
the event type and identity. One event held by the consumer is outside the
queued budget; audio already sent still requires the client to stop playback.
### Measure latency
`EventCollector` accumulates events for assertions and metrics:
@@ -315,6 +349,8 @@ single-consumer async iterator that ends after `session.closed` /
`session.expired`; rejected commands come back as `ErrorEvent`s on it.
Every event renders the wire JSON with `to_realtime()`, so anything written
against the WebSocket protocol works unchanged on the typed stream.
If opening fails, the abandoned handle is closed. Late engine events do not
reopen a locally closed handle.
`vllm_omni.clients.inline_duplex.InlineDuplexClient` wraps a `DuplexOmni`
behind the `DuplexClient` API, so the same application code runs in-process
@@ -573,7 +609,7 @@ formats (`pcm16`, `pcm_s16le`, `s16le`, `pcm_f32le`, `g711_ulaw`,
| `session.created` | 2 | Session opened. OpenAI name; adds `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`id` is the server-allocated session id, `epoch`, `turn_id`, `playback`, `capabilities`, ...). |
| `session.updated` | 2 | Echo of the effective session config after every `session.update` (same extended `session` object). |
| `session.heartbeat_ack` | 3 | Reply to `session.heartbeat`. |
| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, ... |
| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, `output_backpressure`, `close_details_exceed_output_limit`, ... |
| `session.resumed` | 3 | Resume accepted; carries the new `attachment_generation` and rotated `resume_token`; journaled events are replayed after it. |
| `session.replaced` | 3 | Sent to the superseded socket when another socket resumes the session. |
| `session.expired` | 3 | Engine lease reaped (`disconnect_grace_expired`, `idle_ttl_expired`); socket closes after it. |
@@ -370,6 +370,7 @@ def test_draining_empty_eos_completes_owning_response_before_shortcut() -> None:
stage_port=port,
model_state=SimpleNamespace(clear_continuation=lambda: None),
services=SimpleNamespace(spawn=lambda *_a, **_k: None),
run=SimpleNamespace(closing=False),
),
out,
close_from_runtime=_noop,
@@ -461,6 +462,7 @@ def test_draining_completion_drops_finished_response_books_only() -> None:
stage_port=_Port(),
model_state=SimpleNamespace(clear_continuation=lambda: None),
services=SimpleNamespace(spawn=lambda *_a, **_k: None),
run=SimpleNamespace(closing=False),
),
out,
close_from_runtime=_noop,
+123
View File
@@ -0,0 +1,123 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import json
from dataclasses import replace
import pytest
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer, DuplexOutputOverflowError
from vllm_omni.engine.duplex.events import AudioDelta, ResponseDone, SessionClosed, TranscriptDelta
from vllm_omni.engine.duplex.realtime_events import RealtimeProjectionState, project_internal_event
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("event_type", "delta_field", "content_field"),
[
("response.text.delta", "delta", "text"),
("response.output_audio.delta", "text", "transcript"),
],
)
async def test_queued_response_creation_keeps_payload_and_byte_count(event_type, delta_field, content_field):
state = RealtimeProjectionState(session_id="s")
events = project_internal_event(state, {"type": "response.created", "response_id": "r", "modalities": ["text"]})
payloads = [event.to_realtime() for event in events]
byte_count = sum(
len(json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode()) for payload in payloads
)
output = DuplexOutputBuffer(max_bytes=byte_count, max_events=len(events))
for event in events:
output.put(event)
assert output.pending_bytes == byte_count
text = "Growing response. " * 512
project_internal_event(state, {"type": event_type, "response_id": "r", delta_field: text})
item = next(iter(state.conversation_items.values()))
assert item["content"][0][content_field] == text
assert output.pending_bytes == byte_count
delivered = [await output.get() for _ in events]
assert [event.to_realtime() for event in delivered] == payloads
assert output.pending_events == output.pending_bytes == 0
@pytest.mark.asyncio
async def test_cancel_removes_only_matching_audio_and_invalidates_held_audio():
output = DuplexOutputBuffer(max_bytes=4096, max_events=8)
held = AudioDelta(response_id="old", epoch=0, delta="AAAA")
text = TranscriptDelta(response_id="old", epoch=0, delta="kept")
other = replace(held, response_id="other")
newer = replace(held, epoch=1)
done = ResponseDone(response_id="old", response={"status": "cancelled"})
for event in (held, replace(held), text, other, newer, done):
output.put(event)
assert await output.get() is held
assert output.is_valid(held)
output.invalidate("old", through_epoch=0)
assert output.pending_events == 4
with output.guard(held) as valid:
assert not valid
assert [await output.get() for _ in range(4)] == [text, other, newer, done]
assert output.pending_events == output.pending_bytes == 0
@pytest.mark.asyncio
async def test_full_reserve_cannot_block_closure_or_reorder_valid_audio():
output = DuplexOutputBuffer(max_bytes=4096, max_events=1, reserve_events=1)
audio = AudioDelta(response_id="r", epoch=0, delta="AAAA")
done = ResponseDone(response_id="r", response={"status": "completed"})
output.put(audio)
output.put(done)
with pytest.raises(DuplexOutputOverflowError):
output.put(done)
closed = SessionClosed(session_id="s", reason="client_close")
output.close(closed)
output.close(SessionClosed(reason="duplicate"))
assert not output.put(audio)
assert [await output.get() for _ in range(4)] == [audio, done, closed, None]
assert output.pending_events == output.pending_bytes == 0
@pytest.mark.asyncio
async def test_oversized_close_still_ends_stream_without_reopening_it():
output = DuplexOutputBuffer(max_bytes=128, max_events=1)
closed = SessionClosed(session_id="s", details={"large": "x" * 5000})
output.close(closed)
terminal = await output.get()
assert isinstance(terminal, SessionClosed)
assert terminal.event_id == closed.event_id
assert terminal.session_id == "s"
assert terminal.reason == "close_details_exceed_output_limit"
assert terminal.details == {}
assert await output.get() is None
output.close(closed)
assert await output.get() is None
def test_byte_limit_rejects_audio_without_changing_pending_budget():
output = DuplexOutputBuffer(max_bytes=128, max_events=8)
with pytest.raises(DuplexOutputOverflowError):
output.put(AudioDelta(delta="AAAA" * 64))
assert output.pending_events == output.pending_bytes == 0
@pytest.mark.asyncio
async def test_cancelled_waiter_can_be_replaced_and_woken_from_another_thread():
output = DuplexOutputBuffer(max_bytes=4096, max_events=8)
abandoned = asyncio.create_task(output.get())
await asyncio.sleep(0)
abandoned.cancel()
with pytest.raises(asyncio.CancelledError):
await abandoned
waiting = asyncio.create_task(output.get())
await asyncio.sleep(0)
event = AudioDelta(response_id="r", epoch=0, delta="AAAA")
await asyncio.to_thread(output.put, event)
assert await asyncio.wait_for(waiting, timeout=2) is event
closing = asyncio.create_task(output.get())
await asyncio.sleep(0)
await asyncio.to_thread(output.close)
assert await asyncio.wait_for(closing, timeout=2) is None
+237 -53
View File
@@ -25,13 +25,17 @@ from vllm_omni.engine.duplex.contracts import (
DuplexStageSubmissionResult,
duplex_resource_request_belongs_to_session,
)
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.events import (
AudioDelta,
DuplexEvent,
ErrorEvent,
ResponseDone,
SessionClosed,
SessionCreated,
SessionExpired,
SessionHeartbeatAck,
TranscriptDelta,
)
from vllm_omni.engine.duplex.messages import (
CloseDuplexSessionMessage,
@@ -297,6 +301,7 @@ class Harness:
output_sink: asyncio.Queue = field(default_factory=asyncio.Queue)
result_sink: asyncio.Queue = field(default_factory=asyncio.Queue)
_control_seq: int = 0
outputs: dict[str, DuplexOutputBuffer] = field(default_factory=dict)
@classmethod
def create(
@@ -330,11 +335,20 @@ class Harness:
await self.manager.shutdown()
async def open(self, session_id: str, config: DuplexSessionConfig | None = None) -> DuplexControlResultMessage:
limits = self.manager.runtime_config
output_buffer = self.outputs.setdefault(
session_id,
DuplexOutputBuffer(
max_bytes=limits.max_pending_output_bytes_per_session,
max_events=limits.max_pending_output_events_per_session,
),
)
await self.manager.handle(
OpenDuplexSessionMessage(
control_id=f"open-{session_id}",
session_id=session_id,
session_config=config or DuplexSessionConfig(model="fake-model"),
output_buffer=output_buffer,
)
)
return await self.result()
@@ -381,8 +395,14 @@ class Harness:
assert isinstance(result, DuplexControlResultMessage)
return result
def events(self, session_id: str | None = None) -> list[DuplexEvent]:
async def events(self, session_id: str | None = None) -> list[DuplexEvent]:
events: list[DuplexEvent] = []
for identity, output in self.outputs.items():
while output.pending_events:
event = await output.get()
assert event is not None and event.session_id == identity
if output.is_valid(event) and (session_id is None or identity == session_id):
events.append(event)
while not self.output_sink.empty():
message = self.output_sink.get_nowait()
assert isinstance(message, DuplexSessionEventMessage)
@@ -396,6 +416,22 @@ class Harness:
assert session is not None
return session
async def begin_output_response(self, session_id: str) -> str:
"""Initialize real response, text and audio projection before filling output."""
runner = self.manager.runners[session_id]
response_id = runner.session.begin_response()
runner.session.append_assistant_text("abcdefghij")
payloads: list[dict[str, object]] = [
{"type": "response.created", "modalities": ["audio", "text"]},
{"type": "response.text.delta", "delta": "abcdefghij"},
{"type": "response.output_audio.delta", "audio": "AAAAAA==", "format": "pcm16"},
]
for payload in payloads:
runner.emit({**payload, "response_id": response_id, "epoch": runner.session.epoch})
await self.events(session_id)
assert not runner.run.closing
return response_id
def capture_submissions(self, session_id: str) -> list[DuplexCommand]:
"""Replace the runner mailbox with a list so admitted commands can be inspected."""
runner = self.manager.runners[session_id]
@@ -433,7 +469,7 @@ async def test_open_answers_with_capabilities_and_emits_session_created() -> Non
assert result.public_session["voice"] == "test"
assert result.error_code is None
events = harness.events("sid-open")
events = await harness.events("sid-open")
assert isinstance(events[0], SessionCreated)
assert events[0].epoch == 0
assert events[0].session["id"] == "sid-open"
@@ -587,7 +623,7 @@ async def test_plugin_runtime_config_rejection_rolls_back_the_open(error: Except
assert harness.manager.get("sid-runtime") is None
assert harness.manager.active_count() == 0
assert harness.stage_port.ensure_calls == []
assert harness.events() == []
assert await harness.events() == []
# Nothing was registered under the rejected id and the slot is free again.
plugin.runtime_config_error = None
@@ -603,7 +639,7 @@ async def test_command_for_unknown_session_emits_unknown_session_error() -> None
async with Harness.create() as harness:
harness.command("sid-missing", Heartbeat(event_id="evt-hb"))
events = harness.events("sid-missing")
events = await harness.events("sid-missing")
assert len(events) == 1
error = events[0]
assert isinstance(error, ErrorEvent)
@@ -618,11 +654,11 @@ async def test_command_for_closed_session_emits_unknown_session_error() -> None:
async with Harness.create() as harness:
await harness.open("sid-closed")
assert (await harness.close("sid-closed")).ok is True
harness.events()
await harness.events()
harness.command("sid-closed", Heartbeat(event_id="evt-late"))
events = harness.events("sid-closed")
events = await harness.events("sid-closed")
assert [type(event) for event in events] == [ErrorEvent]
assert events[0].code == "unknown_session"
assert events[0].related_event_id == "evt-late"
@@ -631,14 +667,14 @@ async def test_command_for_closed_session_emits_unknown_session_error() -> None:
async def test_admitted_command_reaches_the_session_runner_in_order() -> None:
async with Harness.create() as harness:
await harness.open("sid-runner")
harness.events()
await harness.events()
harness.clock.advance(5.0)
harness.command("sid-runner", Heartbeat(event_id="evt-1"))
harness.command("sid-runner", Heartbeat(event_id="evt-2"))
await _settle()
events = harness.events("sid-runner")
events = await harness.events("sid-runner")
assert [type(event) for event in events] == [SessionHeartbeatAck, SessionHeartbeatAck]
assert all(event.epoch == 0 for event in events)
assert harness.session("sid-runner").lease.last_activity == 5.0
@@ -647,7 +683,7 @@ async def test_admitted_command_reaches_the_session_runner_in_order() -> None:
async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None:
async with Harness.create(max_pending_input_bytes_per_session=8) as harness:
await harness.open("sid-bytes")
harness.events()
await harness.events()
session = harness.session("sid-bytes")
submitted = harness.capture_submissions("sid-bytes")
assert session.reserve_input_bytes(6, limit=8)
@@ -655,7 +691,7 @@ async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None:
harness.command("sid-bytes", AppendAudio(audio=b"pcm", event_id="evt-over"))
harness.command("sid-bytes", AppendAudio(audio=b"pc", event_id="evt-fits"))
events = harness.events("sid-bytes")
events = await harness.events("sid-bytes")
assert len(events) == 1
assert isinstance(events[0], ErrorEvent)
assert events[0].code == "input_backpressure"
@@ -670,7 +706,7 @@ async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None:
async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues_it() -> None:
async with Harness.create(max_pending_turns_per_session=1) as harness:
await harness.open("sid-turns")
harness.events()
await harness.events()
session = harness.session("sid-turns")
submitted = harness.capture_submissions("sid-turns")
@@ -679,7 +715,7 @@ async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues
assert [command.event_id for command in submitted] == ["evt-commit-1"]
assert session.pending_input_turns == 1
events = harness.events("sid-turns")
events = await harness.events("sid-turns")
assert len(events) == 1
assert isinstance(events[0], ErrorEvent)
assert events[0].code == "input_backpressure"
@@ -689,7 +725,7 @@ async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues
session.release_pending_turn()
harness.command("sid-turns", Commit(event_id="evt-commit-3"))
assert [command.event_id for command in submitted] == ["evt-commit-1", "evt-commit-3"]
assert harness.events("sid-turns") == []
assert await harness.events("sid-turns") == []
def test_accepts_only_duplex_message_types() -> None:
@@ -704,7 +740,14 @@ def test_accepts_only_duplex_message_types() -> None:
)
try:
config = DuplexSessionConfig()
assert manager.accepts(OpenDuplexSessionMessage(control_id="c", session_id="s", session_config=config))
assert manager.accepts(
OpenDuplexSessionMessage(
control_id="c",
session_id="s",
session_config=config,
output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
)
)
assert manager.accepts(CloseDuplexSessionMessage(control_id="c", session_id="s"))
assert manager.accepts(ResumeDuplexSessionMessage(control_id="c", session_id="s", expected_lease_generation=0))
assert manager.accepts(TouchDuplexSessionMessage(control_id="c", session_id="s", activity="heartbeat"))
@@ -729,7 +772,7 @@ def test_accepts_only_duplex_message_types() -> None:
async def test_close_releases_stage_resources_and_is_idempotent() -> None:
async with Harness.create() as harness:
await harness.open("sid-close")
harness.events()
await harness.events()
session = harness.session("sid-close")
request_id = stage0_request_id("sid-close")
session.bind_stage_request(0, request_id, fence=session.fence)
@@ -750,7 +793,7 @@ async def test_close_releases_stage_resources_and_is_idempotent() -> None:
assert harness.manager.runner_for_request_id(request_id) is None
assert harness.stage_port.cleanup_calls == [([request_id], True), ([reserved_context.request_id], False)]
assert harness.plugin.data_plane.closed_sessions == ["sid-close"]
closed = [event for event in harness.events("sid-close") if isinstance(event, SessionClosed)]
closed = [event for event in await harness.events("sid-close") if isinstance(event, SessionClosed)]
assert len(closed) == 1
assert closed[0].reason == "client_close"
assert closed[0].is_terminal is True
@@ -760,7 +803,7 @@ async def test_close_releases_stage_resources_and_is_idempotent() -> None:
assert again.lease_generation is None
assert again.public_session is None
assert harness.stage_port.cleanup_calls[2:] == []
assert harness.events() == []
assert await harness.events() == []
async def test_close_retains_the_admission_slot_until_stage_cleanup_succeeds() -> None:
@@ -858,7 +901,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None:
async with Harness.create(idle_ttl_s=2.0, disconnect_grace_s=1.0) as harness:
await harness.open("sid-idle")
await harness.open("sid-active")
harness.events()
await harness.events()
harness.clock.advance(1.0)
assert (await harness.touch("sid-active", DuplexLeaseActivity.HEARTBEAT.value)).ok is True
harness.clock.advance(1.1)
@@ -868,7 +911,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None:
assert harness.manager.get("sid-idle") is None
assert harness.manager.get("sid-active") is not None
assert harness.manager.active_count() == 1
expired = harness.events("sid-idle")
expired = await harness.events("sid-idle")
assert len(expired) == 1
assert isinstance(expired[0], SessionExpired)
assert expired[0].reason == "idle_ttl_expired"
@@ -878,7 +921,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None:
"session_id": "sid-idle",
"reason": "idle_ttl_expired",
}
assert harness.events("sid-active") == []
assert await harness.events("sid-active") == []
assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-idle")], False)]
assert harness.result_sink.empty()
@@ -888,7 +931,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None:
async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() -> None:
async with Harness.create(idle_ttl_s=300.0, disconnect_grace_s=1.0) as harness:
await harness.open("sid-detached")
harness.events()
await harness.events()
assert (await harness.touch("sid-detached", DuplexLeaseActivity.DETACH.value)).ok is True
harness.clock.advance(0.5)
assert await harness.manager.reap_expired() == 0
@@ -897,7 +940,7 @@ async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() ->
assert await harness.manager.reap_expired() == 1
assert harness.manager.get("sid-detached") is None
expired = harness.events("sid-detached")
expired = await harness.events("sid-detached")
assert [type(event) for event in expired] == [SessionExpired]
assert expired[0].reason == "disconnect_grace_expired"
assert (await harness.resume("sid-detached", expected_lease_generation=0)).error_code == "unknown_session"
@@ -906,7 +949,7 @@ async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() ->
async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is_admitted() -> None:
async with Harness.create(max_sessions=1) as harness:
await harness.open("sid-closing")
harness.events()
await harness.events()
session = harness.session("sid-closing")
session.bind_stage_request(0, stage0_request_id("sid-closing"), fence=session.fence)
harness.stage_port.cleanup_gate = asyncio.Event()
@@ -915,11 +958,11 @@ async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is
await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0)
# The runner is gone but the stage request is still being aborted:
# no session.closed yet, so a client cannot race the admission slot.
assert harness.events("sid-closing") == []
assert await harness.events("sid-closing") == []
harness.stage_port.cleanup_gate.set()
assert (await close_task).ok is True
events = harness.events("sid-closing")
events = await harness.events("sid-closing")
assert [type(event) for event in events] == [SessionClosed]
assert events[0].reason == "client_close"
assert (await harness.open("sid-replacement")).ok is True
@@ -928,7 +971,7 @@ async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is
async def test_wire_close_command_frees_the_slot_before_session_closed() -> None:
async with Harness.create(max_sessions=1) as harness:
await harness.open("sid-wire-close")
harness.events()
await harness.events()
session = harness.session("sid-wire-close")
request_id = stage0_request_id("sid-wire-close")
session.bind_stage_request(0, request_id, fence=session.fence)
@@ -937,7 +980,7 @@ async def test_wire_close_command_frees_the_slot_before_session_closed() -> None
harness.command("sid-wire-close", CloseSession(reason="client_close"))
await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0)
assert harness.manager.get("sid-wire-close") is None
assert harness.events("sid-wire-close") == []
assert await harness.events("sid-wire-close") == []
harness.stage_port.cleanup_gate.set()
for _ in range(50):
@@ -945,7 +988,7 @@ async def test_wire_close_command_frees_the_slot_before_session_closed() -> None
break
await asyncio.sleep(0.01)
assert harness.stage_port.cleanup_calls == [([request_id], True)]
events = harness.events("sid-wire-close")
events = await harness.events("sid-wire-close")
assert [type(event) for event in events] == [SessionClosed]
assert events[0].reason == "client_close"
assert (await harness.open("sid-replacement")).ok is True
@@ -954,7 +997,7 @@ async def test_wire_close_command_frees_the_slot_before_session_closed() -> None
async def test_wire_close_stops_the_worker_and_refuses_later_control_ops() -> None:
async with Harness.create(max_sessions=1) as harness:
await harness.open("sid-wire")
harness.events()
await harness.events()
runner = harness.manager.runners["sid-wire"]
worker = runner._worker
assert worker is not None
@@ -975,7 +1018,7 @@ async def test_wire_close_stops_the_worker_and_refuses_later_control_ops() -> No
break
await asyncio.sleep(0.01)
assert worker.done(), "the mailbox worker must exit after a wire close"
events = harness.events("sid-wire")
events = await harness.events("sid-wire")
assert [type(event) for event in events] == [ErrorEvent, SessionClosed] or [
type(event) for event in events
] == [SessionClosed, ErrorEvent]
@@ -1003,7 +1046,7 @@ async def test_concurrent_opens_cannot_exceed_max_sessions() -> None:
async def test_runtime_close_without_stage_requests_frees_the_slot() -> None:
async with Harness.create(max_sessions=1) as harness:
await harness.open("sid-runtime")
harness.events()
await harness.events()
runner = harness.manager.runners["sid-runtime"]
runner.session.release_all_requests()
@@ -1014,7 +1057,7 @@ async def test_runtime_close_without_stage_requests_frees_the_slot() -> None:
await asyncio.sleep(0.01)
assert harness.manager.get("sid-runtime") is None
assert [type(event) for event in harness.events("sid-runtime")] == [SessionClosed]
assert [type(event) for event in await harness.events("sid-runtime")] == [SessionClosed]
assert (await harness.open("sid-next")).ok is True
@@ -1028,7 +1071,7 @@ async def test_close_of_a_forgotten_session_is_idempotent() -> None:
async def test_append_bytes_are_reserved_at_admission_until_the_runner_dequeues() -> None:
async with Harness.create(max_sessions=1, max_pending_input_bytes_per_session=8) as harness:
await harness.open("sid-bytes")
harness.events()
await harness.events()
session = harness.session("sid-bytes")
runner = harness.manager.runners["sid-bytes"]
# Park the worker so admitted commands stay in the mailbox.
@@ -1040,7 +1083,7 @@ async def test_append_bytes_are_reserved_at_admission_until_the_runner_dequeues(
harness.command("sid-bytes", AppendAudio(audio=b"12345"))
assert session.pending_input_bytes == 5
harness.command("sid-bytes", AppendAudio(audio=b"12345"))
errors = [event for event in harness.events("sid-bytes") if isinstance(event, ErrorEvent)]
errors = [event for event in await harness.events("sid-bytes") if isinstance(event, ErrorEvent)]
assert [error.code for error in errors] == ["input_backpressure"]
gate.set()
@@ -1049,7 +1092,7 @@ async def test_append_admission_counts_audio_and_video_frame_bytes() -> None:
"""Manager reserves len(audio)+Σlen(frame); video bytes count toward the same limit."""
async with Harness.create(max_sessions=1, max_pending_input_bytes_per_session=20) as harness:
await harness.open("sid-av")
harness.events()
await harness.events()
session = harness.session("sid-av")
runner = harness.manager.runners["sid-av"]
gate = asyncio.Event()
@@ -1072,7 +1115,7 @@ async def test_append_admission_counts_audio_and_video_frame_bytes() -> None:
"sid-av",
AppendAudio(audio=b"x" * 10, video_frames=("yyyyyyyyyy",), event_id="evt-over"),
)
errors = [event for event in harness.events("sid-av") if isinstance(event, ErrorEvent)]
errors = [event for event in await harness.events("sid-av") if isinstance(event, ErrorEvent)]
assert [error.code for error in errors] == ["input_backpressure"]
assert errors[0].related_event_id == "evt-over"
assert session.pending_input_bytes == expected
@@ -1082,7 +1125,7 @@ async def test_append_admission_counts_audio_and_video_frame_bytes() -> None:
async def test_expired_session_retains_the_admission_slot_until_cleanup_succeeds() -> None:
async with Harness.create(max_sessions=1, idle_ttl_s=1.0) as harness:
await harness.open("sid-expired")
harness.events()
await harness.events()
session = harness.session("sid-expired")
request_id = stage0_request_id("sid-expired")
session.bind_stage_request(0, request_id, fence=session.fence)
@@ -1093,14 +1136,14 @@ async def test_expired_session_retains_the_admission_slot_until_cleanup_succeeds
assert harness.manager.get("sid-expired") is None
assert session.state == DuplexSessionState.CLOSED
assert [type(event) for event in harness.events("sid-expired")] == [SessionExpired]
assert [type(event) for event in await harness.events("sid-expired")] == [SessionExpired]
assert (await harness.open("sid-replacement")).error_code == "resource_exhausted"
assert await harness.manager.reap_expired() == 1
assert harness.stage_port.cleanup_calls == [([request_id], True)] * 2
assert session.resource_request_ids() == []
assert harness.events("sid-expired") == []
assert await harness.events("sid-expired") == []
assert (await harness.open("sid-replacement")).ok is True
@@ -1131,13 +1174,13 @@ async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None:
async with Harness.create(stage_port=OneStuckStagePort(), idle_ttl_s=1.0) as harness:
await harness.open("sid-stuck")
await harness.open("sid-fine")
harness.events()
await harness.events()
harness.clock.advance(2.0)
assert await harness.manager.reap_expired() == 1
assert harness.manager.active_count() == 0
assert {type(event) for event in harness.events()} == {SessionExpired}
assert {type(event) for event in await harness.events()} == {SessionExpired}
assert (await harness.open("sid-fine")).ok is True
assert (await harness.open("sid-stuck")).error_code == "session_exists"
@@ -1150,7 +1193,7 @@ async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None:
async def test_close_sessions_for_request_ids_expires_the_owner_and_retries_cleanup() -> None:
async with Harness.create(max_sessions=1) as harness:
await harness.open("sid-request")
harness.events()
await harness.events()
session = harness.session("sid-request")
request_id = stage0_request_id("sid-request")
session.bind_stage_request(0, request_id, fence=session.fence)
@@ -1162,7 +1205,7 @@ async def test_close_sessions_for_request_ids_expires_the_owner_and_retries_clea
assert harness.manager.get("sid-request") is None
assert session.lease.terminal_reason == "request_cleanup"
await _settle()
expired = harness.events("sid-request")
expired = await harness.events("sid-request")
assert [type(event) for event in expired] == [SessionExpired]
assert expired[0].reason == "request_cleanup"
assert session.state == DuplexSessionState.CLOSED
@@ -1256,7 +1299,12 @@ async def test_control_dispatch_is_ordered_per_session_without_blocking_other_se
blocked_config = DuplexSessionConfig(model="fake-model", instructions="sid-blocked")
manager.dispatch(
OpenDuplexSessionMessage(control_id="blocked-open", session_id="sid-blocked", session_config=blocked_config)
OpenDuplexSessionMessage(
control_id="blocked-open",
session_id="sid-blocked",
session_config=blocked_config,
output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
)
)
await asyncio.wait_for(plugin.runtime_config_started.wait(), timeout=1.0)
manager.dispatch(
@@ -1271,6 +1319,7 @@ async def test_control_dispatch_is_ordered_per_session_without_blocking_other_se
control_id="independent-open",
session_id="sid-independent",
session_config=DuplexSessionConfig(model="fake-model"),
output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
)
)
manager.dispatch(
@@ -1338,16 +1387,18 @@ async def test_shutdown_closes_every_runner_and_stops_dispatch_tasks() -> None:
await harness.open("sid-a")
await harness.open("sid-b")
sessions = [harness.session("sid-a"), harness.session("sid-b")]
harness.events()
await harness.events()
await harness.manager.shutdown()
assert harness.manager.active_count() == 0
assert harness.manager.runners == {}
assert harness.manager._outputs == {}
assert harness.manager._output_failed == set()
assert all(session.state == DuplexSessionState.CLOSED for session in sessions)
assert harness.manager.runner_for_request_id(stage0_request_id("sid-b")) is None
# Shutdown is silent: no session.closed / session.expired is emitted.
assert harness.events() == []
assert await harness.events() == []
assert harness.stage_port.cleanup_calls == []
@@ -1414,12 +1465,12 @@ async def test_expiry_emits_a_terminal_event_even_when_a_close_deferred_it() ->
runner = harness.manager.runners["sid-deferred"]
# What the runner does for a wire close: tear down, defer the terminal.
await runner.close("client_close", emit_closed=False)
harness.events("sid-deferred")
await harness.events("sid-deferred")
await runner.expire("request_cleanup", emit_expired=True)
terminal = [
event for event in harness.events("sid-deferred") if isinstance(event, SessionClosed | SessionExpired)
event for event in await harness.events("sid-deferred") if isinstance(event, SessionClosed | SessionExpired)
]
assert len(terminal) == 1, terminal
assert isinstance(terminal[0], SessionExpired)
@@ -1466,7 +1517,7 @@ async def test_a_runtime_close_racing_a_manager_close_still_emits_one_terminal()
"""
async with Harness.create(max_sessions=1) as harness:
await harness.open("sid-race")
harness.events()
await harness.events()
session = harness.session("sid-race")
session.bind_stage_request(0, stage0_request_id("sid-race"), fence=session.fence)
runner = harness.manager.runners["sid-race"]
@@ -1486,7 +1537,9 @@ async def test_a_runtime_close_racing_a_manager_close_still_emits_one_terminal()
await harness.close("sid-race")
terminal = [event for event in harness.events("sid-race") if isinstance(event, SessionClosed | SessionExpired)]
terminal = [
event for event in await harness.events("sid-race") if isinstance(event, SessionClosed | SessionExpired)
]
assert len(terminal) == 1, terminal
@@ -1498,7 +1551,9 @@ async def test_a_terminal_event_is_emitted_once_even_if_expiry_runs_twice() -> N
await runner.expire("idle_ttl_expired", emit_expired=True)
await runner.expire("idle_ttl_expired", emit_expired=True)
terminal = [event for event in harness.events("sid-once") if isinstance(event, SessionClosed | SessionExpired)]
terminal = [
event for event in await harness.events("sid-once") if isinstance(event, SessionClosed | SessionExpired)
]
assert len(terminal) == 1, terminal
@@ -1590,12 +1645,12 @@ async def test_wire_close_of_an_idle_session_still_emits_session_closed() -> Non
"""
async with Harness.create(max_sessions=1) as harness:
await harness.open("sid-idle-close")
harness.events()
await harness.events()
harness.command("sid-idle-close", CloseSession(reason="client_close"))
events: list[DuplexEvent] = []
for _ in range(200):
events.extend(harness.events("sid-idle-close"))
events.extend(await harness.events("sid-idle-close"))
if events:
break
await asyncio.sleep(0.01)
@@ -1725,11 +1780,13 @@ async def test_an_open_cancelled_after_admission_releases_the_runner_and_its_sta
control_id="open-cancelled",
session_id="sid-cancelled",
session_config=DuplexSessionConfig(model="fake-model"),
output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
)
)
)
await asyncio.wait_for(sink.entered.wait(), timeout=1.0)
assert "sid-cancelled" in harness.manager.runners
assert "sid-cancelled" in harness.manager._outputs
assert [context.request_id for context in harness.stage_port.ensure_calls] == [
stage0_request_id("sid-cancelled")
]
@@ -1740,6 +1797,7 @@ async def test_an_open_cancelled_after_admission_releases_the_runner_and_its_sta
harness.manager._result_sink = original_sink
assert "sid-cancelled" not in harness.manager.runners
assert "sid-cancelled" not in harness.manager._outputs
assert harness.manager.active_count() == 0
assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-cancelled")], False)]
assert stage0_request_id("sid-cancelled") not in harness.manager._request_index
@@ -1764,6 +1822,7 @@ async def test_an_open_cancelled_again_during_its_rollback_leaves_the_stage_clea
control_id="open-cancelled-twice",
session_id="sid-cancelled",
session_config=DuplexSessionConfig(model="fake-model"),
output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
)
)
)
@@ -1779,6 +1838,7 @@ async def test_an_open_cancelled_again_during_its_rollback_leaves_the_stage_clea
open_task.cancel()
await asyncio.wait_for(shutdown_started.wait(), timeout=1.0)
assert "sid-cancelled" not in harness.manager._outputs
open_task.cancel() # lands inside the rollback, before the stage cleanup
with pytest.raises(asyncio.CancelledError):
await open_task
@@ -1875,3 +1935,127 @@ async def test_a_resume_replayed_under_its_control_id_reports_the_generation_it_
other = await harness.resume("sid-replay", expected_lease_generation=0, control_id="rpc-b")
assert other.ok is False and other.error_code == "session_resume_conflict"
async def test_output_overflow_closes_only_its_session_and_releases_capacity(mocker) -> None:
async with Harness.create(max_sessions=2, max_pending_output_events_per_session=8) as harness:
assert (await harness.open("slow")).ok
assert (await harness.open("normal")).ok
await harness.events()
assert not harness.manager.runners["slow"].run.closing
assert not harness.manager.runners["normal"].run.closing
slow = harness.session("slow")
normal = harness.session("normal")
slow_response = await harness.begin_output_response("slow")
normal_response = normal.begin_response()
request_id = stage0_request_id("slow")
slow.bind_request(request_id)
slow.mark_audio_sent(1000)
slow.acknowledge_playback(500, 500)
close_stream = mocker.spy(harness.plugin.data_plane, "close_stream")
for _ in range(5):
harness.manager.runners["slow"].emit(
{
"type": "response.output_audio.delta",
"response_id": slow_response,
"audio": "AAAAAA==",
"format": "pcm16",
}
)
harness.manager.emit(normal, [AudioDelta(response_id=normal_response, delta="CCCC")])
await asyncio.wait_for(asyncio.gather(*tuple(harness.manager._dispatched_control_tasks)), timeout=2.0)
events = await harness.events()
failed = [event for event in events if event.session_id == "slow"]
# Queued transcript deltas remain ordered; stale audio and the failed
# response's undelivered completion markers are discarded.
assert [type(event) for event in failed] == [TranscriptDelta] * 4 + [ErrorEvent, ResponseDone, SessionClosed]
assert failed[-3].code == "output_backpressure"
done = failed[-2]
assert isinstance(done, ResponseDone)
assert done.response_id == slow_response and done.status == "failed"
metadata = done.response["metadata"]
assert isinstance(metadata, dict) and metadata["committed"] is False
assert failed[-1].reason == "output_backpressure"
assert slow.history == ()
assert slow.active_response_id is None
assert harness.stage_port.abort_calls == [[request_id]]
close_stream.assert_called_once_with(request_id)
delivered = [event for event in events if event.session_id == "normal"]
assert len(delivered) == 1 and isinstance(delivered[0], AudioDelta)
assert delivered[0].response_id == normal_response and delivered[0].delta == "CCCC"
assert harness.manager.get("slow") is None
assert harness.manager.get("normal") is normal
assert harness.manager.active_count() == 1
assert (await harness.open("replacement")).ok
@pytest.mark.parametrize(
"capacity,end_before_emit,expected_status,closed",
[
(8, True, "failed", True),
(8, False, "failed", True),
(14, True, "completed", True),
(64, True, "completed", False),
],
ids=["already-ended-overflow", "active-overflow", "overflow-after-ending", "sufficient-capacity"],
)
async def test_completion_projection_retains_exactly_one_response_done(
capacity: int, end_before_emit: bool, expected_status: str, closed: bool
) -> None:
async with Harness.create(max_pending_output_events_per_session=capacity) as harness:
assert (await harness.open("completion")).ok
await harness.events()
response_id = await harness.begin_output_response("completion")
runner = harness.manager.runners["completion"]
for _ in range(4):
runner.emit(
{
"type": "response.output_audio.delta",
"response_id": response_id,
"audio": "AAAAAA==",
"format": "pcm16",
}
)
assert harness.outputs["completion"].pending_events == 8
runner.session.mark_audio_sent(1000)
runner.session.acknowledge_playback(500, 500)
if end_before_emit:
message = runner.session.end_response(commit_text=True)
assert message == {"role": "assistant", "content": "abcde"}
runner.session.register_history_item(f"item_{response_id}", message)
assert runner.session.history == (message,)
runner.emit(
{"type": "response.done", "response_id": response_id, "status": "completed", "committed": end_before_emit}
)
await asyncio.wait_for(asyncio.gather(*tuple(harness.manager._dispatched_control_tasks)), timeout=2.0)
events = await harness.events()
endings = [event for event in events if isinstance(event, ResponseDone)]
assert len(endings) == 1
assert endings[0].response_id == response_id and endings[0].status == expected_status
metadata = endings[0].response["metadata"]
assert isinstance(metadata, dict) and metadata["committed"] is (expected_status == "completed")
assert any(isinstance(event, SessionClosed) for event in events) is closed
if expected_status == "failed":
assert [type(event) for event in events] == [TranscriptDelta] * 4 + [
ErrorEvent,
ResponseDone,
SessionClosed,
]
assert endings[0].response["status_details"] == {"type": "failed", "reason": "output_backpressure"}
assert runner.session.history == ()
else:
assert [type(event) for event in events[:8]] == [AudioDelta, TranscriptDelta] * 4
assert runner.session.history == ({"role": "assistant", "content": "abcde"},)
if closed:
# The ending was accepted; overflowing on rate_limits.updated
# must not replace it, emit a second ending or undo its history.
assert [type(event) for event in events[-3:]] == [ResponseDone, ErrorEvent, SessionClosed]
assert events[-3] is endings[0]
else:
assert not any(isinstance(event, ErrorEvent) for event in events)
assert events[-2] is endings[0]
assert events[-1].type == "rate_limits.updated"
+84 -6
View File
@@ -6,7 +6,7 @@
The runner is driven exactly the way ``DuplexOrchestrator`` drives it: typed
commands go through ``DuplexSessionManager.dispatch``, stage outputs are pushed
with ``runner.on_stage_output`` and everything the session says is read back
from the manager's output sink as typed events. The stage port is a recording
from the session's output buffer as typed events. The stage port is a recording
fake; the model plugin is the real MiniCPM-o 4.5 one so append planning and
output projection are exercised end to end.
"""
@@ -38,7 +38,8 @@ from vllm_omni.engine.duplex.contracts import (
DuplexStageSubmissionResult,
duplex_resource_request_id,
)
from vllm_omni.engine.duplex.events import DuplexEvent
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent
from vllm_omni.engine.duplex.messages import (
CloseDuplexSessionMessage,
DuplexControlResultMessage,
@@ -72,6 +73,8 @@ class RecordingStagePort(DuplexStagePort):
self.submissions: list[DuplexStageSubmission] = []
self.cleanups: list[tuple[list[str], bool]] = []
self.aborts: list[list[str]] = []
self.abort_started = asyncio.Event()
self.abort_gate: asyncio.Event | None = None
self.fail_submit: Exception | None = None
#: When set, ``submit`` parks on it after signalling ``submit_started``.
self.submit_gate: asyncio.Event | None = None
@@ -105,6 +108,9 @@ class RecordingStagePort(DuplexStagePort):
async def abort_requests(self, request_ids: list[str]) -> None:
self.aborts.append(list(request_ids))
self.abort_started.set()
if self.abort_gate is not None:
await self.abort_gate.wait()
def _fake_encode_audio(audio: object, sample_rate_hz: int, response_format: str, speed: float | None) -> str | None:
@@ -120,6 +126,7 @@ class Harness:
manager: DuplexSessionManager
port: RecordingStagePort
output: asyncio.Queue[Any]
output_buffer: DuplexOutputBuffer
results: asyncio.Queue[Any]
runner: DuplexSessionRunner
events: list[DuplexEvent] = field(default_factory=list)
@@ -139,6 +146,12 @@ class Harness:
collected: list[DuplexEvent] = []
while True:
drained = False
while self.output_buffer.pending_events:
event = await self.output_buffer.get()
assert event is not None
if self.output_buffer.is_valid(event):
collected.append(event)
drained = True
while not self.output.empty():
message = self.output.get_nowait()
if isinstance(message, DuplexSessionEventMessage):
@@ -215,12 +228,17 @@ async def open_harness(
port = RecordingStagePort(stage_count=stage_count)
output: asyncio.Queue[Any] = asyncio.Queue()
results: asyncio.Queue[Any] = asyncio.Queue()
limits = runtime_config or DuplexSessionRuntimeConfig()
output_buffer = DuplexOutputBuffer(
max_bytes=limits.max_pending_output_bytes_per_session,
max_events=limits.max_pending_output_events_per_session,
)
manager = DuplexSessionManager(
plugin=plugin,
stage_port=port,
output_sink=output,
result_sink=results,
runtime_config=runtime_config or DuplexSessionRuntimeConfig(),
runtime_config=limits,
model_config=None,
log_stats=log_stats,
clock=clock,
@@ -232,10 +250,21 @@ async def open_harness(
instructions="You are a concise assistant.",
extra_body=body,
)
await manager.handle(OpenDuplexSessionMessage(control_id="c-open", session_id=SESSION_ID, session_config=config))
await manager.handle(
OpenDuplexSessionMessage(
control_id="c-open", session_id=SESSION_ID, session_config=config, output_buffer=output_buffer
)
)
result = await asyncio.wait_for(results.get(), timeout=2.0)
assert isinstance(result, DuplexControlResultMessage) and result.ok, result
harness = Harness(manager=manager, port=port, output=output, results=results, runner=manager.runners[SESSION_ID])
harness = Harness(
manager=manager,
port=port,
output=output,
output_buffer=output_buffer,
results=results,
runner=manager.runners[SESSION_ID],
)
await harness.settle()
return harness
@@ -710,6 +739,8 @@ async def test_stale_epoch_output_is_dropped_after_barge_in() -> None:
async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> None:
h = await open_harness()
release_abort = asyncio.Event()
h.port.abort_gate = release_abort
try:
await h.run(append_audio())
request_id = h.stage0_request_id()
@@ -730,7 +761,24 @@ async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> N
h.session.request_resources[(1, "still-live")] = DuplexRequestResource(
stage_id=1, request_id="still-live", fence=older, submitted=True
)
events = await h.run(commands.BargeIn())
h.manager.emit(
h.session,
[
AudioDelta(response_id="resp-old", delta="AAAA"),
AudioDelta(response_id="resp-old", delta="BBBB"),
AudioDelta(response_id=h.session.active_response_id, delta="CCCC"),
],
)
held = await h.output_buffer.get()
assert held is not None and h.output_buffer.is_valid(held)
h.submit(commands.BargeIn())
await asyncio.wait_for(h.port.abort_started.wait(), timeout=2.0)
# Both queued audio and an event held by the consumer become stale
# before the stage abort can yield, including the older draining turn.
assert not h.output_buffer.is_valid(held)
assert h.output_buffer.pending_events == 0
release_abort.set()
events = await h.settle()
assert h.port.aborts == [[request_id, "duplex-drain-tts"]]
assert not h.session.is_draining_request("duplex-drain-tts")
done_ids = [
@@ -744,6 +792,36 @@ async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> N
assert (0, request_id) not in h.session.request_resources
assert (1, "still-live") in h.session.request_resources
finally:
release_abort.set()
await close_harness(h)
@pytest.mark.asyncio
async def test_close_invalidates_draining_audio_before_stage_abort_finishes() -> None:
h = await open_harness()
release_abort = asyncio.Event()
h.port.abort_gate = release_abort
try:
h.session.bind_draining_request("duplex-drain-tts", "resp-old")
h.manager.emit(
h.session,
[AudioDelta(response_id="resp-old", delta="AAAA"), AudioDelta(response_id="resp-old", delta="BBBB")],
)
held = await h.output_buffer.get()
assert held is not None and h.output_buffer.is_valid(held)
h.submit(commands.CloseSession())
await asyncio.wait_for(h.port.abort_started.wait(), timeout=2.0)
# Session close suppresses audio.cancelled notifications; it must still
# invalidate the pending output without relying on those later events.
assert not h.output_buffer.is_valid(held)
assert h.output_buffer.pending_events == 0
assert h.port.aborts == [["duplex-drain-tts"]]
release_abort.set()
events = await h.settle()
assert "session.closed" in types(events)
assert h.manager.active_count() == 0
finally:
release_abort.set()
await close_harness(h)
+10 -3
View File
@@ -14,6 +14,7 @@ import pytest
from vllm_omni.engine.duplex.commands import Heartbeat
from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.messages import (
CloseDuplexSessionMessage,
DuplexControlResultMessage,
@@ -72,14 +73,16 @@ async def test_open_session_runs_the_correlated_rpc_with_the_typed_config() -> N
result = _ok("open", capabilities=DuplexCapabilities(), public_session={"id": "sid"}, lease_generation=0)
engine = _engine(result)
config = DuplexSessionConfig(model="m")
output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
returned = await engine.open_session_async("sid", config, timeout=3.0)
returned = await engine.open_session_async("sid", config, output_buffer=output_buffer, timeout=3.0)
assert returned is result
((key, message, timeout),) = engine.rpc_client.calls
assert key == ("duplex", message.control_id)
assert isinstance(message, OpenDuplexSessionMessage)
assert message.session_id == "sid" and message.session_config is config
assert message.output_buffer is output_buffer
assert timeout == 3.0
# A momentarily full request queue is backpressure, not a failed open: the
# control RPC blocks rather than raising an untyped queue.Full at the caller.
@@ -116,7 +119,9 @@ async def test_failed_control_result_maps_to_a_typed_session_error() -> None:
)
)
with pytest.raises(DuplexSessionError) as excinfo:
await engine.open_session_async("sid", DuplexSessionConfig())
await engine.open_session_async(
"sid", DuplexSessionConfig(), output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
)
assert excinfo.value.code == "resource_exhausted"
assert excinfo.value.retryable is True
assert excinfo.value.session_id == "sid"
@@ -133,7 +138,9 @@ async def test_engine_error_and_timeout_map_to_session_errors() -> None:
engine = _engine(TimeoutError("duplex open timed out for session sid"))
with pytest.raises(DuplexSessionError) as excinfo:
await engine.open_session_async("sid", DuplexSessionConfig())
await engine.open_session_async(
"sid", DuplexSessionConfig(), output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
)
assert excinfo.value.code == "timeout" and excinfo.value.retryable is True
+22 -9
View File
@@ -24,6 +24,7 @@ from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.engine.duplex import commands
from vllm_omni.engine.duplex.config import DuplexSessionConfig, DuplexSessionState
from vllm_omni.engine.duplex.contracts import DuplexFence, duplex_resource_request_id
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.messages import (
CloseDuplexSessionMessage,
DuplexControlResultMessage,
@@ -78,10 +79,15 @@ def _build(
return orchestrator, clients, rpc_q, output_q
def _open_message(extra_body: dict[str, object] | None = None) -> OpenDuplexSessionMessage:
def _open_message(
extra_body: dict[str, object] | None = None,
*,
output_buffer: DuplexOutputBuffer | None = None,
) -> OpenDuplexSessionMessage:
return OpenDuplexSessionMessage(
control_id=f"open-{SESSION_ID}",
session_id=SESSION_ID,
output_buffer=output_buffer or DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
session_config=DuplexSessionConfig(
model="openbmb/MiniCPM-o-4_5",
modalities=["text"],
@@ -369,12 +375,13 @@ async def test_forwarded_stage_requests_are_bound_and_barge_in_aborts_them() ->
@pytest.mark.asyncio
async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue() -> None:
orchestrator, _, rpc_q, output_q = _build(stages=2)
await _open(orchestrator, rpc_q)
output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
await _open(orchestrator, rpc_q, output_buffer=output_buffer)
request_id = _stage0_request_id()
await _submit(orchestrator, _append_audio())
request_state = orchestrator.request_states[request_id]
while not output_q.empty():
output_q.get_nowait()
while output_buffer.pending_events:
await output_buffer.get()
consumed = await orchestrator._intercept_stage_output(1, 0, _tts_output(request_id), request_state, None, None)
await _settle(orchestrator)
@@ -382,8 +389,9 @@ async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue
assert consumed is True
session = orchestrator.session_manager.get(SESSION_ID)
assert session is not None and session.active_response_id is not None
types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]]
types = [(await output_buffer.get()).type for _ in range(output_buffer.pending_events)]
assert "response.created" in types and "response.output_audio.delta" in types
assert output_q.empty()
orphan = DuplexOrchestratorRequestState(
request_id="duplex-s.b3RoZXI.e.0.r.stage0",
@@ -401,7 +409,8 @@ async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue
@pytest.mark.asyncio
async def test_forward_failure_closes_the_owning_session() -> None:
orchestrator, clients, rpc_q, output_q = _build(stages=2)
await _open(orchestrator, rpc_q)
output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
await _open(orchestrator, rpc_q, output_buffer=output_buffer)
request_id = _stage0_request_id()
await _submit(orchestrator, _append_audio())
request_state = orchestrator.request_states[request_id]
@@ -417,7 +426,8 @@ async def test_forward_failure_closes_the_owning_session() -> None:
assert SESSION_ID not in orchestrator.session_manager.runners
assert session.state == DuplexSessionState.CLOSED
assert orchestrator.session_manager.active_count() == 0
types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]]
types = [(await output_buffer.get()).type for _ in range(output_buffer.pending_events)]
types.extend(output_q.get_nowait().event.type for _ in range(output_q.qsize()))
# The stage failure is reported before the session expires, never after.
assert types.index("error") < types.index("session.expired")
assert types[-1] == "session.expired"
@@ -428,7 +438,8 @@ async def test_aura_forward_failure_keeps_the_session() -> None:
from vllm_omni.model_executor.models.aura_omni.duplex.plugin import AuraDuplexPlugin
orchestrator, clients, rpc_q, output_q = _build(stages=2, plugin=AuraDuplexPlugin(_encode_audio))
await _open(orchestrator, rpc_q)
output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
await _open(orchestrator, rpc_q, output_buffer=output_buffer)
request_id = next(iter(orchestrator.request_states))
await _submit(orchestrator, _append_audio())
request_state = orchestrator.request_states[request_id]
@@ -442,7 +453,8 @@ async def test_aura_forward_failure_keeps_the_session() -> None:
assert request_id not in orchestrator.request_states
assert SESSION_ID in orchestrator.session_manager.runners
assert session.state != DuplexSessionState.CLOSED
types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]]
types = [(await output_buffer.get()).type for _ in range(output_buffer.pending_events)]
types.extend(output_q.get_nowait().event.type for _ in range(output_q.qsize()))
assert "error" in types
assert "session.expired" not in types
await orchestrator.session_manager.shutdown()
@@ -604,6 +616,7 @@ async def test_turn_plugin_processes_multimodal_prompt_before_stage_submission(m
control_id="open-qwen",
session_id=SESSION_ID,
session_config=DuplexSessionConfig(model="qwen", modalities=["text", "audio"]),
output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
)
)
assert (await rpc_q.get()).ok
@@ -18,6 +18,7 @@ from fastapi import WebSocketDisconnect
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.engine.duplex import commands
from vllm_omni.engine.duplex.config import DuplexCapabilities
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent, SessionClosed, SessionCreated
from vllm_omni.engine.duplex.messages import DuplexSessionError
from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request
@@ -94,10 +95,13 @@ class FakeHandle:
self.closed = False
self.close_reasons: list[str] = []
self.commands: list[commands.DuplexCommand] = []
self._outbox: asyncio.Queue[DuplexEvent] = asyncio.Queue()
self._outbox = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
def deliver(self, event: DuplexEvent) -> None:
self._outbox.put_nowait(event)
self._outbox.put(event)
def output_guard(self, event: DuplexEvent):
return self._outbox.guard(event)
async def submit(self, command: commands.DuplexCommand) -> None:
if self.closed:
@@ -107,6 +111,10 @@ class FakeHandle:
async def events(self):
while True:
event = await self._outbox.get()
if event is None:
return
if not self._outbox.is_valid(event):
continue
yield event
if event.is_terminal:
return
@@ -9,6 +9,8 @@ import json
import pytest
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.events import AudioDelta, ResponseDone
from vllm_omni.entrypoints.duplex.session_attachment import (
DuplexEventJournal,
DuplexJournalGapError,
@@ -22,6 +24,36 @@ from vllm_omni.entrypoints.duplex.session_attachment import (
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.asyncio
@pytest.mark.parametrize("journal", [False, True])
async def test_audio_invalidated_while_waiting_for_send_is_not_sequenced_or_journaled(mocker, journal: bool):
registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096)
output = DuplexOutputBuffer(max_bytes=4096, max_events=8)
sent: list[dict[str, object]] = []
accepted = mocker.Mock()
await registry.create("s", send=mocker.AsyncMock(side_effect=sent.append), close=mocker.AsyncMock())
audio = AudioDelta(session_id="s", response_id="r", epoch=0, delta="AAAA")
output.put(audio)
assert await output.get() is audio
state = registry._sessions["s"]
async with state.outbound_lock:
delivery = asyncio.create_task(
registry.send_event(
"s", audio.to_realtime(), journal=journal, on_accepted=accepted, event_guard=lambda: output.guard(audio)
)
)
await asyncio.sleep(0)
assert not delivery.done()
output.invalidate("r", through_epoch=0)
assert await delivery is None
accepted.assert_not_called()
done = ResponseDone(response_id="r", response={"status": "cancelled"})
await registry.send_event("s", done.to_realtime())
assert [(entry["type"], entry["server_event_seq"]) for entry in sent] == [("response.done", 1)]
assert [entry.payload["event_id"] for entry in state.journal.replay_after(0)] == [done.event_id]
class _Clock:
def __init__(self) -> None:
self.value = 0.0
+35 -1
View File
@@ -15,6 +15,7 @@ from tests.entrypoints.test_omni_entrypoints import FakeAsyncOmniEngine
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.engine.duplex import commands
from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.events import (
AudioDelta,
DuplexEvent,
@@ -41,6 +42,7 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
self.duplex_session_config = DuplexSessionRuntimeConfig(max_sessions=2)
self.duplex_capabilities = DuplexCapabilities(supports_session_resume=True)
self.opened: list[tuple[str, DuplexSessionConfig]] = []
self.outputs: dict[str, DuplexOutputBuffer] = {}
self.closed: list[tuple[str, str]] = []
self.resumed: list[tuple[str, int]] = []
#: The control id of every resume RPC, in order (a replay repeats one).
@@ -64,10 +66,11 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
def _result(self, operation: str, session_id: str, **fields: Any) -> DuplexControlResultMessage:
return DuplexControlResultMessage(control_id="c", operation=operation, session_id=session_id, ok=True, **fields)
async def open_session_async(self, session_id, session_config, *, timeout=None):
async def open_session_async(self, session_id, session_config, *, output_buffer, timeout=None):
self.opened.append((session_id, session_config))
if self.open_error is not None:
raise self.open_error
self.outputs[session_id] = output_buffer
return self._result(
"open",
session_id,
@@ -119,7 +122,11 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
def emit(self, session_id: str, event: DuplexEvent) -> None:
# The session manager binds the session identity to every event it emits.
event = replace(event, session_id=session_id)
if isinstance(event, SessionClosed):
self.outputs.pop(session_id, None)
self.output_q.put(DuplexSessionEventMessage(session_id=session_id, event=event))
elif session_id in self.outputs:
self.outputs[session_id].put(event)
def _make_omni(monkeypatch: pytest.MonkeyPatch) -> tuple[DuplexOmni, FakeDuplexEngine]:
@@ -152,6 +159,7 @@ async def test_open_session_allocates_the_id_and_ignores_client_ids(monkeypatch)
assert handle.session_id != "mine"
((session_id, config),) = engine.opened
assert session_id == handle.session_id
assert engine.outputs[session_id] is handle._outbox
assert isinstance(config, DuplexSessionConfig)
assert config.model == "dummy-model" and config.instructions == "hi"
assert handle.capabilities is engine.duplex_capabilities
@@ -169,6 +177,32 @@ async def test_open_session_allocates_the_id_and_ignores_client_ids(monkeypatch)
omni.shutdown()
@pytest.mark.asyncio
async def test_events_skips_audio_invalidated_after_dequeue(monkeypatch) -> None:
omni, engine = _make_omni(monkeypatch)
try:
handle = await omni.open_session()
output = engine.outputs[handle.session_id]
original_get = output.get
async def get_then_invalidate() -> DuplexEvent | None:
event = await original_get()
if isinstance(event, AudioDelta) and event.response_id == "cancelled":
output.invalidate("cancelled", through_epoch=0)
return event
monkeypatch.setattr(output, "get", get_then_invalidate)
engine.emit(handle.session_id, AudioDelta(response_id="cancelled", epoch=0, delta="old"))
engine.emit(handle.session_id, AudioDelta(response_id="current", epoch=1, delta="new"))
events = await _collect(handle, 1)
assert isinstance(events[0], AudioDelta)
assert events[0].response_id == "current" and events[0].delta == "new"
finally:
omni.shutdown()
@pytest.mark.asyncio
async def test_open_failure_drops_the_pending_handle_and_raises_the_engine_code(monkeypatch) -> None:
omni, engine = _make_omni(monkeypatch)
@@ -27,6 +27,7 @@ from vllm_omni.engine.duplex.commands import (
UpdateSession,
)
from vllm_omni.engine.duplex.config import DuplexSessionConfig
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.messages import OpenDuplexSessionMessage
from vllm_omni.engine.duplex.plugin import DuplexRuntimeConfigError
from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
@@ -77,6 +78,7 @@ async def open_qwen():
port = RecordingStagePort(stage_count=3)
output: asyncio.Queue[Any] = asyncio.Queue()
results: asyncio.Queue[Any] = asyncio.Queue()
output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
manager = DuplexSessionManager(
plugin=plugin,
stage_port=port,
@@ -86,10 +88,21 @@ async def open_qwen():
model_config=None,
)
config = DuplexSessionConfig(model="qwen", modalities=["text", "audio"], overlap_policy="barge_in_on_speech")
await manager.handle(OpenDuplexSessionMessage(control_id="open", session_id=SESSION_ID, session_config=config))
await manager.handle(
OpenDuplexSessionMessage(
control_id="open", session_id=SESSION_ID, session_config=config, output_buffer=output_buffer
)
)
result = await results.get()
assert result.ok, result
harness = Harness(manager, port, output, results, manager.runners[SESSION_ID])
harness = Harness(
manager=manager,
port=port,
output=output,
output_buffer=output_buffer,
results=results,
runner=manager.runners[SESSION_ID],
)
await harness.settle()
return harness
@@ -218,14 +218,17 @@ def test_the_error_code_vocabulary_is_tier3() -> None:
assert not hasattr(realtime_errors, "REALTIME_ERROR_TYPES_BY_CODE")
def test_the_tier1_error_event_reports_only_openai_classes() -> None:
@pytest.mark.parametrize("code", ["resource_exhausted", "output_backpressure"])
def test_the_tier1_error_event_reports_only_openai_classes(code: str) -> None:
from vllm_omni.protocol.duplex import events as duplex_wire_events
from vllm_omni.protocol.realtime import events as realtime_events
# Tier 1 knows the envelope shape but not our code vocabulary.
assert realtime_events.ErrorEvent(code="resource_exhausted").error_type == "invalid_request_error"
assert realtime_events.ErrorEvent(code=code).error_type == "invalid_request_error"
# Tier 2 resolves it through our table.
assert duplex_wire_events.ErrorEvent(code="resource_exhausted").error_type == "rate_limit_error"
error = duplex_wire_events.ErrorEvent(code=code)
assert error.error_type == "rate_limit_error"
assert error.to_realtime()["error"]["type"] == "rate_limit_error"
@pytest.mark.parametrize("name", _TIER1_COMMANDS + _TIER2_COMMANDS + _TIER3_COMMANDS)
+4
View File
@@ -536,6 +536,8 @@ class DuplexSessionRuntimeConfig:
resume_replay_max_bytes_per_session: int = 8 * 1024 * 1024
max_pending_input_bytes_per_session: int = 16 * 1024 * 1024
max_pending_turns_per_session: int = 4
max_pending_output_bytes_per_session: int = 2 * 1024 * 1024
max_pending_output_events_per_session: int = 512
max_sessions: int = 1
# Unread by the plugin framework. It used to bound the per-session
# completed-append table that made a retried append RPC submit once; the
@@ -561,6 +563,8 @@ class DuplexSessionRuntimeConfig:
"resume_replay_max_bytes_per_session": self.resume_replay_max_bytes_per_session,
"max_pending_input_bytes_per_session": self.max_pending_input_bytes_per_session,
"max_pending_turns_per_session": self.max_pending_turns_per_session,
"max_pending_output_bytes_per_session": self.max_pending_output_bytes_per_session,
"max_pending_output_events_per_session": self.max_pending_output_events_per_session,
"max_sessions": self.max_sessions,
"completed_append_cache_size": self.completed_append_cache_size,
}
+242
View File
@@ -0,0 +1,242 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Bounded, single-consumer delivery across the engine and caller threads.
Only queued, unsequenced audio can be removed. The producer still rejects
new stale model output; this buffer does not retain a growing response-id
tombstone table. It tracks just the most recently dequeued event until the
next ``get()`` so a consumer can recheck it immediately before delivery.
Limits cover queued events, with a separate small termination reserve and
one final session-closure notification independent of that reserve. A
consumer may additionally hold one dequeued event, at most one queue budget
in size. Already journaled or delivered events are outside this buffer.
"""
from __future__ import annotations
import asyncio
import json
import threading
from collections import deque
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass, replace
from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent, ErrorEvent, ResponseDone, SessionClosed
_MAX_TERMINAL_BYTES = 4096
class DuplexOutputOverflowError(RuntimeError):
"""One session exceeded its pending output budget; valid output was not dropped."""
@dataclass(frozen=True, slots=True)
class _PendingEvent:
event: DuplexEvent
size: int
reserved: bool
class DuplexOutputBuffer:
"""Nonblocking producer, one async consumer, and response-scoped audio removal.
``get()`` releases tracking of its previous result. The consumer must
finish delivering that result before requesting another. ``guard()``
makes the final validity check and synchronous journal recording atomic
with invalidation; never await while holding that guard.
"""
def __init__(
self,
*,
max_bytes: int,
max_events: int,
reserve_bytes: int = 64 * 1024,
reserve_events: int = 8,
) -> None:
if min(max_bytes, max_events, reserve_bytes, reserve_events) <= 0:
raise ValueError("output buffer limits must be positive")
self._max_bytes = max_bytes
self._max_events = max_events
self._reserve_bytes = reserve_bytes
self._reserve_events = reserve_events
self._lock = threading.Lock()
self._pending: deque[_PendingEvent] = deque()
self._bytes = 0
self._reserved_bytes = self._reserved_events = 0
self._held: DuplexEvent | None = None
self._waiter: tuple[asyncio.AbstractEventLoop, asyncio.Future[None]] | None = None
self._closed = False
self._terminal: SessionClosed | None = None
@property
def pending_bytes(self) -> int:
with self._lock:
return self._bytes + self._reserved_bytes
@property
def pending_events(self) -> int:
with self._lock:
return len(self._pending)
def put(self, event: DuplexEvent) -> bool:
"""Append without blocking; overflow leaves the queue unchanged.
Return false for late output after closure. A session terminal closes
through its own single slot, even when the ordinary reserve is full.
Error and response terminals can use the finite reserve.
They retain FIFO order and cannot overtake still-valid media. The
caller must stop the affected session after ordinary overflow instead
of repeatedly generating errors until the reserve also overflows.
"""
if isinstance(event, SessionClosed):
self.close(event)
return True
payload = event.to_realtime()
audio_size = 0
if isinstance(event, AudioDelta):
# The producer has already encoded this as base64 ASCII. Count it
# directly without serializing the large string or decoding PCM.
audio_size = len(event.delta)
payload["delta"] = ""
size = audio_size + len(json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8"))
with self._lock:
if self._closed:
return False
reserved = self._bytes + size > self._max_bytes or (
len(self._pending) - self._reserved_events >= self._max_events
)
if reserved and (
not isinstance(event, ErrorEvent | ResponseDone)
or self._reserved_bytes + size > self._reserve_bytes
or self._reserved_events >= self._reserve_events
):
raise DuplexOutputOverflowError("duplex session pending output limit exceeded")
self._pending.append(_PendingEvent(event=event, size=size, reserved=reserved))
if reserved:
self._reserved_bytes += size
self._reserved_events += 1
else:
self._bytes += size
waiter = self._waiter
self._waiter = None
self._notify(waiter)
return True
@staticmethod
def _matches(event: DuplexEvent, response_id: str, through_epoch: int) -> bool:
return (
isinstance(event, AudioDelta)
and event.response_id == response_id
and event.epoch is not None
and event.epoch <= through_epoch
)
def invalidate(self, response_id: str, through_epoch: int) -> None:
"""Remove queued audio for one accepted cancellation.
Non-audio events and other responses retain their exact order. An
already dequeued matching event becomes invalid too. The producer
must not enqueue further stale events after this call.
"""
with self._lock:
kept: deque[_PendingEvent] = deque()
for pending in self._pending:
if self._matches(pending.event, response_id, through_epoch):
self._release(pending)
else:
kept.append(pending)
self._pending = kept
if self._held is not None and self._matches(self._held, response_id, through_epoch):
self._held = None
def _is_valid(self, event: DuplexEvent) -> bool:
return not isinstance(event, AudioDelta) or event is self._held
def is_valid(self, event: DuplexEvent) -> bool:
"""Check the current dequeued event; use ``guard`` for an atomic handoff."""
with self._lock:
return self._is_valid(event)
@contextmanager
def guard(self, event: DuplexEvent) -> Iterator[bool]:
"""Hold validity stable during synchronous sequencing; never await here."""
with self._lock:
yield self._is_valid(event)
def _release(self, pending: _PendingEvent) -> None:
if pending.reserved:
self._reserved_bytes -= pending.size
self._reserved_events -= 1
else:
self._bytes -= pending.size
async def get(self) -> DuplexEvent | None:
"""Return the next event, or ``None`` once a closed buffer is drained."""
loop = asyncio.get_running_loop()
with self._lock:
self._held = None
while True:
with self._lock:
if self._pending:
pending = self._pending.popleft()
self._release(pending)
self._held = pending.event
return pending.event
if self._closed:
if self._terminal is not None:
terminal = self._terminal
self._terminal = None
return terminal
return None
if self._waiter is not None:
raise RuntimeError("duplex output buffer already has a waiting consumer")
future: asyncio.Future[None] = loop.create_future()
self._waiter = (loop, future)
try:
await future
finally:
with self._lock:
if self._waiter is not None and self._waiter[1] is future:
self._waiter = None
def close(self, event: SessionClosed | None = None) -> None:
"""Finish once, preserving FIFO output before an optional final notification.
The single final slot cannot be consumed by errors or response endings.
Oversized close details are replaced by a compact notification retaining
the generated session/event identity and terminal type. Duplicate closure
and late output are harmless, including after a local
shutdown that already ended consumption without a session event.
"""
if event is not None:
size = len(json.dumps(event.to_realtime(), ensure_ascii=False, separators=(",", ":")).encode("utf-8"))
if size > _MAX_TERMINAL_BYTES:
event = replace(event, reason="close_details_exceed_output_limit", details={})
with self._lock:
if self._closed:
return
self._closed = True
self._terminal = event
waiter = self._waiter
self._waiter = None
self._notify(waiter)
@staticmethod
def _notify(waiter: tuple[asyncio.AbstractEventLoop, asyncio.Future[None]] | None) -> None:
if waiter is None:
return
loop, future = waiter
def wake() -> None:
if not future.done():
future.set_result(None)
try:
loop.call_soon_threadsafe(wake)
except RuntimeError:
if not loop.is_closed():
raise
+6 -2
View File
@@ -5,8 +5,9 @@
Sessions live inside the engine (``DuplexSessionRunner`` on the orchestrator
loop). The API layer only opens/closes/resumes/touches sessions through
correlated RPC and pushes ``DuplexCommand`` objects one-way; every session
output travels back as a ``DuplexSessionEventMessage``.
correlated RPC and pushes ``DuplexCommand`` objects one-way. Open binds a
bounded in-process output buffer; closure notifications and errors without
a live session travel back as ``DuplexSessionEventMessage``.
"""
from __future__ import annotations
@@ -18,6 +19,7 @@ from vllm_omni.engine.messages import EngineQueueMessage
if TYPE_CHECKING:
from vllm_omni.engine.duplex.commands import DuplexCommand
from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.events import DuplexEvent
@@ -44,6 +46,8 @@ class OpenDuplexSessionMessage(EngineQueueMessage, kw_only=True):
session_id: str
#: Already normalized by ``DuplexOmni``; the queue is in-process, so the object crosses as is.
session_config: DuplexSessionConfig
#: In-process, bounded output shared with the single consumer; no media in the engine-wide queue.
output_buffer: DuplexOutputBuffer
class CloseDuplexSessionMessage(EngineQueueMessage, kw_only=True):
+3 -1
View File
@@ -804,7 +804,9 @@ def _project(state: RealtimeProjectionState, event: dict[str, object]) -> list[D
"status": "in_progress",
"content": [],
}
state.conversation_items[item_id] = item
# Later deltas grow the state item. Keep creation events independent
# so their queued payloads retain the size measured at admission.
state.conversation_items[item_id] = {**item, "content": []}
events = [
_response_created_event(event),
*_conversation_item_added_events(state, item),
+55 -3
View File
@@ -17,13 +17,15 @@ runner injects it as ``promote_deferred_overlap``.
from __future__ import annotations
from collections.abc import Callable
from collections.abc import Callable, Mapping
from dataclasses import replace
from vllm_omni.engine.duplex.config import DuplexSessionState
from vllm_omni.engine.duplex.events import (
DOMAIN_TERMINAL_EVENTS,
MODEL_OUTPUT_EVENTS,
DuplexEvent,
ResponseDone,
error_event,
)
from vllm_omni.engine.duplex.realtime_events import (
@@ -84,8 +86,51 @@ class SessionEmitter:
# ------------------------------------------------------------------ #
def emit_events(self, events: list[DuplexEvent]) -> None:
for event in events:
self._ctx.manager.emit(self._ctx.session, event)
self._ctx.manager.emit(self._ctx.session, events)
def fail_output(self, message: str, *, response_id: str | None, terminal: ResponseDone | None) -> list[DuplexEvent]:
"""Fail undelivered output, returning recovery events without sending them again."""
session = self._ctx.session
self._ctx.manager.invalidate_output(session.session_id, response_id, through_epoch=session.epoch)
if session.active_response_id is not None:
# Close still needs the request binding to abort the real engine work.
session.end_response(commit_text=False, preserve_request=True)
events: list[DuplexEvent] = [error_event("output_backpressure", message)]
if response_id is None:
return events
# Normal completion can commit just before its projected ending is
# delivered. An ending converted to failure must undo that commit.
session.delete_history_item(f"item_{response_id}")
payload = {
"type": "response.done",
"session_id": session.session_id,
"response_id": response_id,
"epoch": session.epoch,
"committed": False,
"status": "failed",
"status_details": {"type": "failed", "reason": "output_backpressure"},
"playback": session.playback.as_dict(),
}
if terminal is None:
projected = project_internal_event(self.require_projector(), payload)
terminal = next((event for event in projected if isinstance(event, ResponseDone)), None)
else:
response = dict(terminal.response)
metadata = response.get("metadata")
response.update(
status="failed",
status_details=payload["status_details"],
metadata={**(metadata if isinstance(metadata, Mapping) else {}), **payload},
)
output = response.get("output")
if isinstance(output, list):
response["output"] = [
{**item, "status": "failed"} if isinstance(item, Mapping) else item for item in output
]
terminal = replace(terminal, response=response)
if terminal is not None:
events.append(terminal)
return events
def emit_error(
self,
@@ -110,6 +155,13 @@ class SessionEmitter:
accepted, deferred_overlap_payload = self._apply_outbound_session_event(payload)
if not accepted:
return
if payload.get("type") == "audio.cancelled":
response_id = payload.get("response_id")
cancelled_epoch = payload.get("cancelled_epoch")
if isinstance(response_id, str) and isinstance(cancelled_epoch, int):
self._ctx.manager.invalidate_output(
self._ctx.session.session_id, response_id, through_epoch=cancelled_epoch
)
self.emit_events(project_internal_event(self.require_projector(), payload))
if deferred_overlap_payload is not None and not self._ctx.run.closing:
precreate_response = self._ctx.model_state.deferred_precreate_response
+91 -13
View File
@@ -11,8 +11,9 @@ state.
Control operations (open / close / resume / touch) arrive as correlated RPC
messages and answer through ``result_sink``; session commands arrive one-way
as ``DuplexSessionCommandMessage`` and are pushed onto the runner's ordered
mailbox after backpressure admission; everything a session emits leaves
through ``output_sink`` as ``DuplexSessionEventMessage``.
mailbox after backpressure admission. Public output uses a bounded buffer
shared with that session's consumer. The engine-wide ``output_sink`` only
carries closure notifications and errors without a live session.
"""
from __future__ import annotations
@@ -34,7 +35,16 @@ from vllm_omni.engine.duplex.contracts import (
duplex_resource_request_belongs_to_session,
duplex_resource_request_id,
)
from vllm_omni.engine.duplex.events import DuplexEvent, ErrorEvent, SessionClosed, SessionExpired, error_event
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer, DuplexOutputOverflowError
from vllm_omni.engine.duplex.events import (
DuplexEvent,
ErrorEvent,
OutputAudioCleared,
ResponseDone,
SessionClosed,
SessionExpired,
error_event,
)
from vllm_omni.engine.duplex.messages import (
CloseDuplexSessionMessage,
DuplexControlResultMessage,
@@ -128,6 +138,8 @@ class DuplexSessionManager:
model_path=getattr(runtime_config, "server_vad_model_path", None)
)
self.runners: dict[str, DuplexSessionRunner] = {}
self._outputs: dict[str, DuplexOutputBuffer] = {}
self._output_failed: set[str] = set()
#: Sessions whose close/expiry began but whose stage cleanup has not finalized;
#: they keep their admission slot until finalized.
self._closing: dict[str, _PendingSessionCleanup] = {}
@@ -270,20 +282,24 @@ class DuplexSessionManager:
if runner is None:
self._emit_raw(
message.session_id,
[
self._error_event(
"unknown_session",
f"Unknown or closed duplex session: {message.session_id}",
command=command,
),
)
],
)
return
session = runner.session
if runner.closing:
self.emit(
session,
[
self._error_event(
"session_closed", f"Duplex session is closing: {session.session_id}", command=command
),
)
],
)
return
if isinstance(command, AppendAudio):
@@ -293,7 +309,7 @@ class DuplexSessionManager:
if modality_error is not None:
self.emit(
session,
self._error_event("invalid_input_modality", modality_error, command=command),
[self._error_event("invalid_input_modality", modality_error, command=command)],
)
return
limit = int(self.runtime_config.max_pending_input_bytes_per_session)
@@ -301,22 +317,26 @@ class DuplexSessionManager:
if not session.reserve_input_bytes(pending_bytes, limit=limit):
self.emit(
session,
[
self._error_event(
"input_backpressure",
"Duplex session has too many pending input bytes",
command=command,
),
)
],
)
return
elif isinstance(command, Commit):
if not session.reserve_pending_turn(limit=int(self.runtime_config.max_pending_turns_per_session)):
self.emit(
session,
[
self._error_event(
"input_backpressure",
"Duplex session has too many pending input turns",
command=command,
),
)
],
)
return
runner.submit(command)
@@ -329,18 +349,72 @@ class DuplexSessionManager:
# Emission #
# ------------------------------------------------------------------ #
def emit(self, session: DuplexEngineSession, event: DuplexEvent) -> None:
"""Bind the session identity to one typed event and push it to the engine output queue."""
self._emit_raw(session.session_id, event, epoch=session.epoch)
def emit(self, session: DuplexEngineSession, events: list[DuplexEvent]) -> None:
"""Bind identity and deliver through the session's bounded output."""
self._emit_raw(session.session_id, events, epoch=session.epoch)
def invalidate_output(self, session_id: str, response_id: str | None, *, through_epoch: int) -> None:
output = self._outputs.get(session_id)
if output is not None and response_id is not None:
output.invalidate(response_id, through_epoch=through_epoch)
def _emit_raw(
self,
session_id: str,
event: DuplexEvent,
events: list[DuplexEvent],
*,
epoch: int | None = None,
) -> None:
ending_delivered = False
for index, event in enumerate(events):
try:
self._deliver(session_id, event, epoch=epoch)
except DuplexOutputOverflowError as exc:
if session_id in self._output_failed:
return
self._output_failed.add(session_id)
runner = self.runners.get(session_id)
if runner is None:
return
# Projection has already marked this batch's ending as emitted.
# Retain it instead of recursively projecting a second ending.
terminal = next((item for item in events[index:] if isinstance(item, ResponseDone)), None)
response_id = terminal.response_id if terminal is not None else event.response_id
if response_id is None and not ending_delivered:
response_id = runner.session.active_response_id
for recovery in runner.fail_output(str(exc), response_id=response_id, terminal=terminal):
try:
self._deliver(session_id, recovery, epoch=epoch)
except DuplexOutputOverflowError:
# A full control reserve may omit an error/response ending;
# session closure still has its independent terminal slot.
pass
return
ending_delivered |= isinstance(event, ResponseDone)
def _deliver(self, session_id: str, event: DuplexEvent, *, epoch: int | None) -> None:
event = replace(event, session_id=session_id, epoch=epoch)
output = self._outputs.get(session_id)
if output is not None and not isinstance(event, SessionClosed):
if isinstance(event, OutputAudioCleared) and epoch is not None:
self.invalidate_output(session_id, event.response_id, through_epoch=epoch)
if session_id in self._output_failed:
# The failure/close path has a finite reserve; stop adding ordinary output.
if not isinstance(event, ErrorEvent) and not (
isinstance(event, ResponseDone) and event.status in {"failed", "cancelled"}
):
return
output.put(event)
return
if isinstance(event, SessionClosed):
self._outputs.pop(session_id, None)
self._output_failed.discard(session_id)
elif not isinstance(event, ErrorEvent):
# A closed/rolled-back session must not move late audio back into
# the unbounded shared queue after its output buffer is removed.
return
# Only closure notifications and errors without a live session use the
# engine-wide queue. Audio never accumulates there before the bounded buffer.
message = DuplexSessionEventMessage(session_id=session_id, event=event)
put_nowait = getattr(self._output_sink, "put_nowait", None)
if callable(put_nowait):
@@ -514,6 +588,7 @@ class DuplexSessionManager:
model_config=self.model_config,
)
self.runners[session_id] = runner
self._outputs[session_id] = message.output_buffer
# Reserve the Stage0 request resource atomically with admission.
self.ensure_stage_request(session, stage_id=0)
runner.start()
@@ -566,6 +641,7 @@ class DuplexSessionManager:
owned_runner = runner is not None and self.runners.get(session_id) is runner
if owned_runner:
self.runners.pop(session_id, None)
self._outputs.pop(session_id, None)
pending: _PendingRequestCleanup | None = None
key: tuple[str, int] | None = None
if session is not None:
@@ -686,7 +762,7 @@ class DuplexSessionManager:
terminal: DuplexEvent = (
SessionExpired(reason=reason) if kind == "expired" else SessionClosed(reason=reason)
)
self.emit(session, terminal)
self.emit(session, [terminal])
async def _finalize_pending_cleanup(
self,
@@ -943,6 +1019,8 @@ class DuplexSessionManager:
self._closing.clear()
self._session_snapshots.clear()
self._request_index.clear()
self._outputs.clear()
self._output_failed.clear()
if self._owns_executor:
self.executor.shutdown(wait=False, cancel_futures=True)
@@ -577,6 +577,8 @@ class ModelChannel:
close_reason, emitted_response = await self._send_model_output_events(
drain_result, expected_epoch=expected_epoch
)
if self._ctx.run.closing:
return
if close_reason is not None:
await self._close_from_runtime(close_reason)
return
@@ -607,7 +609,7 @@ class ModelChannel:
expected_epoch: int | None = None,
) -> tuple[str | None, bool]:
session = self._ctx.session
if expected_epoch is not None and session.epoch != expected_epoch:
if self._ctx.run.closing or (expected_epoch is not None and session.epoch != expected_epoch):
return None, False
close_reason: str | None = None
emitted_response = False
@@ -624,7 +626,7 @@ class ModelChannel:
)
emitted_response = emitted_response or did_emit
close_reason = close_reason or close_reason_for_result
if expected_epoch is not None and session.epoch != expected_epoch:
if self._ctx.run.closing or (expected_epoch is not None and session.epoch != expected_epoch):
return None, emitted_response
return close_reason, emitted_response
@@ -636,7 +638,7 @@ class ModelChannel:
str(model_result.get("error_code")),
str(model_result.get("error") or "Duplex native data-plane error"),
)
if response_id is None:
if self._ctx.run.closing or response_id is None:
return
session.end_response(commit_text=False)
self._out.emit(
@@ -734,6 +736,8 @@ class ModelChannel:
close_reason: str | None = None
emitted_response = False
self._end_active_response_before_future_model_turn(model_turn_id=model_turn_id)
if self._ctx.run.closing:
return close_reason, emitted_response
if (
session.active_response_id is not None
and model_turn_id is not None
@@ -776,10 +780,14 @@ class ModelChannel:
payload["response_id"] = response_id
self._attach_runtime_metadata(payload, model_result)
self._out.emit(payload)
if self._ctx.run.closing:
return close_reason, emitted_response
if model_result.get("abort_data_plane_request") is True and isinstance(data_plane_request_id, str):
# stage_port.abort_requests expects a list of ids; a bare str is
# iterated as characters and never matches the prewarmed binding.
await self._abort_request([data_plane_request_id], notify=False)
if self._ctx.run.closing:
return close_reason, emitted_response
if response_id is not None:
if not auto_response and self.response_continuations_remaining(response_id):
self._ctx.services.spawn(
@@ -822,7 +830,7 @@ class ModelChannel:
data_plane = self._ctx.plugin.data_plane
close_reason: str | None = None
emitted_response = False
if expected_epoch is not None and session.epoch != expected_epoch:
if self._ctx.run.closing or (expected_epoch is not None and session.epoch != expected_epoch):
return close_reason, emitted_response
data_plane_request_id = model_result.get("data_plane_request_id")
if isinstance(data_plane_request_id, str) and data_plane.is_terminal(data_plane_request_id):
@@ -932,6 +940,8 @@ class ModelChannel:
return close_reason, emitted_response
if draining_response_id is None:
self._end_active_response_before_future_model_turn(model_turn_id=model_turn_id)
if self._ctx.run.closing:
return close_reason, emitted_response
if (
draining_response_id is None
and session.active_response_id is not None
@@ -955,6 +965,8 @@ class ModelChannel:
if response_request_metrics:
created_payload["response_request_metrics"] = response_request_metrics
self._out.emit(created_payload)
if self._ctx.run.closing:
return close_reason, emitted_response
stage_metrics = model_result.get("stage_metrics")
response_stage_metrics = session.accumulate_response_stage_metrics(
stage_metrics if isinstance(stage_metrics, Mapping) else None
@@ -976,6 +988,8 @@ class ModelChannel:
response_request_metrics=response_request_metrics,
)
self._out.emit(speak_payload)
if self._ctx.run.closing:
return close_reason, emitted_response
target_id = draining_response_id if draining_response_id not in (None, session.active_response_id) else None
previous_sent_ms = session.playback_for_response(target_id).sent_ms
text_chars_before_append = len(session.assistant_transcript(target_id))
@@ -1052,6 +1066,8 @@ class ModelChannel:
response_request_metrics=response_request_metrics,
)
self._out.emit(payload)
if self._ctx.run.closing:
return close_reason, emitted_response
if (
not end_of_turn
and model_result.get("stage_role") == "tts"
@@ -1233,7 +1249,7 @@ class ModelChannel:
response_id = session.active_response_id
response_epoch = session.epoch
response_turn_id = model_turn_id if model_turn_id is not None else session.active_response_turn_id
if expected_epoch is not None and response_epoch != expected_epoch:
if self._ctx.run.closing or (expected_epoch is not None and response_epoch != expected_epoch):
return
model_state.clear_continuation()
payload: dict[str, object] = {
@@ -1246,6 +1262,8 @@ class ModelChannel:
if response_id is not None:
payload["response_id"] = response_id
self._out.emit(payload)
if self._ctx.run.closing:
return
if response_id is None:
if session.epoch == response_epoch and response_turn_id is not None:
session.complete_model_turn(response_turn_id)
@@ -1277,7 +1295,8 @@ class ModelChannel:
) -> bool:
session = self._ctx.session
stale_common_owner = (
session.state == DuplexSessionState.CLOSED
self._ctx.run.closing
or session.state == DuplexSessionState.CLOSED
or session.active_request_id != request_id
or (expected_epoch is not None and session.epoch != expected_epoch)
)
@@ -1372,6 +1391,6 @@ class ModelChannel:
except Exception as exc:
logger.exception("Failed to schedule duplex native response continuation: %s", exc)
scheduled = False
if scheduled:
if scheduled and not self._ctx.run.closing:
model_state.continuation_owner_id = owner_id
model_state.continuation_units = count + 1
+13
View File
@@ -64,6 +64,7 @@ from vllm_omni.engine.duplex.contracts import (
from vllm_omni.engine.duplex.events import (
DuplexEvent,
InputCleared,
ResponseDone,
SessionExpired,
SessionHeartbeatAck,
)
@@ -347,6 +348,15 @@ class DuplexSessionRunner:
"""Whether ``session.closed`` / ``session.expired`` already left this runner."""
return self.run.closed_emitted
def fail_output(self, message: str, *, response_id: str | None, terminal: ResponseDone | None) -> list[DuplexEvent]:
"""Experimental overflow policy: fail this response and close only its session."""
if self.closing:
return []
events = self.out.fail_output(message, response_id=response_id, terminal=terminal)
self._begin_close("output_backpressure")
self.manager.close_from_runner(self, "output_backpressure")
return events
def mark_closed_emitted(self) -> None:
"""Claim the session's one terminal event for the caller.
@@ -1474,6 +1484,9 @@ class DuplexSessionRunner:
if older_abort_ids:
session.release_resources_for_request_ids(older_abort_ids)
new_epoch, old_playback = helpers.advance_barge_in_epoch(session)
self.manager.invalidate_output(session.session_id, old_response_id, through_epoch=old_epoch)
for response_id, _, _ in draining_cancels:
self.manager.invalidate_output(session.session_id, response_id, through_epoch=old_epoch)
if old_request_id is not None:
# Release projector/parser cursors so cancelled epochs do not
# accumulate until the whole session closes.
+8 -1
View File
@@ -37,6 +37,7 @@ from vllm_omni.engine.orchestrator import OrchestratorBase
if TYPE_CHECKING:
from vllm_omni.engine.duplex.commands import DuplexCommand
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
logger = init_logger(__name__)
@@ -151,6 +152,7 @@ class DuplexOmniEngine(AsyncOmniEngine):
session_id: str,
session_config: DuplexSessionConfig,
*,
output_buffer: DuplexOutputBuffer,
timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
) -> DuplexControlResultMessage:
control_id = uuid.uuid4().hex
@@ -159,6 +161,7 @@ class DuplexOmniEngine(AsyncOmniEngine):
control_id=control_id,
session_id=session_id,
session_config=session_config,
output_buffer=output_buffer,
),
control_id=control_id,
operation="open",
@@ -171,10 +174,14 @@ class DuplexOmniEngine(AsyncOmniEngine):
session_id: str,
session_config: DuplexSessionConfig,
*,
output_buffer: DuplexOutputBuffer,
timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
) -> DuplexControlResultMessage:
loop = asyncio.get_running_loop()
return await loop.run_in_executor(None, lambda: self._open_session(session_id, session_config, timeout=timeout))
return await loop.run_in_executor(
None,
lambda: self._open_session(session_id, session_config, output_buffer=output_buffer, timeout=timeout),
)
def _close_session(
self,
+8 -4
View File
@@ -22,6 +22,7 @@ import json
from collections.abc import Mapping
from contextlib import suppress
from dataclasses import dataclass, replace
from functools import partial
from fastapi import WebSocket, WebSocketDisconnect
from vllm.logger import init_logger
@@ -478,7 +479,7 @@ class OmniDuplexSessionHandler:
attachment_generation=credentials.attachment_generation,
resume_token=credentials.resume_token,
)
await self._send_event(session_id, event)
await self._send_event(session_id, event, handle=handle)
if isinstance(event, SessionClosed):
close_reason = event.reason or event.type
break
@@ -523,19 +524,22 @@ class OmniDuplexSessionHandler:
with suppress(Exception):
await asyncio.wait_for(asyncio.shield(pump), _PUMP_DRAIN_TIMEOUT_S)
async def _send_event(self, session_id: str, event: DuplexEvent) -> None:
async def _send_event(self, session_id: str, event: DuplexEvent, *, handle: DuplexSessionHandle) -> None:
payload = event.to_realtime()
journal = not isinstance(event, _UNJOURNALED_EVENTS) and session_id not in self._resync_required_sessions
event_guard = partial(handle.output_guard, event)
try:
try:
await self._attachment_registry.send_event(session_id, payload, journal=journal)
await self._attachment_registry.send_event(
session_id, payload, journal=journal, event_guard=event_guard
)
except DuplexJournalOverflowError:
first_overflow = session_id not in self._resync_required_sessions
self._resync_required_sessions.add(session_id)
if first_overflow:
resync = SessionResyncRequired(session_id=session_id, reason="journal_overflow")
await self._attachment_registry.send_event(session_id, resync.to_realtime(), journal=False)
await self._attachment_registry.send_event(session_id, payload, journal=False)
await self._attachment_registry.send_event(session_id, payload, journal=False, event_guard=event_guard)
except KeyError:
# Attachment already closed (takeover or teardown); the journal is gone.
pass
@@ -11,6 +11,7 @@ import secrets
import time
from collections import deque
from collections.abc import Awaitable, Callable, Mapping
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, field
from types import MappingProxyType
@@ -291,6 +292,7 @@ class DuplexSessionAttachmentRegistry:
*,
journal: bool = True,
on_accepted: Callable[[], None] | None = None,
event_guard: Callable[[], AbstractContextManager[bool]] | None = None,
) -> JournalEntry | None:
"""Sequence and dispatch one event to the current attachment.
@@ -309,6 +311,12 @@ class DuplexSessionAttachmentRegistry:
async with self._lock:
if self._sessions.get(session_id) is not state:
raise KeyError(f"unknown duplex attachment session: {session_id}")
# The producer may invalidate held audio while this coroutine
# waits for an attachment lock. Commit order before sequencing;
# release the cross-thread guard before awaiting network I/O.
with event_guard() if event_guard is not None else nullcontext(True) as valid:
if not valid:
return None
entry = state.journal.record(payload) if journal else None
attachment = state.attachment
wire_payload = dict(entry.payload) if entry is not None else dict(payload)
+18 -8
View File
@@ -27,7 +27,7 @@ from __future__ import annotations
import asyncio
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
from contextlib import suppress
from contextlib import AbstractContextManager, suppress
from typing import Any
from uuid import uuid4
@@ -37,6 +37,7 @@ from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.engine.duplex import commands as duplex_commands
from vllm_omni.engine.duplex.commands import DuplexCommand
from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig, ResponseCreateOptions
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.events import DuplexEvent, SessionClosed
from vllm_omni.engine.duplex.messages import (
DuplexControlResultMessage,
@@ -66,7 +67,11 @@ class DuplexSessionHandle:
self.capabilities: DuplexCapabilities = DuplexCapabilities()
self.public_session: dict[str, object] = {}
self.lease_generation: int = 0
self._outbox: asyncio.Queue[DuplexEvent | None] = asyncio.Queue()
limits = omni.duplex_session_config
self._outbox = DuplexOutputBuffer(
max_bytes=limits.max_pending_output_bytes_per_session,
max_events=limits.max_pending_output_events_per_session,
)
self._closed = False
self._close_reason: str | None = None
self._closed_event = asyncio.Event()
@@ -219,12 +224,10 @@ class DuplexSessionHandle:
self._consumer_active = True
try:
while True:
if self._closed and self._outbox.empty():
return
event = await self._outbox.get()
if event is None:
if self._closed and self._outbox.empty():
return
if not self._outbox.is_valid(event):
continue
yield event
if event.is_terminal:
@@ -236,6 +239,10 @@ class DuplexSessionHandle:
await self._closed_event.wait()
return self._close_reason or "closed"
def output_guard(self, event: DuplexEvent) -> AbstractContextManager[bool]:
"""Recheck held audio while committing its delivery; never await inside this guard."""
return self._outbox.guard(event)
async def __aenter__(self) -> DuplexSessionHandle:
return self
@@ -256,7 +263,7 @@ class DuplexSessionHandle:
self.lease_generation = int(result.lease_generation)
def _deliver(self, event: DuplexEvent) -> None:
self._outbox.put_nowait(event)
self._outbox.put(event)
if isinstance(event, SessionClosed): # SessionExpired is a SessionClosed
self._mark_closed(event.reason or event.type)
@@ -266,7 +273,7 @@ class DuplexSessionHandle:
self._closed = True
self._close_reason = reason
self._closed_event.set()
self._outbox.put_nowait(None)
self._outbox.close()
class DuplexOmni(AsyncOmni):
@@ -356,8 +363,11 @@ class DuplexOmni(AsyncOmni):
self._handles[session_id] = handle
self._final_output_handler()
try:
result = await self.engine.open_session_async(session_id, session_config, timeout=timeout)
result = await self.engine.open_session_async(
session_id, session_config, output_buffer=handle._outbox, timeout=timeout
)
except BaseException:
handle._mark_closed("open_failed")
if self._handles.get(session_id) is handle:
self._handles.pop(session_id, None)
# Dropping the handle only forgets the id here. The engine may still
+1
View File
@@ -45,6 +45,7 @@ REALTIME_ERROR_TYPES_BY_CODE: dict[str, str] = {
"runtime_touch_failed": "server_error",
"engine_error": "server_error",
"input_backpressure": "rate_limit_error",
"output_backpressure": "rate_limit_error",
"response_already_active": "invalid_request_error",
"response_not_active": "invalid_request_error",
"response_create_without_input": "invalid_request_error",