[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:
1 parent
3f20252151
commit
f781f7bb08
26 files changed
+1151
-128
No files matched your search
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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
|
||||
@@ -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):
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in new issue
Block a user