[Full Duplex] Feat: Support Full-Duplex realtime runtime & add MiniCPM-o 4.5 demo (#3907)
Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: lishunyang12 <lishunyang03@gmail.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: lishunyang12 <lishunyang03@gmail.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: StormMelody <22371184@zju.edu.cn> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
142 files changed
+48389
-1013
No files matched your search
@@ -1,159 +1,139 @@
|
||||
# MiniCPM-o 4.5: Online serving
|
||||
|
||||
OpenAI-compatible `/v1/chat/completions` serving for **MiniCPM-o 4.5**, plus a
|
||||
Gradio UI and curl / Python clients.
|
||||
This directory contains MiniCPM-o 4.5 online serving demos for vLLM-Omni.
|
||||
Inputs can include text, image, audio, or video; outputs are text and optional
|
||||
24 kHz speech.
|
||||
|
||||
Inputs: text / image / audio / video. Outputs: text and optional **24 kHz** speech.
|
||||
For the experimental native duplex runtime architecture, lifecycle invariants,
|
||||
capability boundary, and validation scope, see
|
||||
[`vllm_omni/experimental/fullduplex/DESIGN.md`](../../../vllm_omni/experimental/fullduplex/DESIGN.md).
|
||||
|
||||
## Installation
|
||||
|
||||
Please refer to [README.md](../../../README.md). Install the talker extra:
|
||||
Install vLLM-Omni with the MiniCPM-o talker dependencies:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[minicpmo]'
|
||||
|
||||
# From a source checkout:
|
||||
pip install -e '.[minicpmo]'
|
||||
```
|
||||
|
||||
## Launch the server
|
||||
The `minicpmo` extra installs `stepaudio2-minicpmo` and its audio dependencies,
|
||||
including `librosa`.
|
||||
|
||||
The deploy config auto-loads via `--omni`; the default
|
||||
`vllm_omni/deploy/minicpmo_4_5.yaml` targets a single-GPU layout (thinker
|
||||
and talker + t2w co-located on GPU 0). For other hardware layouts pick
|
||||
one of the deploy variants below.
|
||||
## Start the backend server
|
||||
|
||||
| deploy config | GPUs | Notes |
|
||||
|---|---|---|
|
||||
| `minicpmo_4_5.yaml` (default) | 1 | Thinker and talker+t2w co-located on GPU0. |
|
||||
| `minicpmo_4_5_2gpu.yaml` | 2 | Thinker on GPU0, talker+t2w on GPU1. |
|
||||
| `minicpmo_4_5_3gpu.yaml` | 3 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. |
|
||||
| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
|
||||
Pick a deploy config that matches your GPU layout:
|
||||
|
||||
Default (single-GPU):
|
||||
| config | GPUs | TP | Notes |
|
||||
|---|---|---|---|
|
||||
| `minicpmo_4_5.yaml` | 1 | 1 | Thinker and talker+t2w co-located on GPU0. |
|
||||
| `minicpmo_4_5_2gpu.yaml` | 2 | 1 | Thinker on GPU0, talker+t2w on GPU1. |
|
||||
| `minicpmo_4_5_3gpu.yaml` | 3 | 2 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. |
|
||||
| `minicpmo_4_5_8x4090.yaml` | 8 | 4 | Thinker 4-way TP on GPU0-3, talker+t2w on GPU4. |
|
||||
| `minicpmo_4_5_3gpu_stage1_replicas.yaml` | 3 | 1 | Experimental validation profile: Thinker on GPU0, two talker+Token2wav replicas on GPU1/2. |
|
||||
| `minicpmo_4_5_4gpu_stage1_replicas.yaml` | 4 | 1 | Experimental validation profile: Thinker on GPU0, three talker+Token2wav replicas on GPU1/2/3. |
|
||||
| `minicpmo_4_5_8x4090_stage1_replicas.yaml` | 8 | 4 | Experimental validation profile: Thinker 4-way TP on GPU0-3, four talker+Token2wav replicas on GPU4-7. |
|
||||
| `minicpmo_4_5_duplex.yaml` | 2 | 1 | Experimental native duplex profile. |
|
||||
|
||||
```bash
|
||||
vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
||||
vllm-omni serve openbmb/MiniCPM-o-4_5 \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \
|
||||
--trust-remote-code \
|
||||
--host 0.0.0.0 --port 8099
|
||||
```
|
||||
|
||||
Other layouts:
|
||||
For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the
|
||||
checkpoint path. To start the experimental native duplex backend, use
|
||||
`vllm_omni/deploy/minicpmo_4_5_duplex.yaml`.
|
||||
|
||||
```bash
|
||||
vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5_3gpu.yaml \
|
||||
--trust-remote-code \
|
||||
--host 0.0.0.0 --port 8099
|
||||
```
|
||||
The `*_stage1_replicas.yaml` files exercise composite Stage1 replica routing
|
||||
and failure recovery. They are validation profiles, not recommended production
|
||||
entrypoints.
|
||||
|
||||
### Stage-based CLI (optional)
|
||||
## Send chat requests
|
||||
|
||||
Stage 0 (thinker + API) and stage 1 (talker) can run in separate processes:
|
||||
|
||||
```bash
|
||||
# Stage 0
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
||||
--trust-remote-code --port 8099 --stage-id 0 \
|
||||
--omni-master-address 127.0.0.1 --omni-master-port 26000
|
||||
|
||||
# Stage 1 (headless)
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
||||
--trust-remote-code --stage-id 1 --headless \
|
||||
--omni-master-address 127.0.0.1 --omni-master-port 26000
|
||||
```
|
||||
|
||||
### Per-stage overrides
|
||||
|
||||
```bash
|
||||
vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code --port 8099 \
|
||||
--stage-overrides '{"0": {"gpu_memory_utilization": 0.65}}'
|
||||
```
|
||||
|
||||
## Send multimodal requests
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/minicpmo
|
||||
```
|
||||
|
||||
### curl
|
||||
From this directory, run the MiniCPM-specific curl or Python clients:
|
||||
|
||||
```bash
|
||||
bash run_curl_multimodal_generation.sh text
|
||||
bash run_curl_multimodal_generation.sh use_image
|
||||
bash run_curl_multimodal_generation.sh use_audio '["text"]' # text-only
|
||||
```
|
||||
bash run_curl_multimodal_generation.sh use_audio '["text"]'
|
||||
|
||||
Text + speech smoke test (TTS needs top-level `chat_template_kwargs`):
|
||||
|
||||
```bash
|
||||
curl http://localhost:8099/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}],
|
||||
"modalities": ["text", "audio"],
|
||||
"chat_template_kwargs": {"use_tts_template": true}
|
||||
}'
|
||||
```
|
||||
|
||||
### Python client
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type use_image \
|
||||
--port 8099 \
|
||||
--host localhost
|
||||
|
||||
# Text-only (faster; no <|tts_bos|>)
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type text \
|
||||
--modalities text \
|
||||
--prompt "Briefly introduce yourself."
|
||||
```
|
||||
|
||||
Shared helpers also work if you pass MiniCPM defaults yourself:
|
||||
|
||||
```bash
|
||||
python ../openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--model openbmb/MiniCPM-o-4_5 \
|
||||
--query-type text \
|
||||
--host localhost \
|
||||
--port 8099
|
||||
```
|
||||
|
||||
(Note: the shared client does **not** set `use_tts_template`; prefer the
|
||||
MiniCPM-specific client above for speech.)
|
||||
Speech output requires `chat_template_kwargs.use_tts_template=true`. Put that
|
||||
field at the request root for curl; the OpenAI Python SDK can merge it from
|
||||
`extra_body`.
|
||||
|
||||
### Gradio demo
|
||||
## Launch the Gradio demo
|
||||
|
||||
```bash
|
||||
bash run_gradio_demo.sh
|
||||
bash examples/online_serving/minicpmo/run_gradio_demo.sh
|
||||
|
||||
# Or:
|
||||
python gradio_demo.py \
|
||||
# Or run the Python entry point directly:
|
||||
python examples/online_serving/minicpmo/gradio_demo.py \
|
||||
--minicpmo45-api-base http://localhost:8099/v1 \
|
||||
--minicpmo45-model openbmb/MiniCPM-o-4_5 \
|
||||
--port 7862
|
||||
```
|
||||
|
||||
Open `http://<host>:7862`. Uncheck **"Generate speech output (TTS)"** for
|
||||
text-only responses.
|
||||
Open `http://<host>:7862` in a browser.
|
||||
|
||||
## Modality control
|
||||
## Run the Realtime duplex CLI demo
|
||||
|
||||
| Modalities | Output |
|
||||
|---|---|
|
||||
| `["text"]` | Text only (no TTS bos) |
|
||||
| `["text", "audio"]` / unset | Text + 24 kHz speech |
|
||||
After the duplex backend is running, stream one WAV through the Realtime
|
||||
WebSocket endpoint:
|
||||
|
||||
Speech requires `chat_template_kwargs.use_tts_template=true` so the chat
|
||||
template appends `<|tts_bos|>`. For **curl**, put that field at the request
|
||||
root; nested `extra_body` is ignored. The OpenAI Python SDK may use
|
||||
`extra_body` because it merges those fields into the root.
|
||||
```bash
|
||||
python examples/online_serving/minicpmo/realtime_duplex_demo.py \
|
||||
--url ws://localhost:8099/v1/realtime?duplex=1 \
|
||||
--model openbmb/MiniCPM-o-4_5 \
|
||||
--input-wav /path/to/input_16k_mono_pcm16.wav \
|
||||
--ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \
|
||||
--output-dir /tmp/minicpmo_realtime_duplex_demo
|
||||
```
|
||||
|
||||
## Notes
|
||||
## Open the experimental browser client
|
||||
|
||||
- Stage 1 is capped at `max_num_seqs: 1` in the deploy YAML (talker shares
|
||||
request-0 audio metadata).
|
||||
- Output audio is base64 WAV in `message.audio.data` (24 kHz mono).
|
||||
- Offline counterpart:
|
||||
[`examples/offline_inference/minicpmo/`](../../offline_inference/minicpmo/)
|
||||
- Recipe:
|
||||
[`recipes/OpenBMB/MiniCPM-o-4_5.md`](../../../recipes/OpenBMB/MiniCPM-o-4_5.md)
|
||||
The browser UI serves the page and proxies the same-origin Realtime WebSocket to
|
||||
the backend:
|
||||
|
||||
```bash
|
||||
python -m examples.online_serving.minicpmo.realtime_web \
|
||||
--port 7862 \
|
||||
--ws-backend ws://127.0.0.1:8099 \
|
||||
--ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav
|
||||
```
|
||||
|
||||
Open `http://<host>:7862/`. When using a reverse proxy, open the URL mapped to
|
||||
port `7862`; the browser derives its WebSocket endpoint relative to that URL.
|
||||
|
||||
If the page proxy serves HTTP but does not forward WebSocket upgrades, point the
|
||||
browser at a separately exposed Realtime endpoint:
|
||||
|
||||
```bash
|
||||
python -m examples.online_serving.minicpmo.realtime_web \
|
||||
--port 7862 \
|
||||
--ws-backend ws://127.0.0.1:8099 \
|
||||
--public-realtime-url wss://public.example/v1/realtime
|
||||
```
|
||||
|
||||
## Validate soft-interrupt behavior
|
||||
|
||||
The soft-interrupt E2E driver defaults to `--validation-mode model-policy`,
|
||||
which checks lifecycle and streaming invariants for arbitrary input audio. The
|
||||
stronger `response-required` mode is diagnostic: it requires a purpose-built
|
||||
two-response WAV, its `--input-sha256`, and an
|
||||
`--expect-second-response-substring` value.
|
||||
|
||||
## Related examples
|
||||
|
||||
- [Offline MiniCPM-o inference](../../offline_inference/minicpmo/)
|
||||
- [MiniCPM-o 4.5 recipe](../../../recipes/OpenBMB/MiniCPM-o-4_5.md)
|
||||
@@ -0,0 +1,399 @@
|
||||
"""Minimal single-input MiniCPM-o 4.5 Realtime duplex demo.
|
||||
|
||||
Run this after starting the duplex server. Strict lifecycle, overlap, and
|
||||
multi-session validation lives under ``tests/e2e/online_serving``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||
PCM16_BYTES_PER_SAMPLE,
|
||||
PCM16_SAMPLE_RATE,
|
||||
RealtimeDuplexClient,
|
||||
RealtimeEventCollector,
|
||||
build_realtime_url,
|
||||
read_pcm16_wav,
|
||||
wait_for,
|
||||
write_pcm16_wav,
|
||||
)
|
||||
|
||||
|
||||
class _StreamingOutputWriter:
|
||||
"""Persist and report output deltas as the client receives them."""
|
||||
|
||||
def __init__(self, output_dir: Path) -> None:
|
||||
self.output_dir = output_dir
|
||||
self.audio_chunk_dir = output_dir / "audio_chunks"
|
||||
self.audio_chunk_paths: list[Path] = []
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.audio_chunk_dir.mkdir(parents=True, exist_ok=True)
|
||||
(self.output_dir / "output.pcm").write_bytes(b"")
|
||||
|
||||
def handle(self, event: dict[str, object]) -> None:
|
||||
event_type = event.get("type")
|
||||
if event_type in {
|
||||
"response.audio_transcript.delta",
|
||||
"response.output_text.delta",
|
||||
}:
|
||||
delta = event.get("delta")
|
||||
if isinstance(delta, str) and delta:
|
||||
print(delta, end="", file=sys.stderr, flush=True)
|
||||
return
|
||||
if event_type != "response.audio.delta":
|
||||
return
|
||||
|
||||
delta = event.get("delta") or event.get("audio")
|
||||
if not isinstance(delta, str) or not delta:
|
||||
return
|
||||
try:
|
||||
pcm16 = base64.b64decode(delta)
|
||||
except ValueError:
|
||||
return
|
||||
if not pcm16:
|
||||
return
|
||||
|
||||
chunk_index = len(self.audio_chunk_paths) + 1
|
||||
chunk_path = self.audio_chunk_dir / f"chunk_{chunk_index:04d}.wav"
|
||||
sample_rate_hz = event.get("sample_rate_hz")
|
||||
if not isinstance(sample_rate_hz, int) or sample_rate_hz <= 0:
|
||||
sample_rate_hz = 24_000
|
||||
with (self.output_dir / "output.pcm").open("ab") as output_pcm:
|
||||
output_pcm.write(pcm16)
|
||||
write_pcm16_wav(chunk_path, pcm16, sample_rate_hz=sample_rate_hz)
|
||||
self.audio_chunk_paths.append(chunk_path)
|
||||
print(
|
||||
f"\n[audio chunk {chunk_index}: {len(pcm16)} bytes -> {chunk_path}]",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
class _StreamingEventCollector(RealtimeEventCollector):
|
||||
def __init__(self, writer: _StreamingOutputWriter) -> None:
|
||||
super().__init__()
|
||||
self._writer = writer
|
||||
|
||||
def add(self, event: dict[str, object], *, received_at_s: float | None = None) -> None:
|
||||
super().add(event, received_at_s=received_at_s)
|
||||
self._writer.handle(self.events[-1])
|
||||
|
||||
|
||||
def _input_committed_index(
|
||||
events: list[dict[str, object]],
|
||||
after_index: int,
|
||||
) -> int | None:
|
||||
for index, event in enumerate(events[max(after_index, 0) :], start=max(after_index, 0)):
|
||||
if event.get("type") == "input_audio_buffer.committed":
|
||||
return index
|
||||
return None
|
||||
|
||||
|
||||
def _post_commit_model_decision(
|
||||
events: list[dict[str, object]],
|
||||
committed_index: int | None,
|
||||
) -> str | None:
|
||||
if committed_index is None:
|
||||
return None
|
||||
for event in events[committed_index + 1 :]:
|
||||
event_type = event.get("type")
|
||||
if event_type == "response.listen":
|
||||
return "listen"
|
||||
if event_type == "response.done":
|
||||
response = event.get("response")
|
||||
if not isinstance(response, dict) or response.get("status") != "cancelled":
|
||||
return "speak"
|
||||
return None
|
||||
|
||||
|
||||
def _latest_model_decision(
|
||||
events: list[dict[str, object]],
|
||||
after_index: int,
|
||||
) -> str | None:
|
||||
decision: str | None = None
|
||||
for event in events[max(after_index, 0) :]:
|
||||
event_type = event.get("type")
|
||||
if event_type == "response.listen":
|
||||
decision = "listen"
|
||||
elif event_type == "response.done":
|
||||
response = event.get("response")
|
||||
if not isinstance(response, dict) or response.get("status") != "cancelled":
|
||||
decision = "speak"
|
||||
return decision
|
||||
|
||||
|
||||
def _chunk_period_ms(events: list[dict[str, object]]) -> int:
|
||||
for event in reversed(events):
|
||||
session = event.get("session")
|
||||
if not isinstance(session, dict):
|
||||
continue
|
||||
capabilities = session.get("capabilities")
|
||||
if not isinstance(capabilities, dict):
|
||||
continue
|
||||
chunk_period_ms = capabilities.get("chunk_period_ms")
|
||||
if isinstance(chunk_period_ms, int) and chunk_period_ms > 0:
|
||||
return chunk_period_ms
|
||||
return 1000
|
||||
|
||||
|
||||
def _has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool:
|
||||
unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000
|
||||
return bool(unit_bytes > 0 and len(pcm16) % unit_bytes)
|
||||
|
||||
|
||||
def _response_in_progress(events: list[dict[str, object]]) -> bool:
|
||||
return sum(event.get("type") == "response.created" for event in events) > sum(
|
||||
event.get("type") == "response.done" for event in events
|
||||
)
|
||||
|
||||
|
||||
def _event_count_after(
|
||||
events: list[dict[str, object]],
|
||||
event_type: str,
|
||||
index: int | None,
|
||||
) -> int:
|
||||
if index is None:
|
||||
return 0
|
||||
return sum(event.get("type") == event_type for event in events[index + 1 :])
|
||||
|
||||
|
||||
def _ref_audio_data_url(path: str | None) -> str | None:
|
||||
if path is None:
|
||||
return None
|
||||
ref_path = Path(path).expanduser()
|
||||
return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii")
|
||||
|
||||
|
||||
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
input_pcm16 = read_pcm16_wav(Path(args.input_wav))
|
||||
if not input_pcm16:
|
||||
raise ValueError("input WAV has no audio")
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
stream_writer = _StreamingOutputWriter(output_dir)
|
||||
|
||||
url = build_realtime_url(
|
||||
args.url,
|
||||
args.model,
|
||||
autostart=False if args.ref_audio else None,
|
||||
session_id=args.session_id,
|
||||
)
|
||||
client = RealtimeDuplexClient(url)
|
||||
client.events = _StreamingEventCollector(stream_writer)
|
||||
async with client:
|
||||
await client.configure(
|
||||
args.model,
|
||||
ref_audio=_ref_audio_data_url(args.ref_audio),
|
||||
session_id=args.session_id,
|
||||
timeout_s=args.timeout_s,
|
||||
)
|
||||
stream_event_cursor = len(client.events.events)
|
||||
await client.stream_pcm16(
|
||||
input_pcm16,
|
||||
chunk_ms=args.chunk_ms,
|
||||
realtime=not args.no_realtime_pacing,
|
||||
)
|
||||
commit_event_cursor = len(client.events.events)
|
||||
stream_decision = _latest_model_decision(client.events.events, stream_event_cursor)
|
||||
input_has_residual_model_unit = _has_residual_model_unit(
|
||||
input_pcm16,
|
||||
chunk_period_ms=_chunk_period_ms(client.events.events),
|
||||
)
|
||||
wait_for_post_commit_decision = False
|
||||
commit_sent_at_s = time.monotonic()
|
||||
await client.commit()
|
||||
wait_error: str | None = None
|
||||
committed_index: int | None = None
|
||||
post_commit_decision: str | None = None
|
||||
try:
|
||||
await wait_for(
|
||||
lambda: _input_committed_index(client.events.events, commit_event_cursor) is not None,
|
||||
timeout_s=args.timeout_s,
|
||||
label="input_audio_buffer.committed",
|
||||
)
|
||||
committed_index = _input_committed_index(client.events.events, commit_event_cursor)
|
||||
stream_decision = _latest_model_decision(client.events.events[: committed_index + 1], stream_event_cursor)
|
||||
wait_for_post_commit_decision = input_has_residual_model_unit or _response_in_progress(
|
||||
client.events.events[: committed_index + 1]
|
||||
)
|
||||
if wait_for_post_commit_decision:
|
||||
await wait_for(
|
||||
lambda: _post_commit_model_decision(client.events.events, committed_index) is not None,
|
||||
timeout_s=args.timeout_s,
|
||||
label="post-commit model decision or response drain",
|
||||
)
|
||||
post_commit_decision = _post_commit_model_decision(client.events.events, committed_index)
|
||||
except TimeoutError as exc:
|
||||
wait_error = str(exc)
|
||||
await client.acknowledge_playback()
|
||||
close_error: str | None = None
|
||||
try:
|
||||
await client.close_session(timeout_s=args.timeout_s)
|
||||
except TimeoutError as exc:
|
||||
close_error = str(exc)
|
||||
|
||||
audio = client.events.audio_bytes()
|
||||
first_text_at_s = client.events.first_received_at(
|
||||
"response.audio_transcript.delta",
|
||||
"response.output_text.delta",
|
||||
after_s=commit_sent_at_s,
|
||||
)
|
||||
first_audio_at_s = client.events.first_received_at(
|
||||
"response.audio.delta",
|
||||
after_s=commit_sent_at_s,
|
||||
)
|
||||
response_created_at_s = client.events.first_received_at(
|
||||
"response.created",
|
||||
after_s=commit_sent_at_s,
|
||||
)
|
||||
response_done_at_s = client.events.first_received_at(
|
||||
"response.done",
|
||||
after_s=commit_sent_at_s,
|
||||
)
|
||||
audio_duration_s = len(audio) / (client.events.output_sample_rate_hz * 2)
|
||||
response_generation_s = (
|
||||
response_done_at_s - response_created_at_s
|
||||
if response_done_at_s is not None and response_created_at_s is not None
|
||||
else None
|
||||
)
|
||||
transcript_deltas = [
|
||||
str(event.get("delta", ""))
|
||||
for event in client.events.events
|
||||
if event.get("type")
|
||||
in {
|
||||
"response.audio_transcript.delta",
|
||||
"response.output_text.delta",
|
||||
}
|
||||
]
|
||||
response_id = client.events.response_ids[0] if client.events.response_ids else None
|
||||
timing = client.events.timing_summary(
|
||||
after_s=commit_sent_at_s,
|
||||
input_committed_at_s=commit_sent_at_s,
|
||||
response_id=response_id,
|
||||
)
|
||||
errors = client.events.errors()
|
||||
if wait_error:
|
||||
errors.append({"type": "client.timeout", "message": wait_error})
|
||||
if close_error:
|
||||
errors.append({"type": "client.timeout", "message": close_error})
|
||||
model_decision = post_commit_decision or stream_decision
|
||||
(output_dir / "events.jsonl").write_text(
|
||||
"".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(output_dir / "output.pcm").write_bytes(audio)
|
||||
if audio:
|
||||
write_pcm16_wav(
|
||||
output_dir / "output.wav",
|
||||
audio,
|
||||
sample_rate_hz=client.events.output_sample_rate_hz,
|
||||
)
|
||||
|
||||
result = {
|
||||
"ok": (
|
||||
client.events.count("session.created") > 0
|
||||
and client.events.count("session.closed") > 0
|
||||
and not errors
|
||||
and model_decision is not None
|
||||
and (bool(audio) or not args.require_audio)
|
||||
),
|
||||
"model_decision": model_decision,
|
||||
"post_commit": {
|
||||
"input_committed_event_index": committed_index,
|
||||
"decision": post_commit_decision,
|
||||
"decision_required": wait_for_post_commit_decision,
|
||||
"input_had_residual_model_unit": input_has_residual_model_unit,
|
||||
"response_listen_count": _event_count_after(
|
||||
client.events.events,
|
||||
"response.listen",
|
||||
committed_index,
|
||||
),
|
||||
"response_done_count": _event_count_after(
|
||||
client.events.events,
|
||||
"response.done",
|
||||
committed_index,
|
||||
),
|
||||
},
|
||||
"audio_bytes": len(audio),
|
||||
"audio_chunk_count": len(stream_writer.audio_chunk_paths),
|
||||
"audio_chunk_files": [str(path) for path in stream_writer.audio_chunk_paths],
|
||||
"output_sample_rate_hz": client.events.output_sample_rate_hz,
|
||||
"latency": {
|
||||
"ttft_ms": (
|
||||
round((first_text_at_s - commit_sent_at_s) * 1000, 2) if first_text_at_s is not None else None
|
||||
),
|
||||
"ttfp_ms": (
|
||||
round((first_audio_at_s - commit_sent_at_s) * 1000, 2) if first_audio_at_s is not None else None
|
||||
),
|
||||
"rtf": (
|
||||
round(response_generation_s / audio_duration_s, 4)
|
||||
if response_generation_s is not None and audio_duration_s > 0
|
||||
else None
|
||||
),
|
||||
"response_generation_ms": (
|
||||
round(response_generation_s * 1000, 2) if response_generation_s is not None else None
|
||||
),
|
||||
"text_stream_ms": (
|
||||
round((response_done_at_s - first_text_at_s) * 1000, 2)
|
||||
if response_done_at_s is not None and first_text_at_s is not None
|
||||
else None
|
||||
),
|
||||
"transcript_delta_count": len(transcript_deltas),
|
||||
"audio_duration_s": round(audio_duration_s, 3),
|
||||
"measurement_origin": "input_audio_buffer.commit send",
|
||||
},
|
||||
"timing": timing,
|
||||
"response_ids": client.events.response_ids,
|
||||
"transcript": "".join(transcript_deltas),
|
||||
"errors": errors,
|
||||
"output_dir": str(output_dir),
|
||||
}
|
||||
(output_dir / "result.json").write_text(
|
||||
json.dumps(result, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1")
|
||||
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
|
||||
parser.add_argument("--session-id")
|
||||
parser.add_argument("--input-wav", required=True)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
required=True,
|
||||
help=(
|
||||
"Reference WAV for the MiniCPM-o duplex assistant voice. "
|
||||
"This demo matches the official flow by always providing a reference audio clip."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--output-dir", default="/tmp/minicpmo_realtime_duplex_demo")
|
||||
parser.add_argument("--chunk-ms", type=int, default=200)
|
||||
parser.add_argument("--timeout-s", type=float, default=60.0)
|
||||
parser.add_argument("--no-realtime-pacing", action="store_true")
|
||||
parser.add_argument("--require-audio", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
result = asyncio.run(run_demo(parse_args()))
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
if not result["ok"]:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Browser demo for the experimental full-duplex runtime."""
|
||||
|
||||
from .server import build_app
|
||||
|
||||
__all__ = ["build_app"]
|
||||
@@ -0,0 +1,4 @@
|
||||
from .server import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,79 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<meta name="color-scheme" content="light">
|
||||
<title>MiniCPM-o 4.5 Voice</title>
|
||||
<link rel="stylesheet" href="static/styles.css">
|
||||
</head>
|
||||
<body>
|
||||
<header class="app-header">
|
||||
<div>
|
||||
<p class="eyebrow">Experimental full-duplex runtime</p>
|
||||
<h1>MiniCPM-o 4.5 Voice</h1>
|
||||
</div>
|
||||
<span id="connectionState" class="status status-offline">Offline</span>
|
||||
</header>
|
||||
|
||||
<main>
|
||||
<section class="call-bar" aria-label="Session controls">
|
||||
<div class="call-actions">
|
||||
<button id="callButton" class="primary" type="button">Start session</button>
|
||||
<button id="muteButton" class="secondary" type="button" disabled>Mute</button>
|
||||
<button id="cameraButton" class="secondary" type="button" disabled>Camera</button>
|
||||
<video id="cameraPreview" muted playsinline style="display:none;width:96px;height:72px;border-radius:8px;object-fit:cover;background:#000"></video>
|
||||
</div>
|
||||
<div style="display:flex;gap:8px;align-items:flex-start;margin-top:8px;width:100%">
|
||||
<label style="display:inline-flex;flex-direction:column;gap:4px;font-size:13px;opacity:.85;white-space:nowrap">
|
||||
Prompt preset
|
||||
<select id="promptPreset" style="font-size:13px;padding:4px 6px;border-radius:6px">
|
||||
<option value="omni" selected>Omni (official omni preset)</option>
|
||||
<option value="chinese_call">中文通话 (official audio preset)</option>
|
||||
<option value="english_call">English Call (official audio preset)</option>
|
||||
<option value="custom">Custom</option>
|
||||
</select>
|
||||
</label>
|
||||
<textarea id="systemPrompt" rows="2" placeholder="System prompt" style="flex:1;font-size:13px;padding:6px 8px;border-radius:6px;resize:vertical">Streaming Omni Conversation.</textarea>
|
||||
</div>
|
||||
<div class="meter" aria-label="Microphone level">
|
||||
<span>Mic</span>
|
||||
<div class="meter-track"><div id="meterFill" class="meter-fill"></div></div>
|
||||
</div>
|
||||
<dl class="runtime-status">
|
||||
<div><dt>Model</dt><dd id="modelState">Idle</dd></div>
|
||||
<div><dt>Playback</dt><dd id="playbackState">Idle</dd></div>
|
||||
<div><dt>Session</dt><dd id="sessionTimer">00:00</dd></div>
|
||||
</dl>
|
||||
</section>
|
||||
|
||||
<section class="conversation-section" aria-labelledby="conversationTitle">
|
||||
<div class="section-heading">
|
||||
<div>
|
||||
<p class="eyebrow">Live transcript</p>
|
||||
<h2 id="conversationTitle">Conversation</h2>
|
||||
</div>
|
||||
<span class="policy-label">Model-controlled listen / speak</span>
|
||||
</div>
|
||||
<div id="conversation" class="conversation" aria-live="polite">
|
||||
<p id="emptyConversation" class="empty-state">Start a session and speak naturally.</p>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<details class="diagnostics" open>
|
||||
<summary>
|
||||
<span>Event log</span>
|
||||
<span id="eventCount" class="event-count">0 events</span>
|
||||
</summary>
|
||||
<div class="log-toolbar">
|
||||
<span id="runtimeDetail">No active connection</span>
|
||||
<button id="clearLogButton" class="icon-button" type="button" title="Clear event log">Clear</button>
|
||||
</div>
|
||||
<pre id="eventLog" class="event-log" aria-live="polite"></pre>
|
||||
</details>
|
||||
</main>
|
||||
|
||||
<script>window.FULL_DUPLEX_CONFIG = __FULL_DUPLEX_CONFIG__;</script>
|
||||
<script src="static/app.js?v=__FULL_DUPLEX_APP_VERSION__"></script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,671 @@
|
||||
(() => {
|
||||
'use strict';
|
||||
|
||||
const config = window.FULL_DUPLEX_CONFIG || {};
|
||||
const callButton = document.getElementById('callButton');
|
||||
const muteButton = document.getElementById('muteButton');
|
||||
const cameraButton = document.getElementById('cameraButton');
|
||||
const cameraPreview = document.getElementById('cameraPreview');
|
||||
const promptPreset = document.getElementById('promptPreset');
|
||||
const systemPromptInput = document.getElementById('systemPrompt');
|
||||
const connectionState = document.getElementById('connectionState');
|
||||
const modelState = document.getElementById('modelState');
|
||||
const playbackState = document.getElementById('playbackState');
|
||||
const sessionTimer = document.getElementById('sessionTimer');
|
||||
const meterFill = document.getElementById('meterFill');
|
||||
const conversation = document.getElementById('conversation');
|
||||
const emptyConversation = document.getElementById('emptyConversation');
|
||||
const eventLog = document.getElementById('eventLog');
|
||||
const eventCount = document.getElementById('eventCount');
|
||||
const runtimeDetail = document.getElementById('runtimeDetail');
|
||||
const clearLogButton = document.getElementById('clearLogButton');
|
||||
|
||||
const INPUT_RATE = 16000;
|
||||
const OUTPUT_RATE = 24000;
|
||||
const SEND_INTERVAL_MS = 200;
|
||||
const ECHO_GUARD_MS = 300;
|
||||
const INITIAL_PLAYBACK_BUFFER_MS = 400;
|
||||
const SESSION_CLOSE_TIMEOUT_MS = 1000;
|
||||
|
||||
// Default prompts mirroring the official MiniCPM-o-Demo presets
|
||||
// (assets/presets/{omni,audio_duplex}/*.yaml).
|
||||
const PROMPT_PRESETS = {
|
||||
omni: 'Streaming Omni Conversation.',
|
||||
chinese_call: '扮演一个具有以上声音特征的助手。请认真、高质量地回复用户的问题。'
|
||||
+ '请用高自然度的方式和用户聊天。你处于双工模式，可以一边听、一边说。'
|
||||
+ '你是由面壁智能开发的人工智能助手：面壁小钢炮。',
|
||||
english_call: 'Replicate the tone and style from the input audio. Your task is to be '
|
||||
+ 'a helpful assistant using this voice pattern. Please answer the user\'s questions '
|
||||
+ 'seriously and in a high quality. Please chat with the user in a high naturalness '
|
||||
+ 'style. You are in duplex mode, where you can listen and speak at the same time.',
|
||||
};
|
||||
|
||||
let socket = null;
|
||||
let mediaStream = null;
|
||||
let captureContext = null;
|
||||
let captureNode = null;
|
||||
let playbackContext = null;
|
||||
let playbackNode = null;
|
||||
let sendTimer = null;
|
||||
let clockTimer = null;
|
||||
let startedAt = 0;
|
||||
let running = false;
|
||||
let muted = false;
|
||||
let assistantActive = false;
|
||||
let captureRate = INPUT_RATE;
|
||||
let cameraStream = null;
|
||||
let cameraTimer = null;
|
||||
let cameraPendingFrame = null;
|
||||
const cameraCanvas = document.createElement('canvas');
|
||||
let playbackRate = OUTPUT_RATE;
|
||||
let pendingCapture = [];
|
||||
let currentResponseId = null;
|
||||
let responseHasAudio = false;
|
||||
let logCount = 0;
|
||||
let liveUserTurn = null;
|
||||
let liveAssistantTurn = null;
|
||||
let sessionCloseResolver = null;
|
||||
|
||||
function staticAssetUrl(path) {
|
||||
const version = String(config.appVersion || '').trim();
|
||||
return version ? `${path}?v=${encodeURIComponent(version)}` : path;
|
||||
}
|
||||
|
||||
if (promptPreset && systemPromptInput) {
|
||||
promptPreset.addEventListener('change', () => {
|
||||
const preset = PROMPT_PRESETS[promptPreset.value];
|
||||
if (preset !== undefined) systemPromptInput.value = preset;
|
||||
});
|
||||
systemPromptInput.addEventListener('input', () => {
|
||||
promptPreset.value = 'custom';
|
||||
});
|
||||
}
|
||||
|
||||
function realtimeUrl() {
|
||||
const url = new URL(config.realtimePath, window.location.href);
|
||||
url.protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
|
||||
url.searchParams.set('duplex', '1');
|
||||
url.searchParams.set('model', config.model || 'openbmb/MiniCPM-o-4_5');
|
||||
url.searchParams.set('minicpmo45_native_duplex', '1');
|
||||
url.searchParams.set('autostart', '0');
|
||||
return url.toString();
|
||||
}
|
||||
|
||||
function setConnection(label, kind) {
|
||||
connectionState.textContent = label;
|
||||
connectionState.className = `status status-${kind}`;
|
||||
}
|
||||
|
||||
function setModel(label) {
|
||||
modelState.textContent = label;
|
||||
}
|
||||
|
||||
function setPlayback(label) {
|
||||
playbackState.textContent = label;
|
||||
}
|
||||
|
||||
function compactEvent(event) {
|
||||
const fields = [];
|
||||
const payload = event.event || event;
|
||||
const responseId = responseIdOf(event);
|
||||
if (responseId) fields.push(`response=${responseId}`);
|
||||
if (event.item_id) fields.push(`item=${event.item_id}`);
|
||||
if (event.code) fields.push(`code=${event.code}`);
|
||||
if (event.response && event.response.status) fields.push(`status=${event.response.status}`);
|
||||
if (payload.played_ms !== undefined) fields.push(`played=${payload.played_ms}ms`);
|
||||
if (payload.committed_ms !== undefined) fields.push(`committed=${payload.committed_ms}ms`);
|
||||
if (payload.history_committed !== undefined) fields.push(`history=${payload.history_committed}`);
|
||||
return fields.join(' ');
|
||||
}
|
||||
|
||||
function appendLog(message, error = false) {
|
||||
const time = new Date().toLocaleTimeString([], { hour12: false });
|
||||
const line = document.createElement('span');
|
||||
if (error) line.className = 'log-error';
|
||||
line.textContent = `${time} ${message}\n`;
|
||||
eventLog.appendChild(line);
|
||||
eventLog.scrollTop = eventLog.scrollHeight;
|
||||
logCount += 1;
|
||||
eventCount.textContent = `${logCount} ${logCount === 1 ? 'event' : 'events'}`;
|
||||
}
|
||||
|
||||
function appendEventLog(event) {
|
||||
const detail = compactEvent(event);
|
||||
appendLog(`${event.type || 'unknown'}${detail ? ` ${detail}` : ''}`, event.type === 'error');
|
||||
}
|
||||
|
||||
function responseIdOf(event) {
|
||||
return event.response_id || (event.response && event.response.id) || null;
|
||||
}
|
||||
|
||||
function ensureTurn(role) {
|
||||
const existing = role === 'user' ? liveUserTurn : liveAssistantTurn;
|
||||
if (existing) return existing;
|
||||
if (emptyConversation) emptyConversation.remove();
|
||||
const row = document.createElement('div');
|
||||
row.className = `turn turn-${role} turn-live`;
|
||||
const label = document.createElement('div');
|
||||
label.className = 'turn-role';
|
||||
label.textContent = role === 'user' ? 'You' : 'Assistant';
|
||||
const text = document.createElement('div');
|
||||
text.className = 'turn-text';
|
||||
row.append(label, text);
|
||||
conversation.appendChild(row);
|
||||
conversation.scrollTop = conversation.scrollHeight;
|
||||
const turn = { row, text, value: '' };
|
||||
if (role === 'user') liveUserTurn = turn;
|
||||
else liveAssistantTurn = turn;
|
||||
return turn;
|
||||
}
|
||||
|
||||
function addTranscript(role, delta) {
|
||||
if (!delta) return;
|
||||
const turn = ensureTurn(role);
|
||||
turn.value += delta;
|
||||
turn.text.textContent = turn.value;
|
||||
conversation.scrollTop = conversation.scrollHeight;
|
||||
}
|
||||
|
||||
function finishTranscript(role, finalText = '') {
|
||||
const turn = role === 'user' ? liveUserTurn : liveAssistantTurn;
|
||||
if (!turn && !finalText) return;
|
||||
const current = turn || ensureTurn(role);
|
||||
if (finalText) {
|
||||
current.value = finalText;
|
||||
current.text.textContent = finalText;
|
||||
}
|
||||
current.row.classList.remove('turn-live');
|
||||
if (role === 'user') liveUserTurn = null;
|
||||
else liveAssistantTurn = null;
|
||||
}
|
||||
|
||||
function bytesToBase64(bytes) {
|
||||
let binary = '';
|
||||
const chunkSize = 0x8000;
|
||||
for (let offset = 0; offset < bytes.length; offset += chunkSize) {
|
||||
binary += String.fromCharCode(...bytes.subarray(offset, offset + chunkSize));
|
||||
}
|
||||
return btoa(binary);
|
||||
}
|
||||
|
||||
function int16ToBase64(pcm) {
|
||||
return bytesToBase64(new Uint8Array(pcm.buffer, pcm.byteOffset, pcm.byteLength));
|
||||
}
|
||||
|
||||
function base64ToBytes(encoded) {
|
||||
const binary = atob(encoded);
|
||||
const bytes = new Uint8Array(binary.length);
|
||||
for (let index = 0; index < binary.length; index += 1) bytes[index] = binary.charCodeAt(index);
|
||||
return bytes;
|
||||
}
|
||||
|
||||
function resampleInt16(input, sourceRate, targetRate) {
|
||||
if (sourceRate === targetRate) return input;
|
||||
const ratio = sourceRate / targetRate;
|
||||
const output = new Int16Array(Math.floor(input.length / ratio));
|
||||
for (let index = 0; index < output.length; index += 1) {
|
||||
const start = Math.floor(index * ratio);
|
||||
const end = Math.max(start + 1, Math.min(input.length, Math.floor((index + 1) * ratio)));
|
||||
let sum = 0;
|
||||
for (let source = start; source < end; source += 1) sum += input[source];
|
||||
output[index] = sum / (end - start);
|
||||
}
|
||||
return output;
|
||||
}
|
||||
|
||||
async function decodeAudioDelta(event) {
|
||||
const encoded = event.delta || (event.response && event.response.audio);
|
||||
if (!encoded) return null;
|
||||
const bytes = base64ToBytes(encoded);
|
||||
const format = String(event.format || event.audio_format || 'pcm16').toLowerCase();
|
||||
const sourceRate = Number(event.sample_rate_hz || event.sample_rate || OUTPUT_RATE);
|
||||
if (format.includes('f32')) {
|
||||
const floats = new Float32Array(bytes.buffer, bytes.byteOffset, Math.floor(bytes.byteLength / 4));
|
||||
const pcm = new Int16Array(floats.length);
|
||||
for (let index = 0; index < floats.length; index += 1) {
|
||||
const sample = Math.max(-1, Math.min(1, floats[index]));
|
||||
pcm[index] = sample < 0 ? sample * 32768 : sample * 32767;
|
||||
}
|
||||
return { pcm, sourceRate };
|
||||
}
|
||||
if (format.includes('wav')) {
|
||||
const decoded = await playbackContext.decodeAudioData(bytes.buffer.slice(0));
|
||||
const channel = decoded.getChannelData(0);
|
||||
const pcm = new Int16Array(channel.length);
|
||||
for (let index = 0; index < channel.length; index += 1) {
|
||||
const sample = Math.max(-1, Math.min(1, channel[index]));
|
||||
pcm[index] = sample < 0 ? sample * 32768 : sample * 32767;
|
||||
}
|
||||
return { pcm, sourceRate: decoded.sampleRate };
|
||||
}
|
||||
return {
|
||||
pcm: new Int16Array(bytes.buffer, bytes.byteOffset, Math.floor(bytes.byteLength / 2)),
|
||||
sourceRate,
|
||||
};
|
||||
}
|
||||
|
||||
function updateMeter(pcm) {
|
||||
let peak = 0;
|
||||
for (let index = 0; index < pcm.length; index += 8) peak = Math.max(peak, Math.abs(pcm[index]));
|
||||
meterFill.style.width = `${Math.min(100, (peak / 32768) * 150).toFixed(0)}%`;
|
||||
}
|
||||
|
||||
function microphoneUploadEnabled() {
|
||||
return running && !muted;
|
||||
}
|
||||
|
||||
function flushCapture() {
|
||||
if (!socket || socket.readyState !== WebSocket.OPEN || pendingCapture.length === 0) return;
|
||||
if (!microphoneUploadEnabled()) {
|
||||
pendingCapture = [];
|
||||
return;
|
||||
}
|
||||
const length = pendingCapture.reduce((total, chunk) => total + chunk.length, 0);
|
||||
const merged = new Int16Array(length);
|
||||
let offset = 0;
|
||||
for (const chunk of pendingCapture) {
|
||||
merged.set(chunk, offset);
|
||||
offset += chunk.length;
|
||||
}
|
||||
pendingCapture = [];
|
||||
const pcm = resampleInt16(merged, captureRate, INPUT_RATE);
|
||||
const appendEvent = {
|
||||
type: 'input_audio_buffer.append',
|
||||
audio: int16ToBase64(pcm),
|
||||
format: 'pcm16',
|
||||
sample_rate_hz: INPUT_RATE,
|
||||
};
|
||||
// Omni duplex: ~1 fps camera frame rides the audio append (official
|
||||
// MiniCPM-o-Demo contract: one base64 JPEG per ~1 s chunk).
|
||||
if (cameraPendingFrame) {
|
||||
appendEvent.video_frames = [cameraPendingFrame];
|
||||
cameraPendingFrame = null;
|
||||
}
|
||||
socket.send(JSON.stringify(appendEvent));
|
||||
}
|
||||
|
||||
function beginAssistant(responseId) {
|
||||
currentResponseId = responseId || currentResponseId;
|
||||
responseHasAudio = false;
|
||||
assistantActive = true;
|
||||
setModel('Speaking');
|
||||
}
|
||||
|
||||
function feedPlayback(decoded, responseId) {
|
||||
if (!decoded || !decoded.pcm || decoded.pcm.length === 0 || !playbackNode) return;
|
||||
const pcm = resampleInt16(decoded.pcm, decoded.sourceRate, playbackRate);
|
||||
responseHasAudio = true;
|
||||
assistantActive = true;
|
||||
setPlayback('Buffering');
|
||||
playbackNode.port.postMessage({
|
||||
type: 'audio',
|
||||
pcm,
|
||||
responseId: responseId || currentResponseId,
|
||||
initialBufferMs: INITIAL_PLAYBACK_BUFFER_MS,
|
||||
}, [pcm.buffer]);
|
||||
}
|
||||
|
||||
function requestPlaybackDrain(responseId) {
|
||||
if (!playbackNode) return;
|
||||
playbackNode.port.postMessage({ type: 'drain', responseId: responseId || currentResponseId });
|
||||
}
|
||||
|
||||
function sendPlaybackAck(responseId, playedMs) {
|
||||
if (!responseId || !socket || socket.readyState !== WebSocket.OPEN || playedMs <= 0) {
|
||||
if (!responseId && playedMs > 0) appendLog('playback ack skipped: missing response id', true);
|
||||
return;
|
||||
}
|
||||
socket.send(JSON.stringify({
|
||||
type: 'playback.ack',
|
||||
response_id: responseId,
|
||||
item_id: `item_${responseId}`,
|
||||
played_ms: playedMs,
|
||||
committed_ms: playedMs,
|
||||
}));
|
||||
}
|
||||
|
||||
function playbackDrained(message) {
|
||||
const responseId = message.responseId || currentResponseId;
|
||||
sendPlaybackAck(responseId, Number(message.playedMs) || 0);
|
||||
setPlayback('Idle');
|
||||
if (message.underrunMs > 0) {
|
||||
appendLog(`playback underrun ${message.underrunMs} ms`);
|
||||
}
|
||||
window.setTimeout(() => {
|
||||
assistantActive = false;
|
||||
currentResponseId = null;
|
||||
responseHasAudio = false;
|
||||
if (running) setModel('Listening');
|
||||
}, ECHO_GUARD_MS);
|
||||
}
|
||||
|
||||
function handleEvent(event) {
|
||||
appendEventLog(event);
|
||||
const responseId = responseIdOf(event);
|
||||
switch (event.type) {
|
||||
case 'session.created':
|
||||
case 'session.updated':
|
||||
setConnection('Connected', 'online');
|
||||
setModel('Listening');
|
||||
break;
|
||||
case 'response.listen':
|
||||
assistantActive = false;
|
||||
setModel('Listening');
|
||||
break;
|
||||
case 'response.created':
|
||||
case 'response.speak':
|
||||
beginAssistant(responseId);
|
||||
break;
|
||||
case 'response.audio.delta':
|
||||
currentResponseId = responseId || currentResponseId;
|
||||
assistantActive = true;
|
||||
setModel('Speaking');
|
||||
decodeAudioDelta(event)
|
||||
.then((decoded) => feedPlayback(decoded, responseId))
|
||||
.catch((error) => appendLog(`audio decode failed: ${error.message || error}`, true));
|
||||
break;
|
||||
case 'response.audio.done':
|
||||
requestPlaybackDrain(responseId);
|
||||
break;
|
||||
case 'response.audio_transcript.delta':
|
||||
addTranscript('assistant', event.delta || '');
|
||||
break;
|
||||
case 'response.audio_transcript.done':
|
||||
finishTranscript('assistant', event.transcript || '');
|
||||
break;
|
||||
case 'conversation.item.input_audio_transcription.delta':
|
||||
addTranscript('user', event.delta || '');
|
||||
break;
|
||||
case 'conversation.item.input_audio_transcription.completed':
|
||||
finishTranscript('user', event.transcript || '');
|
||||
break;
|
||||
case 'response.done':
|
||||
finishTranscript('assistant');
|
||||
if (!responseHasAudio) requestPlaybackDrain(responseId);
|
||||
break;
|
||||
case 'playback.acknowledged':
|
||||
{
|
||||
const acknowledgement = event.event || event;
|
||||
runtimeDetail.textContent = `Playback committed ${acknowledgement.committed_ms || 0} ms`;
|
||||
}
|
||||
break;
|
||||
case 'session.closed':
|
||||
if (sessionCloseResolver) sessionCloseResolver();
|
||||
sessionCloseResolver = null;
|
||||
break;
|
||||
case 'error':
|
||||
setConnection('Error', 'error');
|
||||
runtimeDetail.textContent = String(event.error || event.code || 'Server error');
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
async function openPlayback() {
|
||||
playbackContext = new (window.AudioContext || window.webkitAudioContext)({ sampleRate: OUTPUT_RATE });
|
||||
playbackRate = playbackContext.sampleRate;
|
||||
await playbackContext.audioWorklet.addModule(staticAssetUrl('static/playback_worklet.js'));
|
||||
playbackNode = new AudioWorkletNode(playbackContext, 'fullduplex-pcm-playback');
|
||||
playbackNode.port.onmessage = (message) => {
|
||||
if (message.data.type === 'playback-started') setPlayback('Playing');
|
||||
else if (message.data.type === 'playback-drained') playbackDrained(message.data);
|
||||
else if (message.data.type === 'playback-underrun') {
|
||||
runtimeDetail.textContent = `Playback underrun ${message.data.underrunMs || 0} ms`;
|
||||
}
|
||||
};
|
||||
playbackNode.connect(playbackContext.destination);
|
||||
await playbackContext.resume();
|
||||
}
|
||||
|
||||
async function openCapture() {
|
||||
mediaStream = await navigator.mediaDevices.getUserMedia({
|
||||
audio: {
|
||||
channelCount: 1,
|
||||
echoCancellation: true,
|
||||
noiseSuppression: true,
|
||||
autoGainControl: true,
|
||||
sampleRate: { ideal: INPUT_RATE },
|
||||
},
|
||||
});
|
||||
try {
|
||||
captureContext = new (window.AudioContext || window.webkitAudioContext)({ sampleRate: INPUT_RATE });
|
||||
} catch (_error) {
|
||||
captureContext = new (window.AudioContext || window.webkitAudioContext)();
|
||||
}
|
||||
captureRate = captureContext.sampleRate;
|
||||
await captureContext.audioWorklet.addModule(staticAssetUrl('static/pcm_worklet.js'));
|
||||
const source = captureContext.createMediaStreamSource(mediaStream);
|
||||
captureNode = new AudioWorkletNode(captureContext, 'fullduplex-pcm-capture');
|
||||
captureNode.port.onmessage = (message) => {
|
||||
const pcm = new Int16Array(message.data);
|
||||
updateMeter(pcm);
|
||||
if (microphoneUploadEnabled()) pendingCapture.push(pcm);
|
||||
};
|
||||
const silentSink = captureContext.createGain();
|
||||
silentSink.gain.value = 0;
|
||||
source.connect(captureNode);
|
||||
captureNode.connect(silentSink).connect(captureContext.destination);
|
||||
await captureContext.resume();
|
||||
}
|
||||
|
||||
function openSocket() {
|
||||
return new Promise((resolve, reject) => {
|
||||
const url = realtimeUrl();
|
||||
socket = new WebSocket(url);
|
||||
let settled = false;
|
||||
socket.onopen = () => {
|
||||
settled = true;
|
||||
const extraBody = {
|
||||
auto_response: true,
|
||||
minicpmo45_native_duplex: true,
|
||||
};
|
||||
const session = {
|
||||
modalities: ['audio', 'text'],
|
||||
voice: 'default',
|
||||
extra_body: extraBody,
|
||||
};
|
||||
// Reference voice for TTS cloning, provided by the server via
|
||||
// --ref-audio (mirrors the official demo's default ref audio).
|
||||
if (config.refAudio) session.ref_audio = config.refAudio;
|
||||
const instructions = systemPromptInput ? systemPromptInput.value.trim() : '';
|
||||
if (instructions) session.instructions = instructions;
|
||||
socket.send(JSON.stringify({ type: 'session.update', session }));
|
||||
runtimeDetail.textContent = `${captureRate} Hz capture / ${playbackRate} Hz playback`;
|
||||
appendLog(`websocket open ${url}`);
|
||||
resolve();
|
||||
};
|
||||
socket.onmessage = (message) => {
|
||||
if (typeof message.data !== 'string') return;
|
||||
try {
|
||||
handleEvent(JSON.parse(message.data));
|
||||
} catch (error) {
|
||||
appendLog(`invalid server event: ${error.message || error}`, true);
|
||||
}
|
||||
};
|
||||
socket.onerror = () => {
|
||||
if (!settled) {
|
||||
settled = true;
|
||||
reject(new Error(`WebSocket connection failed: ${url}`));
|
||||
}
|
||||
};
|
||||
socket.onclose = (event) => {
|
||||
appendLog(`websocket closed code=${event.code}`);
|
||||
if (running) {
|
||||
setConnection('Disconnected', 'error');
|
||||
stopSession();
|
||||
}
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
function formatElapsed(seconds) {
|
||||
const minutes = Math.floor(seconds / 60);
|
||||
return `${String(minutes).padStart(2, '0')}:${String(seconds % 60).padStart(2, '0')}`;
|
||||
}
|
||||
|
||||
function startClock() {
|
||||
startedAt = Date.now();
|
||||
clockTimer = window.setInterval(() => {
|
||||
sessionTimer.textContent = formatElapsed(Math.floor((Date.now() - startedAt) / 1000));
|
||||
}, 1000);
|
||||
}
|
||||
|
||||
async function startSession() {
|
||||
if (running) return;
|
||||
callButton.disabled = true;
|
||||
setConnection('Connecting', 'connecting');
|
||||
runtimeDetail.textContent = 'Requesting microphone access';
|
||||
try {
|
||||
await openPlayback();
|
||||
await openCapture();
|
||||
await openSocket();
|
||||
running = true;
|
||||
muted = false;
|
||||
assistantActive = false;
|
||||
sendTimer = window.setInterval(flushCapture, SEND_INTERVAL_MS);
|
||||
startClock();
|
||||
callButton.textContent = 'End session';
|
||||
callButton.classList.add('is-active');
|
||||
muteButton.disabled = false;
|
||||
cameraButton.disabled = false;
|
||||
setConnection('Connected', 'online');
|
||||
setModel('Listening');
|
||||
appendLog('session started');
|
||||
} catch (error) {
|
||||
appendLog(`start failed: ${error.message || error}`, true);
|
||||
setConnection('Error', 'error');
|
||||
runtimeDetail.textContent = String(error.message || error);
|
||||
await stopSession();
|
||||
} finally {
|
||||
callButton.disabled = false;
|
||||
}
|
||||
}
|
||||
|
||||
async function startCamera() {
|
||||
if (cameraStream) return;
|
||||
cameraStream = await navigator.mediaDevices.getUserMedia({ video: true, audio: false });
|
||||
cameraPreview.srcObject = cameraStream;
|
||||
cameraPreview.style.display = '';
|
||||
await cameraPreview.play().catch(() => {});
|
||||
// Official omni-duplex cadence: one JPEG (quality 0.7) per ~1 s chunk,
|
||||
// no client-side resize (the server normalizes at scale_resolution=448).
|
||||
cameraTimer = window.setInterval(() => {
|
||||
if (!cameraStream || cameraPreview.videoWidth === 0) return;
|
||||
cameraCanvas.width = cameraPreview.videoWidth;
|
||||
cameraCanvas.height = cameraPreview.videoHeight;
|
||||
cameraCanvas.getContext('2d').drawImage(cameraPreview, 0, 0);
|
||||
cameraPendingFrame = cameraCanvas.toDataURL('image/jpeg', 0.7).split(',')[1];
|
||||
}, 1000);
|
||||
cameraButton.textContent = 'Camera off';
|
||||
cameraButton.classList.add('is-active');
|
||||
appendLog('camera on (1 fps omni frames)');
|
||||
}
|
||||
|
||||
function stopCamera() {
|
||||
if (cameraTimer !== null) clearInterval(cameraTimer);
|
||||
cameraTimer = null;
|
||||
if (cameraStream) {
|
||||
for (const track of cameraStream.getTracks()) track.stop();
|
||||
}
|
||||
cameraStream = null;
|
||||
cameraPendingFrame = null;
|
||||
cameraPreview.srcObject = null;
|
||||
cameraPreview.style.display = 'none';
|
||||
cameraButton.textContent = 'Camera';
|
||||
cameraButton.classList.remove('is-active');
|
||||
}
|
||||
|
||||
cameraButton.addEventListener('click', () => {
|
||||
if (cameraStream) {
|
||||
stopCamera();
|
||||
appendLog('camera off');
|
||||
return;
|
||||
}
|
||||
startCamera().catch((error) => appendLog(`camera failed: ${error.message || error}`, true));
|
||||
});
|
||||
|
||||
function waitForSessionClosed(targetSocket, timeoutMs) {
|
||||
if (!targetSocket || targetSocket.readyState !== WebSocket.OPEN) return Promise.resolve();
|
||||
return new Promise((resolve) => {
|
||||
let done = false;
|
||||
const finish = () => {
|
||||
if (done) return;
|
||||
done = true;
|
||||
if (sessionCloseResolver === finish) sessionCloseResolver = null;
|
||||
resolve();
|
||||
};
|
||||
sessionCloseResolver = finish;
|
||||
window.setTimeout(finish, timeoutMs);
|
||||
});
|
||||
}
|
||||
|
||||
async function stopSession({ terminal = true } = {}) {
|
||||
running = false;
|
||||
assistantActive = false;
|
||||
pendingCapture = [];
|
||||
if (sendTimer !== null) clearInterval(sendTimer);
|
||||
if (clockTimer !== null) clearInterval(clockTimer);
|
||||
sendTimer = null;
|
||||
clockTimer = null;
|
||||
if (socket) {
|
||||
const closingSocket = socket;
|
||||
socket = null;
|
||||
closingSocket.onclose = null;
|
||||
if (terminal && closingSocket.readyState === WebSocket.OPEN) {
|
||||
const closed = waitForSessionClosed(closingSocket, SESSION_CLOSE_TIMEOUT_MS);
|
||||
closingSocket.send(JSON.stringify({ type: 'session.close' }));
|
||||
await closed;
|
||||
}
|
||||
closingSocket.close(1000, 'client stop');
|
||||
}
|
||||
if (playbackNode) playbackNode.port.postMessage({ type: 'clear' });
|
||||
if (mediaStream) {
|
||||
for (const track of mediaStream.getTracks()) track.stop();
|
||||
}
|
||||
mediaStream = null;
|
||||
stopCamera();
|
||||
cameraButton.disabled = true;
|
||||
if (captureContext) await captureContext.close().catch(() => {});
|
||||
if (playbackContext) await playbackContext.close().catch(() => {});
|
||||
captureContext = null;
|
||||
captureNode = null;
|
||||
playbackContext = null;
|
||||
playbackNode = null;
|
||||
currentResponseId = null;
|
||||
responseHasAudio = false;
|
||||
meterFill.style.width = '0%';
|
||||
sessionTimer.textContent = '00:00';
|
||||
callButton.textContent = 'Start session';
|
||||
callButton.classList.remove('is-active');
|
||||
muteButton.textContent = 'Mute';
|
||||
muteButton.classList.remove('is-active');
|
||||
muteButton.disabled = true;
|
||||
setConnection('Offline', 'offline');
|
||||
setModel('Idle');
|
||||
setPlayback('Idle');
|
||||
if (runtimeDetail.textContent.startsWith('Playback committed')) return;
|
||||
if (!runtimeDetail.textContent.startsWith('start failed')) runtimeDetail.textContent = 'No active connection';
|
||||
}
|
||||
|
||||
function toggleMute() {
|
||||
if (!running) return;
|
||||
muted = !muted;
|
||||
pendingCapture = [];
|
||||
muteButton.textContent = muted ? 'Unmute' : 'Mute';
|
||||
muteButton.classList.toggle('is-active', muted);
|
||||
appendLog(muted ? 'microphone muted' : 'microphone unmuted');
|
||||
}
|
||||
|
||||
callButton.addEventListener('click', () => {
|
||||
if (running) stopSession();
|
||||
else startSession();
|
||||
});
|
||||
muteButton.addEventListener('click', toggleMute);
|
||||
clearLogButton.addEventListener('click', () => {
|
||||
eventLog.textContent = '';
|
||||
logCount = 0;
|
||||
eventCount.textContent = '0 events';
|
||||
});
|
||||
window.addEventListener('beforeunload', () => { stopSession({ terminal: false }); });
|
||||
})();
|
||||
@@ -0,0 +1,15 @@
|
||||
class FullDuplexPcmCapture extends AudioWorkletProcessor {
|
||||
process(inputs) {
|
||||
const channel = inputs[0] && inputs[0][0];
|
||||
if (!channel || channel.length === 0) return true;
|
||||
const pcm = new Int16Array(channel.length);
|
||||
for (let index = 0; index < channel.length; index += 1) {
|
||||
const sample = Math.max(-1, Math.min(1, channel[index]));
|
||||
pcm[index] = sample < 0 ? sample * 32768 : sample * 32767;
|
||||
}
|
||||
this.port.postMessage(pcm.buffer, [pcm.buffer]);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
registerProcessor('fullduplex-pcm-capture', FullDuplexPcmCapture);
|
||||
@@ -0,0 +1,153 @@
|
||||
class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
constructor() {
|
||||
super();
|
||||
this.queue = [];
|
||||
this.offset = 0;
|
||||
this.playedFrames = 0;
|
||||
this.underrunFrames = 0;
|
||||
this.drain = null;
|
||||
this.started = false;
|
||||
this.activeResponseId = null;
|
||||
this.initialBufferFrames = Math.round(sampleRate * 0.2);
|
||||
this.bufferWaitFrames = this.initialBufferFrames;
|
||||
this.rebuffering = false;
|
||||
this.fadeFrames = Math.max(1, Math.round(sampleRate * 0.005));
|
||||
this.fadeInFrames = 0;
|
||||
this.port.onmessage = (event) => this.handleMessage(event.data || {});
|
||||
}
|
||||
|
||||
handleMessage(message) {
|
||||
if (message.type === 'audio' && message.pcm) {
|
||||
const wasEmpty = this.queue.length === 0;
|
||||
if (!this.started && !this.activeResponseId) {
|
||||
this.activeResponseId = message.responseId || null;
|
||||
}
|
||||
if (!this.started && Number.isFinite(message.initialBufferMs)) {
|
||||
this.initialBufferFrames = Math.max(0, Math.round((sampleRate * message.initialBufferMs) / 1000));
|
||||
}
|
||||
this.queue.push(message.pcm);
|
||||
if (!this.started && wasEmpty && !this.rebuffering) {
|
||||
this.bufferWaitFrames = this.initialBufferFrames;
|
||||
}
|
||||
} else if (message.type === 'drain') {
|
||||
this.drain = { responseId: message.responseId || null };
|
||||
if (!this.started && this.bufferedFrames() > 0) {
|
||||
this.bufferWaitFrames = 0;
|
||||
this.startPlayback();
|
||||
}
|
||||
this.notifyIfDrained();
|
||||
} else if (message.type === 'clear') {
|
||||
this.queue = [];
|
||||
this.offset = 0;
|
||||
this.playedFrames = 0;
|
||||
this.underrunFrames = 0;
|
||||
this.drain = null;
|
||||
this.started = false;
|
||||
this.activeResponseId = null;
|
||||
this.bufferWaitFrames = this.initialBufferFrames;
|
||||
this.rebuffering = false;
|
||||
this.fadeInFrames = 0;
|
||||
}
|
||||
}
|
||||
|
||||
bufferedFrames() {
|
||||
return this.queue.reduce((total, pcm, index) => (
|
||||
total + pcm.length - (index === 0 ? this.offset : 0)
|
||||
), 0);
|
||||
}
|
||||
|
||||
startPlayback() {
|
||||
if (this.started) return;
|
||||
this.started = true;
|
||||
this.fadeInFrames = this.fadeFrames;
|
||||
if (this.playedFrames === 0) {
|
||||
this.port.postMessage({ type: 'playback-started', responseId: this.activeResponseId });
|
||||
}
|
||||
}
|
||||
|
||||
notifyIfDrained() {
|
||||
if (!this.drain || this.queue.length > 0) return;
|
||||
this.port.postMessage({
|
||||
type: 'playback-drained',
|
||||
responseId: this.drain.responseId,
|
||||
playedMs: Math.round((this.playedFrames * 1000) / sampleRate),
|
||||
underrunMs: Math.round((this.underrunFrames * 1000) / sampleRate),
|
||||
});
|
||||
this.playedFrames = 0;
|
||||
this.underrunFrames = 0;
|
||||
this.drain = null;
|
||||
this.started = false;
|
||||
this.activeResponseId = null;
|
||||
this.bufferWaitFrames = this.initialBufferFrames;
|
||||
this.rebuffering = false;
|
||||
this.fadeInFrames = 0;
|
||||
}
|
||||
|
||||
reportUnderrun() {
|
||||
this.port.postMessage({
|
||||
type: 'playback-underrun',
|
||||
responseId: this.activeResponseId,
|
||||
underrunMs: Math.round((this.underrunFrames * 1000) / sampleRate),
|
||||
});
|
||||
}
|
||||
|
||||
process(_inputs, outputs) {
|
||||
const output = outputs[0][0];
|
||||
output.fill(0);
|
||||
if (!this.started) {
|
||||
if (this.rebuffering && !this.drain) {
|
||||
this.underrunFrames += output.length;
|
||||
this.reportUnderrun();
|
||||
}
|
||||
if (this.queue.length > 0 && this.bufferWaitFrames > 0) {
|
||||
this.bufferWaitFrames = Math.max(0, this.bufferWaitFrames - output.length);
|
||||
return true;
|
||||
}
|
||||
if (this.queue.length > 0) {
|
||||
this.startPlayback();
|
||||
this.rebuffering = false;
|
||||
}
|
||||
}
|
||||
if (!this.started) {
|
||||
this.notifyIfDrained();
|
||||
return true;
|
||||
}
|
||||
let target = 0;
|
||||
while (target < output.length && this.queue.length > 0) {
|
||||
const pcm = this.queue[0];
|
||||
const count = Math.min(output.length - target, pcm.length - this.offset);
|
||||
for (let index = 0; index < count; index += 1) {
|
||||
let sample = pcm[this.offset + index] / 32768;
|
||||
if (this.fadeInFrames > 0) {
|
||||
const elapsed = this.fadeFrames - this.fadeInFrames;
|
||||
sample *= elapsed / this.fadeFrames;
|
||||
this.fadeInFrames -= 1;
|
||||
}
|
||||
output[target + index] = sample;
|
||||
}
|
||||
target += count;
|
||||
this.offset += count;
|
||||
this.playedFrames += count;
|
||||
if (this.offset >= pcm.length) {
|
||||
this.queue.shift();
|
||||
this.offset = 0;
|
||||
}
|
||||
}
|
||||
if (target < output.length && !this.drain) {
|
||||
const fadeCount = Math.min(target, this.fadeFrames);
|
||||
for (let index = 0; index < fadeCount; index += 1) {
|
||||
output[target - fadeCount + index] *= (fadeCount - index - 1) / fadeCount;
|
||||
}
|
||||
this.underrunFrames += output.length - target;
|
||||
this.reportUnderrun();
|
||||
this.started = false;
|
||||
this.rebuffering = true;
|
||||
this.bufferWaitFrames = this.initialBufferFrames;
|
||||
this.fadeInFrames = 0;
|
||||
}
|
||||
this.notifyIfDrained();
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
registerProcessor('fullduplex-pcm-playback', FullDuplexPcmPlayback);
|
||||
@@ -0,0 +1,127 @@
|
||||
:root {
|
||||
color-scheme: light;
|
||||
font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
|
||||
color: #20252b;
|
||||
background: #f4f6f7;
|
||||
font-synthesis: none;
|
||||
}
|
||||
|
||||
* { box-sizing: border-box; }
|
||||
|
||||
body { margin: 0; min-width: 320px; }
|
||||
|
||||
button, summary { font: inherit; }
|
||||
|
||||
button:focus-visible, summary:focus-visible {
|
||||
outline: 3px solid rgba(21, 101, 192, .25);
|
||||
outline-offset: 2px;
|
||||
}
|
||||
|
||||
.app-header {
|
||||
min-height: 104px;
|
||||
padding: 24px clamp(20px, 5vw, 64px);
|
||||
background: #ffffff;
|
||||
border-bottom: 1px solid #d9dee3;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
gap: 20px;
|
||||
}
|
||||
|
||||
h1, h2, p { margin: 0; }
|
||||
h1 { font-size: 24px; line-height: 1.2; letter-spacing: 0; }
|
||||
h2 { font-size: 18px; line-height: 1.3; letter-spacing: 0; }
|
||||
|
||||
.eyebrow {
|
||||
margin-bottom: 5px;
|
||||
color: #68717b;
|
||||
font-size: 11px;
|
||||
font-weight: 700;
|
||||
text-transform: uppercase;
|
||||
letter-spacing: .08em;
|
||||
}
|
||||
|
||||
main { width: min(1040px, calc(100% - 40px)); margin: 28px auto 48px; }
|
||||
|
||||
.status {
|
||||
min-width: 88px;
|
||||
padding: 7px 10px;
|
||||
border: 1px solid;
|
||||
border-radius: 6px;
|
||||
text-align: center;
|
||||
font-size: 13px;
|
||||
font-weight: 700;
|
||||
}
|
||||
|
||||
.status-offline { color: #616971; background: #f0f2f3; border-color: #cbd1d6; }
|
||||
.status-connecting { color: #8a4b00; background: #fff5df; border-color: #e6b76a; }
|
||||
.status-online { color: #176a3a; background: #e9f7ef; border-color: #7dc49b; }
|
||||
.status-error { color: #a12222; background: #fff0f0; border-color: #df8f8f; }
|
||||
|
||||
.call-bar {
|
||||
min-height: 86px;
|
||||
padding: 16px 18px;
|
||||
background: #252b31;
|
||||
color: #ffffff;
|
||||
border-radius: 6px;
|
||||
display: grid;
|
||||
grid-template-columns: auto minmax(150px, 1fr) auto;
|
||||
align-items: center;
|
||||
gap: 24px;
|
||||
}
|
||||
|
||||
.call-actions { display: flex; gap: 8px; }
|
||||
button { border-radius: 5px; border: 1px solid transparent; cursor: pointer; }
|
||||
button:disabled { opacity: .46; cursor: default; }
|
||||
.primary { padding: 10px 15px; color: #ffffff; background: #16834b; font-weight: 700; }
|
||||
.primary.is-active { background: #b43636; }
|
||||
.secondary { padding: 10px 13px; color: #ffffff; background: #39414a; border-color: #59636d; }
|
||||
.secondary.is-active { color: #2b2525; background: #ffc9c9; border-color: #e58f8f; }
|
||||
|
||||
.meter { min-width: 0; display: flex; align-items: center; gap: 10px; color: #c8cfd6; font-size: 12px; }
|
||||
.meter-track { height: 8px; flex: 1; overflow: hidden; background: #444d56; border-radius: 4px; }
|
||||
.meter-fill { width: 0; height: 100%; background: #54bd7d; transition: width 70ms linear; }
|
||||
|
||||
.runtime-status { margin: 0; display: flex; gap: 24px; }
|
||||
.runtime-status div { min-width: 68px; }
|
||||
.runtime-status dt { color: #9fa8b1; font-size: 10px; text-transform: uppercase; }
|
||||
.runtime-status dd { margin: 3px 0 0; font-size: 13px; font-weight: 700; font-variant-numeric: tabular-nums; }
|
||||
|
||||
.conversation-section { margin-top: 22px; padding: 22px 0 10px; }
|
||||
.section-heading { display: flex; justify-content: space-between; align-items: end; gap: 16px; padding: 0 2px 14px; border-bottom: 1px solid #cfd5da; }
|
||||
.policy-label { color: #5e6871; font-size: 12px; }
|
||||
.conversation { min-height: 260px; max-height: 440px; overflow-y: auto; padding: 18px 2px; }
|
||||
.empty-state { color: #7a838b; font-size: 14px; }
|
||||
|
||||
.turn { display: grid; grid-template-columns: 92px minmax(0, 1fr); gap: 12px; padding: 12px 0; border-bottom: 1px solid #e0e4e7; }
|
||||
.turn-role { padding-top: 2px; color: #68717b; font-size: 11px; font-weight: 800; text-transform: uppercase; }
|
||||
.turn-user .turn-role { color: #1565c0; }
|
||||
.turn-assistant .turn-role { color: #a45117; }
|
||||
.turn-text { min-width: 0; white-space: pre-wrap; overflow-wrap: anywhere; line-height: 1.55; }
|
||||
.turn-live .turn-text::after { content: ""; display: inline-block; width: 6px; height: 14px; margin-left: 4px; background: #9098a0; vertical-align: -2px; animation: blink 1s step-end infinite; }
|
||||
@keyframes blink { 50% { opacity: 0; } }
|
||||
|
||||
.diagnostics { margin-top: 18px; border-top: 1px solid #bfc6cc; border-bottom: 1px solid #bfc6cc; background: #ffffff; }
|
||||
.diagnostics summary { padding: 13px 14px; display: flex; justify-content: space-between; cursor: pointer; font-weight: 700; list-style: none; }
|
||||
.diagnostics summary::-webkit-details-marker { display: none; }
|
||||
.event-count { color: #6f7880; font-size: 12px; font-weight: 500; }
|
||||
.log-toolbar { min-height: 40px; padding: 7px 12px; border-top: 1px solid #e0e4e7; display: flex; align-items: center; justify-content: space-between; gap: 12px; color: #69727a; font-size: 12px; }
|
||||
.icon-button { padding: 5px 9px; color: #394149; background: #f4f6f7; border-color: #cdd3d8; }
|
||||
.event-log { height: 200px; margin: 0; padding: 12px 14px; overflow: auto; background: #171b1f; color: #d9e0e6; font: 11px/1.55 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; white-space: pre-wrap; }
|
||||
.log-error { color: #ff9f9f; }
|
||||
|
||||
@media (max-width: 760px) {
|
||||
.app-header { min-height: 88px; padding: 18px 20px; }
|
||||
h1 { font-size: 20px; }
|
||||
main { width: min(100% - 24px, 620px); margin-top: 16px; }
|
||||
.call-bar { grid-template-columns: 1fr; gap: 15px; }
|
||||
.call-actions button { flex: 1; }
|
||||
.runtime-status { justify-content: space-between; gap: 10px; }
|
||||
.section-heading { align-items: start; flex-direction: column; gap: 7px; }
|
||||
.conversation { min-height: 230px; }
|
||||
.turn { grid-template-columns: 72px minmax(0, 1fr); }
|
||||
}
|
||||
|
||||
@media (prefers-reduced-motion: reduce) {
|
||||
*, *::before, *::after { animation-duration: .01ms !important; animation-iteration-count: 1 !important; }
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Static host and same-origin Realtime proxy for the full-duplex demo."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import uvicorn
|
||||
import websockets
|
||||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
||||
from fastapi.responses import HTMLResponse, Response
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
APP_DIR = Path(__file__).parent / "app"
|
||||
STATIC_DIR = APP_DIR / "static"
|
||||
|
||||
|
||||
def _join_ws_url(base: str, path: str, query: str) -> str:
|
||||
return base.rstrip("/") + path + (("?" + query) if query else "")
|
||||
|
||||
|
||||
async def _pump_client_to_backend(client: WebSocket, backend) -> None:
|
||||
try:
|
||||
while True:
|
||||
message = await client.receive()
|
||||
if message["type"] == "websocket.disconnect":
|
||||
await backend.close()
|
||||
return
|
||||
if message.get("text") is not None:
|
||||
await backend.send(message["text"])
|
||||
elif message.get("bytes") is not None:
|
||||
await backend.send(message["bytes"])
|
||||
except (WebSocketDisconnect, websockets.ConnectionClosed):
|
||||
with contextlib.suppress(Exception):
|
||||
await backend.close()
|
||||
|
||||
|
||||
async def _pump_backend_to_client(client: WebSocket, backend) -> None:
|
||||
try:
|
||||
async for message in backend:
|
||||
if isinstance(message, bytes):
|
||||
await client.send_bytes(message)
|
||||
else:
|
||||
await client.send_text(message)
|
||||
except (WebSocketDisconnect, websockets.ConnectionClosed):
|
||||
return
|
||||
except RuntimeError as exc:
|
||||
if "websocket.send" in str(exc) and "websocket.close" in str(exc):
|
||||
return
|
||||
raise
|
||||
|
||||
|
||||
def _expected_proxy_close(exc: BaseException) -> bool:
|
||||
if isinstance(exc, (WebSocketDisconnect, websockets.ConnectionClosed, asyncio.CancelledError)):
|
||||
return True
|
||||
return isinstance(exc, RuntimeError) and "websocket.send" in str(exc) and "websocket.close" in str(exc)
|
||||
|
||||
|
||||
def build_app(
|
||||
*,
|
||||
ws_backend: str = "ws://127.0.0.1:8099",
|
||||
model: str = "openbmb/MiniCPM-o-4_5",
|
||||
public_realtime_url: str | None = None,
|
||||
ref_audio: str | None = None,
|
||||
) -> FastAPI:
|
||||
app = FastAPI(title="Experimental Full-Duplex Web Demo")
|
||||
index_path = APP_DIR / "index.html"
|
||||
app_version_hash = hashlib.sha256()
|
||||
for asset_path in (
|
||||
STATIC_DIR / "app.js",
|
||||
STATIC_DIR / "pcm_worklet.js",
|
||||
STATIC_DIR / "playback_worklet.js",
|
||||
):
|
||||
app_version_hash.update(asset_path.read_bytes())
|
||||
app_version = app_version_hash.hexdigest()[:12]
|
||||
|
||||
ref_audio_uri: str | None = None
|
||||
if ref_audio:
|
||||
ref_path = Path(ref_audio)
|
||||
if not ref_path.is_file():
|
||||
raise SystemExit(f"--ref-audio not found: {ref_audio}")
|
||||
encoded = base64.b64encode(ref_path.read_bytes()).decode("ascii")
|
||||
ref_audio_uri = f"data:audio/wav;base64,{encoded}"
|
||||
|
||||
@app.get("/", response_class=HTMLResponse)
|
||||
def index() -> HTMLResponse:
|
||||
config = json.dumps(
|
||||
{
|
||||
"model": model,
|
||||
"realtimePath": public_realtime_url or "v1/realtime",
|
||||
"refAudio": ref_audio_uri,
|
||||
"appVersion": app_version,
|
||||
},
|
||||
ensure_ascii=True,
|
||||
)
|
||||
html = (
|
||||
index_path.read_text(encoding="utf-8")
|
||||
.replace("__FULL_DUPLEX_CONFIG__", config)
|
||||
.replace("__FULL_DUPLEX_APP_VERSION__", app_version)
|
||||
)
|
||||
return HTMLResponse(html, headers={"Cache-Control": "no-store"})
|
||||
|
||||
@app.get("/healthz")
|
||||
def healthz() -> Response:
|
||||
return Response(content="ok", media_type="text/plain")
|
||||
|
||||
@app.websocket("/v1/realtime")
|
||||
async def realtime_proxy(websocket: WebSocket) -> None:
|
||||
await websocket.accept()
|
||||
query = urlencode(websocket.query_params.multi_items())
|
||||
backend_url = _join_ws_url(ws_backend, "/v1/realtime", query)
|
||||
logger.info("Proxying Realtime WebSocket to %s", backend_url)
|
||||
try:
|
||||
async with websockets.connect(
|
||||
backend_url,
|
||||
max_size=64 * 1024 * 1024,
|
||||
) as backend:
|
||||
tasks = {
|
||||
asyncio.create_task(_pump_client_to_backend(websocket, backend)),
|
||||
asyncio.create_task(_pump_backend_to_client(websocket, backend)),
|
||||
}
|
||||
done, pending = await asyncio.wait(
|
||||
tasks,
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
for result in await asyncio.gather(*done, return_exceptions=True):
|
||||
if isinstance(result, BaseException) and not _expected_proxy_close(result):
|
||||
raise result
|
||||
except (WebSocketDisconnect, websockets.ConnectionClosed):
|
||||
return
|
||||
except Exception:
|
||||
logger.exception("Realtime WebSocket proxy failed")
|
||||
with contextlib.suppress(Exception):
|
||||
await websocket.close(code=1011)
|
||||
|
||||
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
|
||||
return app
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=7862)
|
||||
parser.add_argument("--ws-backend", default="ws://127.0.0.1:8099")
|
||||
parser.add_argument(
|
||||
"--public-realtime-url",
|
||||
help="Browser-visible ws:// or wss:// Realtime URL; defaults to the same-origin proxy.",
|
||||
)
|
||||
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
required=True,
|
||||
help=(
|
||||
"Reference voice wav for TTS voice cloning, e.g. the "
|
||||
"official MiniCPM-o-Demo assets/ref_audio/ref_minicpm_signature.wav"
|
||||
),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
uvicorn.run(
|
||||
build_app(
|
||||
ws_backend=args.ws_backend,
|
||||
model=args.model,
|
||||
public_realtime_url=args.public_realtime_url,
|
||||
ref_audio=args.ref_audio,
|
||||
),
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
log_level="info",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,26 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_importing_scheduler_output_does_not_eagerly_import_scheduler_implementations() -> None:
|
||||
package_init = Path(__file__).parents[3] / "vllm_omni/core/sched/__init__.py"
|
||||
module = ast.parse(package_init.read_text())
|
||||
eager_modules = {
|
||||
node.module
|
||||
for node in ast.walk(module)
|
||||
if isinstance(node, ast.ImportFrom)
|
||||
and node.level == 1
|
||||
and node.module in {"omni_ar_scheduler", "omni_generation_scheduler"}
|
||||
}
|
||||
|
||||
assert eager_modules == set()
|
||||
assert any(isinstance(node, ast.FunctionDef) and node.name == "__getattr__" for node in module.body)
|
||||
@@ -677,6 +677,28 @@ def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter
|
||||
assert "req-zombie" not in adapter.finished_requests
|
||||
|
||||
|
||||
def test_finish_requests_releases_active_stream_slot(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=1)
|
||||
aborted = _req("req-aborted", RequestStatus.RUNNING)
|
||||
waiting = _req("req-waiting", RequestStatus.WAITING)
|
||||
waiting_queue = DummyWaitingQueue([waiting])
|
||||
running_queue = []
|
||||
adapter._active_streams[aborted.request_id] = aborted
|
||||
adapter._held_non_active.append(aborted)
|
||||
|
||||
adapter.finish_requests(
|
||||
[aborted.request_id],
|
||||
RequestStatus.FINISHED_ABORTED,
|
||||
{aborted.request_id: aborted},
|
||||
)
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
|
||||
assert aborted.request_id not in adapter._active_streams
|
||||
assert [request.request_id for request in adapter._held_non_active] == []
|
||||
assert list(adapter._active_streams) == [waiting.request_id]
|
||||
assert waiting.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
|
||||
|
||||
def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1)
|
||||
zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,388 @@
|
||||
"""Run two MiniCPM-o Realtime duplex demo processes concurrently.
|
||||
|
||||
This E2E driver covers the same process boundary users exercise manually:
|
||||
two independent ``realtime_duplex_demo.py`` invocations, two distinct input
|
||||
WAVs, and two distinct output directories.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import uuid
|
||||
import wave
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
DEMO_PATH = REPO_ROOT / "examples/online_serving/minicpmo/realtime_duplex_demo.py"
|
||||
AUDIO_DELTA_EVENTS = {"response.audio.delta", "response.output_audio.delta"}
|
||||
|
||||
|
||||
def _canonical_path(path: str) -> Path:
|
||||
return Path(path).expanduser().resolve(strict=False)
|
||||
|
||||
|
||||
def _sha256_file(path: Path) -> str:
|
||||
hasher = hashlib.sha256()
|
||||
with path.open("rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
||||
hasher.update(chunk)
|
||||
return hasher.hexdigest()
|
||||
|
||||
|
||||
def _validate_pair_args(args: argparse.Namespace) -> dict[str, object]:
|
||||
input_a = _canonical_path(args.input_wav_a)
|
||||
input_b = _canonical_path(args.input_wav_b)
|
||||
output_a = _canonical_path(args.output_dir_a)
|
||||
output_b = _canonical_path(args.output_dir_b)
|
||||
if output_a == output_b:
|
||||
raise ValueError("output directories must be different")
|
||||
if not input_a.is_file():
|
||||
raise ValueError(f"input WAV does not exist: {input_a}")
|
||||
if not input_b.is_file():
|
||||
raise ValueError(f"input WAV does not exist: {input_b}")
|
||||
sha_a = _sha256_file(input_a)
|
||||
sha_b = _sha256_file(input_b)
|
||||
if sha_a == sha_b:
|
||||
raise ValueError("input WAV files must have different content")
|
||||
return {
|
||||
"input_wav_a": str(input_a),
|
||||
"input_wav_b": str(input_b),
|
||||
"input_sha256_a": sha_a,
|
||||
"input_sha256_b": sha_b,
|
||||
"output_dir_a": str(output_a),
|
||||
"output_dir_b": str(output_b),
|
||||
}
|
||||
|
||||
|
||||
def _event_response_id(event: dict[str, object]) -> str | None:
|
||||
response_id = event.get("response_id")
|
||||
if isinstance(response_id, str):
|
||||
return response_id
|
||||
response = event.get("response")
|
||||
if isinstance(response, dict) and isinstance(response.get("id"), str):
|
||||
return response["id"]
|
||||
return None
|
||||
|
||||
|
||||
def _read_json(path: Path) -> dict[str, object]:
|
||||
if not path.is_file():
|
||||
return {}
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
|
||||
def _read_events(path: Path) -> list[dict[str, object]]:
|
||||
if not path.is_file():
|
||||
return []
|
||||
events: list[dict[str, object]] = []
|
||||
for line in path.read_text(encoding="utf-8").splitlines():
|
||||
if not line.strip():
|
||||
continue
|
||||
event = json.loads(line)
|
||||
if isinstance(event, dict):
|
||||
events.append(event)
|
||||
return events
|
||||
|
||||
|
||||
def _response_ids_from_artifacts(result: dict[str, object], events: list[dict[str, object]]) -> list[str]:
|
||||
response_ids = result.get("response_ids")
|
||||
if isinstance(response_ids, list):
|
||||
ids = [item for item in response_ids if isinstance(item, str)]
|
||||
if ids:
|
||||
return ids
|
||||
seen: list[str] = []
|
||||
for event in events:
|
||||
response_id = _event_response_id(event)
|
||||
if response_id is not None and response_id not in seen:
|
||||
seen.append(response_id)
|
||||
return seen
|
||||
|
||||
|
||||
def _audio_delta_duration_ms(event: dict[str, object]) -> float | None:
|
||||
encoded = event.get("delta") or event.get("audio")
|
||||
if not isinstance(encoded, str) or not encoded:
|
||||
return None
|
||||
try:
|
||||
raw = base64.b64decode(encoded)
|
||||
except ValueError:
|
||||
return None
|
||||
fmt = str(event.get("format") or event.get("audio_format") or "pcm16").lower()
|
||||
sample_rate = event.get("sample_rate_hz") or event.get("sample_rate") or 24_000
|
||||
if not isinstance(sample_rate, int | float) or sample_rate <= 0:
|
||||
return None
|
||||
if "wav" in fmt:
|
||||
try:
|
||||
with wave.open(io.BytesIO(raw), "rb") as wav:
|
||||
return round((wav.getnframes() * 1000.0) / wav.getframerate(), 3)
|
||||
except (EOFError, wave.Error):
|
||||
return None
|
||||
bytes_per_sample = 4 if "f32" in fmt else 2
|
||||
return round((len(raw) * 1000.0) / (bytes_per_sample * float(sample_rate)), 3)
|
||||
|
||||
|
||||
def _audio_delta_received_at_s(event: dict[str, object]) -> float | None:
|
||||
raw = event.get("_client_received_at_s")
|
||||
if isinstance(raw, int | float):
|
||||
return float(raw)
|
||||
return None
|
||||
|
||||
|
||||
def _session_audio_contract_ok(
|
||||
*,
|
||||
audio_delta_count: int,
|
||||
delta_response_ids: list[str],
|
||||
response_ids: list[str],
|
||||
first_audio_delta_before_done: bool,
|
||||
) -> bool:
|
||||
if audio_delta_count <= 0 or not response_ids:
|
||||
return False
|
||||
if len(delta_response_ids) != audio_delta_count:
|
||||
return False
|
||||
delta_id_set = set(delta_response_ids)
|
||||
if len(delta_id_set) != 1:
|
||||
return False
|
||||
return first_audio_delta_before_done and next(iter(delta_id_set)) in set(response_ids)
|
||||
|
||||
|
||||
def _summarize_session(
|
||||
*,
|
||||
label: str,
|
||||
input_wav: str,
|
||||
output_dir: str,
|
||||
returncode: int,
|
||||
stdout: str,
|
||||
stderr: str,
|
||||
) -> dict[str, object]:
|
||||
output_path = Path(output_dir)
|
||||
result = _read_json(output_path / "result.json")
|
||||
events = _read_events(output_path / "events.jsonl")
|
||||
audio_delta_indices = [index for index, event in enumerate(events) if event.get("type") in AUDIO_DELTA_EVENTS]
|
||||
done_indices = [index for index, event in enumerate(events) if event.get("type") == "response.done"]
|
||||
response_ids = _response_ids_from_artifacts(result, events)
|
||||
delta_response_ids = [
|
||||
response_id
|
||||
for response_id in (_event_response_id(events[index]) for index in audio_delta_indices)
|
||||
if response_id is not None
|
||||
]
|
||||
first_audio_delta_before_done = bool(audio_delta_indices and done_indices) and min(audio_delta_indices) < min(
|
||||
done_indices
|
||||
)
|
||||
audio_delta_events = [events[index] for index in audio_delta_indices]
|
||||
first_received_at_s = next(
|
||||
(
|
||||
received_at
|
||||
for received_at in (_audio_delta_received_at_s(event) for event in audio_delta_events)
|
||||
if received_at is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
audio_delta_timings = []
|
||||
for event in audio_delta_events:
|
||||
received_at_s = _audio_delta_received_at_s(event)
|
||||
audio_delta_timings.append(
|
||||
{
|
||||
"response_id": _event_response_id(event),
|
||||
"received_at_s": received_at_s,
|
||||
"arrival_offset_ms": (
|
||||
round((received_at_s - first_received_at_s) * 1000.0, 3)
|
||||
if received_at_s is not None and first_received_at_s is not None
|
||||
else None
|
||||
),
|
||||
"audio_duration_ms": _audio_delta_duration_ms(event),
|
||||
}
|
||||
)
|
||||
all_audio_deltas_same_response = (
|
||||
len(delta_response_ids) == len(audio_delta_indices) and len(set(delta_response_ids)) == 1
|
||||
)
|
||||
response_id_nonempty = bool(response_ids)
|
||||
response_audio_contract_ok = _session_audio_contract_ok(
|
||||
audio_delta_count=len(audio_delta_indices),
|
||||
delta_response_ids=delta_response_ids,
|
||||
response_ids=response_ids,
|
||||
first_audio_delta_before_done=first_audio_delta_before_done,
|
||||
)
|
||||
return {
|
||||
"label": label,
|
||||
"input_wav": input_wav,
|
||||
"output_dir": output_dir,
|
||||
"returncode": returncode,
|
||||
"ok": returncode == 0 and result.get("ok") is True and response_audio_contract_ok,
|
||||
"result_ok": result.get("ok"),
|
||||
"response_ids": response_ids,
|
||||
"response_id_nonempty": response_id_nonempty,
|
||||
"audio_bytes": result.get("audio_bytes", 0),
|
||||
"audio_delta_count": len(audio_delta_indices),
|
||||
"audio_delta_response_ids": delta_response_ids,
|
||||
"audio_delta_timings": audio_delta_timings,
|
||||
"all_audio_deltas_same_response": all_audio_deltas_same_response,
|
||||
"first_audio_delta_before_done": first_audio_delta_before_done,
|
||||
"response_audio_contract_ok": response_audio_contract_ok,
|
||||
"event_count": len(events),
|
||||
"stdout_tail": stdout[-4000:],
|
||||
"stderr_tail": stderr[-4000:],
|
||||
}
|
||||
|
||||
|
||||
def _identity_isolation_ok(sessions: list[dict[str, object]]) -> bool:
|
||||
seen: set[str] = set()
|
||||
for session in sessions:
|
||||
ids = {response_id for response_id in session.get("response_ids", []) if isinstance(response_id, str)}
|
||||
if not ids:
|
||||
return False
|
||||
if ids & seen:
|
||||
return False
|
||||
seen.update(ids)
|
||||
return True
|
||||
|
||||
|
||||
async def _run_demo_process(
|
||||
*,
|
||||
label: str,
|
||||
args: argparse.Namespace,
|
||||
input_wav: str,
|
||||
output_dir: str,
|
||||
) -> dict[str, object]:
|
||||
command = [
|
||||
sys.executable,
|
||||
str(DEMO_PATH),
|
||||
"--url",
|
||||
args.url,
|
||||
"--model",
|
||||
args.model,
|
||||
"--input-wav",
|
||||
input_wav,
|
||||
"--output-dir",
|
||||
output_dir,
|
||||
"--chunk-ms",
|
||||
str(args.chunk_ms),
|
||||
"--timeout-s",
|
||||
str(args.timeout_s),
|
||||
"--session-id",
|
||||
f"duplex-pair-{label}-{uuid.uuid4().hex}",
|
||||
]
|
||||
if args.no_realtime_pacing:
|
||||
command.append("--no-realtime-pacing")
|
||||
if args.require_audio:
|
||||
command.append("--require-audio")
|
||||
command.extend(["--ref-audio", str(args.ref_audio)])
|
||||
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*command,
|
||||
cwd=str(REPO_ROOT),
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
try:
|
||||
stdout_bytes, stderr_bytes = await asyncio.wait_for(
|
||||
process.communicate(),
|
||||
timeout=args.timeout_s + 15.0,
|
||||
)
|
||||
except TimeoutError:
|
||||
process.kill()
|
||||
stdout_bytes, stderr_bytes = await process.communicate()
|
||||
return _summarize_session(
|
||||
label=label,
|
||||
input_wav=input_wav,
|
||||
output_dir=output_dir,
|
||||
returncode=-1,
|
||||
stdout=stdout_bytes.decode(errors="replace"),
|
||||
stderr=(stderr_bytes.decode(errors="replace") + "\nclient timeout"),
|
||||
)
|
||||
return _summarize_session(
|
||||
label=label,
|
||||
input_wav=input_wav,
|
||||
output_dir=output_dir,
|
||||
returncode=process.returncode if process.returncode is not None else -1,
|
||||
stdout=stdout_bytes.decode(errors="replace"),
|
||||
stderr=stderr_bytes.decode(errors="replace"),
|
||||
)
|
||||
|
||||
|
||||
async def run_pair(args: argparse.Namespace) -> dict[str, object]:
|
||||
validated = _validate_pair_args(args)
|
||||
Path(validated["output_dir_a"]).mkdir(parents=True, exist_ok=True)
|
||||
Path(validated["output_dir_b"]).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
sessions = await asyncio.gather(
|
||||
_run_demo_process(
|
||||
label="a",
|
||||
args=args,
|
||||
input_wav=str(validated["input_wav_a"]),
|
||||
output_dir=str(validated["output_dir_a"]),
|
||||
),
|
||||
_run_demo_process(
|
||||
label="b",
|
||||
args=args,
|
||||
input_wav=str(validated["input_wav_b"]),
|
||||
output_dir=str(validated["output_dir_b"]),
|
||||
),
|
||||
)
|
||||
identity_isolation_ok = _identity_isolation_ok(sessions)
|
||||
min_audio_delta_ok = all(
|
||||
int(session.get("audio_delta_count", 0)) >= args.min_audio_deltas_per_session for session in sessions
|
||||
)
|
||||
response_audio_contract_ok = all(session.get("response_audio_contract_ok") is True for session in sessions)
|
||||
result = {
|
||||
"ok": (
|
||||
all(session.get("ok") is True for session in sessions)
|
||||
and identity_isolation_ok
|
||||
and min_audio_delta_ok
|
||||
and response_audio_contract_ok
|
||||
),
|
||||
"input_wavs_distinct": True,
|
||||
"output_dirs_distinct": True,
|
||||
"identity_isolation_ok": identity_isolation_ok,
|
||||
"min_audio_deltas_per_session": args.min_audio_deltas_per_session,
|
||||
"min_audio_delta_ok": min_audio_delta_ok,
|
||||
"response_audio_contract_ok": response_audio_contract_ok,
|
||||
"inputs": {
|
||||
"a": validated["input_wav_a"],
|
||||
"b": validated["input_wav_b"],
|
||||
"sha256_a": validated["input_sha256_a"],
|
||||
"sha256_b": validated["input_sha256_b"],
|
||||
},
|
||||
"sessions": sessions,
|
||||
}
|
||||
summary_output = (
|
||||
Path(args.summary_output) if args.summary_output else Path(validated["output_dir_a"]).parent / "summary.json"
|
||||
)
|
||||
summary_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
summary_output.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return result
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1")
|
||||
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
|
||||
parser.add_argument("--input-wav-a", required=True)
|
||||
parser.add_argument("--input-wav-b", required=True)
|
||||
parser.add_argument("--ref-audio", required=True)
|
||||
parser.add_argument("--output-dir-a", required=True)
|
||||
parser.add_argument("--output-dir-b", required=True)
|
||||
parser.add_argument("--summary-output")
|
||||
parser.add_argument("--chunk-ms", type=int, default=200)
|
||||
parser.add_argument("--timeout-s", type=float, default=120.0)
|
||||
parser.add_argument("--no-realtime-pacing", action="store_true")
|
||||
parser.add_argument("--require-audio", action="store_true")
|
||||
parser.add_argument("--min-audio-deltas-per-session", type=int, default=2)
|
||||
args = parser.parse_args()
|
||||
if args.min_audio_deltas_per_session < 0:
|
||||
parser.error("--min-audio-deltas-per-session must be non-negative")
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
result = asyncio.run(run_pair(parse_args()))
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
raise SystemExit(0 if result["ok"] else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,556 @@
|
||||
"""Concurrent MiniCPM-o Realtime duplex and resumable-session E2E driver."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import sys
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
import websockets
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
if str(SCRIPT_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(SCRIPT_DIR))
|
||||
|
||||
from minicpmo_realtime_duplex_scenarios import _url_with_model, run_demo # noqa: E402
|
||||
|
||||
|
||||
def _with_resume_mode(url: str) -> str:
|
||||
parts = urlsplit(url)
|
||||
query = dict(parse_qsl(parts.query, keep_blank_values=True))
|
||||
query["resume"] = "1"
|
||||
return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
|
||||
|
||||
|
||||
def _response_ids(result: dict[str, object]) -> set[str]:
|
||||
values = result.get("completed_response_ids")
|
||||
return {value for value in values if isinstance(value, str)} if isinstance(values, list) else set()
|
||||
|
||||
|
||||
def _event_session_id(event: dict[str, object]) -> str | None:
|
||||
session_id = event.get("session_id")
|
||||
if isinstance(session_id, str):
|
||||
return session_id
|
||||
inner = event.get("event")
|
||||
if isinstance(inner, dict):
|
||||
nested_session_id = inner.get("session_id")
|
||||
if isinstance(nested_session_id, str):
|
||||
return nested_session_id
|
||||
return None
|
||||
|
||||
|
||||
def _error_code(event: dict[str, object]) -> str | None:
|
||||
error = event.get("error")
|
||||
if isinstance(error, dict) and isinstance(error.get("code"), str):
|
||||
return error["code"]
|
||||
code = event.get("code")
|
||||
return code if isinstance(code, str) else None
|
||||
|
||||
|
||||
def _validate_identity_isolation(results: list[dict[str, object]]) -> bool:
|
||||
seen: set[str] = set()
|
||||
for result in results:
|
||||
current = _response_ids(result)
|
||||
if current & seen:
|
||||
return False
|
||||
seen.update(current)
|
||||
return True
|
||||
|
||||
|
||||
def _validate_semantic_isolation(
|
||||
results: list[dict[str, object]],
|
||||
*,
|
||||
input_wavs: list[str],
|
||||
expected_tokens: list[str],
|
||||
) -> bool:
|
||||
if not input_wavs:
|
||||
return True
|
||||
if len(results) != len(input_wavs):
|
||||
return False
|
||||
input_hashes = [hashlib.sha256(Path(path).read_bytes()).digest() for path in input_wavs]
|
||||
if len(set(input_hashes)) != len(input_hashes):
|
||||
return False
|
||||
if not expected_tokens:
|
||||
return True
|
||||
if len(expected_tokens) != len(results):
|
||||
return False
|
||||
for result, expected_token in zip(results, expected_tokens, strict=True):
|
||||
details = result.get("transcript_integrity")
|
||||
transcripts = (
|
||||
[str(item.get("transcript", "")) for item in details if isinstance(item, dict)]
|
||||
if isinstance(details, list)
|
||||
else []
|
||||
)
|
||||
if expected_token not in "".join(transcripts):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def _receive_until(ws, event_type: str, *, timeout_s: float) -> tuple[dict[str, object], list[dict[str, object]]]:
|
||||
async def receive() -> tuple[dict[str, object], list[dict[str, object]]]:
|
||||
events: list[dict[str, object]] = []
|
||||
while True:
|
||||
raw = await ws.recv()
|
||||
if not isinstance(raw, str):
|
||||
continue
|
||||
event = json.loads(raw)
|
||||
if not isinstance(event, dict):
|
||||
continue
|
||||
events.append(event)
|
||||
if event.get("type") == event_type:
|
||||
return event, events
|
||||
|
||||
return await asyncio.wait_for(receive(), timeout=timeout_s)
|
||||
|
||||
|
||||
async def _open_admission_session(
|
||||
args: argparse.Namespace,
|
||||
session_id: str,
|
||||
) -> tuple[object, dict[str, object]]:
|
||||
url = _url_with_model(args.url, args.model, session_id=session_id)
|
||||
ws = await websockets.connect(url, max_size=64 * 1024 * 1024)
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"session_id": session_id,
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
created, _ = await _receive_until(ws, "session.created", timeout_s=args.timeout_s)
|
||||
return ws, created
|
||||
|
||||
|
||||
async def _close_admission_session(ws, *, timeout_s: float) -> None:
|
||||
await ws.send(json.dumps({"type": "session.close"}))
|
||||
await _receive_until(ws, "session.closed", timeout_s=timeout_s)
|
||||
await ws.close()
|
||||
|
||||
|
||||
async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, object]:
|
||||
if limit < 1:
|
||||
raise ValueError("admission limit must be positive")
|
||||
prefix = f"admission-{uuid.uuid4().hex}"
|
||||
accepted: list[tuple[object, dict[str, object]]] = []
|
||||
replacement: tuple[object, dict[str, object]] | None = None
|
||||
overflow_code = None
|
||||
try:
|
||||
for index in range(limit):
|
||||
accepted.append(await _open_admission_session(args, f"{prefix}-accepted-{index}"))
|
||||
|
||||
overflow_id = f"{prefix}-overflow"
|
||||
overflow_url = _url_with_model(args.url, args.model, session_id=overflow_id)
|
||||
async with websockets.connect(overflow_url, max_size=64 * 1024 * 1024) as overflow:
|
||||
await overflow.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"session_id": overflow_id,
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
error, _ = await _receive_until(overflow, "error", timeout_s=args.timeout_s)
|
||||
overflow_code = _error_code(error)
|
||||
|
||||
first_ws, _ = accepted.pop(0)
|
||||
await _close_admission_session(first_ws, timeout_s=args.timeout_s)
|
||||
replacement = await _open_admission_session(args, f"{prefix}-replacement")
|
||||
|
||||
first_capabilities = accepted[0][1].get("session") if accepted else replacement[1].get("session")
|
||||
capabilities = first_capabilities.get("capabilities") if isinstance(first_capabilities, dict) else None
|
||||
advertised_multi = capabilities.get("supports_multi_session") if isinstance(capabilities, dict) else None
|
||||
admission_mode = capabilities.get("session_admission_mode") if isinstance(capabilities, dict) else None
|
||||
return {
|
||||
"ok": (
|
||||
overflow_code == "resource_exhausted"
|
||||
and replacement[1].get("type") == "session.created"
|
||||
and admission_mode == "engine_managed"
|
||||
),
|
||||
"configured_limit": limit,
|
||||
"accepted_before_rejection": limit,
|
||||
"overflow_error_code": overflow_code,
|
||||
"replacement_accepted": True,
|
||||
"advertised_multi_session": advertised_multi,
|
||||
"session_admission_mode": admission_mode,
|
||||
}
|
||||
finally:
|
||||
cleanup = list(accepted)
|
||||
if replacement is not None:
|
||||
cleanup.append(replacement)
|
||||
for ws, _ in cleanup:
|
||||
try:
|
||||
await _close_admission_session(ws, timeout_s=args.timeout_s)
|
||||
except Exception:
|
||||
await ws.close()
|
||||
|
||||
|
||||
async def _resume_probe(
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
session_id: str,
|
||||
expect_expired: bool = False,
|
||||
) -> dict[str, object]:
|
||||
url = _url_with_model(args.url, args.model, session_id=session_id)
|
||||
async with websockets.connect(url, max_size=64 * 1024 * 1024) as first:
|
||||
await first.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s)
|
||||
token = created.get("resume_token")
|
||||
incarnation = created.get("incarnation")
|
||||
generation = created.get("attachment_generation")
|
||||
if not isinstance(token, str) or not isinstance(incarnation, int):
|
||||
raise RuntimeError("session.created omitted resumable credentials")
|
||||
last_seq = max(
|
||||
(
|
||||
event.get("server_event_seq", 0)
|
||||
for event in first_events
|
||||
if isinstance(event.get("server_event_seq"), int)
|
||||
),
|
||||
default=0,
|
||||
)
|
||||
|
||||
delay_s = args.expire_after_s if expect_expired else args.resume_after_ms / 1000
|
||||
if delay_s > 0:
|
||||
await asyncio.sleep(delay_s)
|
||||
|
||||
resume_url = _with_resume_mode(url)
|
||||
async with websockets.connect(resume_url, max_size=64 * 1024 * 1024) as second:
|
||||
await second.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.resume",
|
||||
"session_id": session_id,
|
||||
"incarnation": incarnation,
|
||||
"resume_token": token,
|
||||
"last_received_server_event_seq": last_seq,
|
||||
}
|
||||
)
|
||||
)
|
||||
if expect_expired:
|
||||
error, error_events = await _receive_until(second, "error", timeout_s=args.timeout_s)
|
||||
error_payload = error.get("error")
|
||||
code = error_payload.get("code") if isinstance(error_payload, dict) else error.get("code")
|
||||
return {
|
||||
"ok": code == "session_resume_expired",
|
||||
"session_id": session_id,
|
||||
"expired": True,
|
||||
"error_code": code,
|
||||
"event_count": len(error_events),
|
||||
}
|
||||
resumed, replay = await _receive_until(second, "session.resumed", timeout_s=args.timeout_s)
|
||||
rotated = resumed.get("resume_token")
|
||||
if not isinstance(rotated, str) or rotated == token:
|
||||
raise RuntimeError("session.resume did not rotate the resume token")
|
||||
await second.send(json.dumps({"type": "session.heartbeat"}))
|
||||
heartbeat, heartbeat_events = await _receive_until(
|
||||
second,
|
||||
"session.heartbeat_ack",
|
||||
timeout_s=args.timeout_s,
|
||||
)
|
||||
await second.send(json.dumps({"type": "session.close"}))
|
||||
closed, close_events = await _receive_until(second, "session.closed", timeout_s=args.timeout_s)
|
||||
|
||||
replay_sequences = [event["server_event_seq"] for event in replay if isinstance(event.get("server_event_seq"), int)]
|
||||
return {
|
||||
"ok": (
|
||||
resumed.get("session_id") == session_id
|
||||
and isinstance(generation, int)
|
||||
and resumed.get("attachment_generation") == generation + 1
|
||||
and _event_session_id(heartbeat) == session_id
|
||||
and _event_session_id(closed) == session_id
|
||||
and replay_sequences == sorted(replay_sequences)
|
||||
),
|
||||
"session_id": session_id,
|
||||
"resumed_session_id": resumed.get("session_id"),
|
||||
"heartbeat_session_id": _event_session_id(heartbeat),
|
||||
"closed_session_id": _event_session_id(closed),
|
||||
"initial_attachment_generation": generation,
|
||||
"resumed_attachment_generation": resumed.get("attachment_generation"),
|
||||
"replayed_event_count": len(replay_sequences),
|
||||
"replayed_event_sequences": replay_sequences,
|
||||
"heartbeat_event_count": len(heartbeat_events),
|
||||
"close_event_count": len(close_events),
|
||||
"token_rotated": True,
|
||||
}
|
||||
|
||||
|
||||
async def _takeover_probe(
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
session_id: str,
|
||||
) -> dict[str, object]:
|
||||
url = _url_with_model(args.url, args.model, session_id=session_id)
|
||||
first = await websockets.connect(url, max_size=64 * 1024 * 1024)
|
||||
second = None
|
||||
try:
|
||||
await first.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s)
|
||||
token = created.get("resume_token")
|
||||
incarnation = created.get("incarnation")
|
||||
generation = created.get("attachment_generation")
|
||||
if not isinstance(token, str) or not isinstance(incarnation, int) or not isinstance(generation, int):
|
||||
raise RuntimeError("session.created omitted takeover credentials")
|
||||
last_seq = max(
|
||||
(
|
||||
event.get("server_event_seq", 0)
|
||||
for event in first_events
|
||||
if isinstance(event.get("server_event_seq"), int)
|
||||
),
|
||||
default=0,
|
||||
)
|
||||
|
||||
second = await websockets.connect(_with_resume_mode(url), max_size=64 * 1024 * 1024)
|
||||
await second.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.resume",
|
||||
"session_id": session_id,
|
||||
"incarnation": incarnation,
|
||||
"resume_token": token,
|
||||
"last_received_server_event_seq": last_seq,
|
||||
}
|
||||
)
|
||||
)
|
||||
resumed, replay_events = await _receive_until(second, "session.resumed", timeout_s=args.timeout_s)
|
||||
replaced, replaced_events = await _receive_until(first, "session.replaced", timeout_s=args.timeout_s)
|
||||
await asyncio.wait_for(first.wait_closed(), timeout=args.timeout_s)
|
||||
|
||||
rejected_old_writes = 0
|
||||
for _ in range(4):
|
||||
try:
|
||||
await first.send(json.dumps({"type": "session.heartbeat"}))
|
||||
except ConnectionClosed:
|
||||
rejected_old_writes += 1
|
||||
|
||||
await second.send(json.dumps({"type": "session.heartbeat"}))
|
||||
heartbeat, heartbeat_events = await _receive_until(
|
||||
second,
|
||||
"session.heartbeat_ack",
|
||||
timeout_s=args.timeout_s,
|
||||
)
|
||||
await second.send(json.dumps({"type": "session.close"}))
|
||||
closed, close_events = await _receive_until(second, "session.closed", timeout_s=args.timeout_s)
|
||||
rotated_token = resumed.get("resume_token")
|
||||
return {
|
||||
"ok": (
|
||||
resumed.get("session_id") == session_id
|
||||
and resumed.get("attachment_generation") == generation + 1
|
||||
and isinstance(rotated_token, str)
|
||||
and rotated_token != token
|
||||
and _event_session_id(replaced) == session_id
|
||||
and replaced.get("attachment_generation") == generation
|
||||
and rejected_old_writes == 4
|
||||
and _event_session_id(heartbeat) == session_id
|
||||
and _event_session_id(closed) == session_id
|
||||
),
|
||||
"session_id": session_id,
|
||||
"initial_attachment_generation": generation,
|
||||
"resumed_attachment_generation": resumed.get("attachment_generation"),
|
||||
"replaced_attachment_generation": replaced.get("attachment_generation"),
|
||||
"token_rotated": isinstance(rotated_token, str) and rotated_token != token,
|
||||
"old_attachment_closed": True,
|
||||
"rejected_old_writes": rejected_old_writes,
|
||||
"replay_event_count": len(replay_events),
|
||||
"replaced_event_count": len(replaced_events),
|
||||
"heartbeat_event_count": len(heartbeat_events),
|
||||
"close_event_count": len(close_events),
|
||||
}
|
||||
finally:
|
||||
if second is not None:
|
||||
await second.close()
|
||||
await first.close()
|
||||
|
||||
|
||||
def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace:
|
||||
validation_mode = "response-required" if args.response_required else "model-policy"
|
||||
input_wav = args.session_input_wav[index] if args.session_input_wav else args.input_wav
|
||||
return SimpleNamespace(
|
||||
url=args.url,
|
||||
model=args.model,
|
||||
session_id=f"multi-{index}-{uuid.uuid4().hex}",
|
||||
input_wav=input_wav,
|
||||
ref_audio=args.ref_audio,
|
||||
turn_input_wav=list(args.turn_input_wav),
|
||||
output_dir=str(Path(args.output_dir) / f"session_{index:02d}"),
|
||||
output_audio_format="pcm16",
|
||||
chunk_ms=args.chunk_ms,
|
||||
realtime_input=args.realtime_input,
|
||||
first_turn_ms=args.first_turn_ms,
|
||||
turn_duration_ms=list(args.turn_duration_ms),
|
||||
first_turn_transcript=f"session {index} input",
|
||||
omit_transcript_hints=True,
|
||||
validation_mode=validation_mode,
|
||||
temperature=args.temperature,
|
||||
scenario="sequential",
|
||||
require_audio=args.response_required,
|
||||
require_distinct_inputs=False,
|
||||
expect_empty_turn=[],
|
||||
short_ack_ms=350,
|
||||
turns=args.turns,
|
||||
timeout_s=args.timeout_s,
|
||||
model_policy_settle_ms=args.model_policy_settle_ms,
|
||||
)
|
||||
|
||||
|
||||
async def run_multi_session(args: argparse.Namespace) -> dict[str, object]:
|
||||
if args.sessions < 1:
|
||||
raise ValueError("--sessions must be positive")
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
resume_result = None
|
||||
if args.disconnect_session_index is not None:
|
||||
if not 0 <= args.disconnect_session_index < args.sessions:
|
||||
raise ValueError("--disconnect-session-index is outside the session range")
|
||||
resume_result = await _resume_probe(
|
||||
args,
|
||||
session_id=f"resume-{args.disconnect_session_index}-{uuid.uuid4().hex}",
|
||||
)
|
||||
takeover_result = None
|
||||
if args.takeover_session_index is not None:
|
||||
if not 0 <= args.takeover_session_index < args.sessions:
|
||||
raise ValueError("--takeover-session-index is outside the session range")
|
||||
takeover_result = await _takeover_probe(
|
||||
args,
|
||||
session_id=f"takeover-{args.takeover_session_index}-{uuid.uuid4().hex}",
|
||||
)
|
||||
expiry_result = None
|
||||
if args.expire_session_index is not None:
|
||||
if not 0 <= args.expire_session_index < args.sessions:
|
||||
raise ValueError("--expire-session-index is outside the session range")
|
||||
expiry_result = await _resume_probe(
|
||||
args,
|
||||
session_id=f"expire-{args.expire_session_index}-{uuid.uuid4().hex}",
|
||||
expect_expired=True,
|
||||
)
|
||||
admission_result = (
|
||||
await _admission_probe(args, limit=args.verify_admission_limit)
|
||||
if args.verify_admission_limit is not None
|
||||
else None
|
||||
)
|
||||
|
||||
session_results = await asyncio.gather(
|
||||
*(run_demo(_demo_args(args, index)) for index in range(args.sessions)),
|
||||
return_exceptions=True,
|
||||
)
|
||||
failures = [repr(result) for result in session_results if isinstance(result, BaseException)]
|
||||
completed = [result for result in session_results if isinstance(result, dict)]
|
||||
identity_isolation_ok = _validate_identity_isolation(completed)
|
||||
semantic_isolation_ok = _validate_semantic_isolation(
|
||||
completed,
|
||||
input_wavs=list(args.session_input_wav),
|
||||
expected_tokens=list(args.session_expected_token),
|
||||
)
|
||||
result = {
|
||||
"ok": (
|
||||
not failures
|
||||
and len(completed) == args.sessions
|
||||
and all(item.get("ok") is True for item in completed)
|
||||
and identity_isolation_ok
|
||||
and semantic_isolation_ok
|
||||
and (resume_result is None or resume_result.get("ok") is True)
|
||||
and (takeover_result is None or takeover_result.get("ok") is True)
|
||||
and (expiry_result is None or expiry_result.get("ok") is True)
|
||||
and (admission_result is None or admission_result.get("ok") is True)
|
||||
),
|
||||
"session_count": args.sessions,
|
||||
"identity_isolation_ok": identity_isolation_ok,
|
||||
"semantic_isolation_ok": semantic_isolation_ok,
|
||||
"resume": resume_result,
|
||||
"takeover": takeover_result,
|
||||
"expiry": expiry_result,
|
||||
"admission": admission_result,
|
||||
"failures": failures,
|
||||
"sessions": completed,
|
||||
}
|
||||
(output_dir / "summary.json").write_text(
|
||||
json.dumps(result, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--url", default="ws://127.0.0.1:8113/v1/realtime?duplex=1")
|
||||
parser.add_argument("--base-url", help="Deprecated alias; /v1/realtime is appended when supplied.")
|
||||
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
|
||||
parser.add_argument("--sessions", type=int, default=2)
|
||||
parser.add_argument("--input-wav", required=True)
|
||||
parser.add_argument("--ref-audio", help="Optional WAV used as the MiniCPM-o voice prompt for every session.")
|
||||
parser.add_argument("--session-input-wav", action="append", default=[])
|
||||
parser.add_argument("--session-expected-token", action="append", default=[])
|
||||
parser.add_argument("--turn-input-wav", action="append", default=[])
|
||||
parser.add_argument("--output-dir", default="/tmp/minicpmo_pr3907_multi_session_e2e")
|
||||
parser.add_argument("--realtime-input", action="store_true")
|
||||
parser.add_argument("--chunk-ms", type=int, default=200)
|
||||
parser.add_argument("--turns", type=int, default=1)
|
||||
parser.add_argument("--first-turn-ms", type=int, default=1400)
|
||||
parser.add_argument("--turn-duration-ms", type=int, action="append", default=[])
|
||||
parser.add_argument("--response-required", action="store_true")
|
||||
parser.add_argument("--temperature", type=float, default=None)
|
||||
parser.add_argument("--disconnect-session-index", type=int)
|
||||
parser.add_argument("--takeover-session-index", type=int)
|
||||
parser.add_argument("--resume-after-ms", type=int, default=1000)
|
||||
parser.add_argument("--expire-session-index", type=int)
|
||||
parser.add_argument("--expire-after-s", type=float, default=6.0)
|
||||
parser.add_argument("--verify-admission-limit", type=int)
|
||||
parser.add_argument("--model-policy-settle-ms", type=int, default=600)
|
||||
parser.add_argument("--timeout-s", type=float, default=120.0)
|
||||
args = parser.parse_args()
|
||||
if args.base_url:
|
||||
args.url = args.base_url.rstrip("/") + "/v1/realtime?duplex=1"
|
||||
if args.session_input_wav and len(args.session_input_wav) != args.sessions:
|
||||
parser.error("provide exactly one --session-input-wav per session")
|
||||
if args.session_expected_token and len(args.session_expected_token) != args.sessions:
|
||||
parser.error("provide exactly one --session-expected-token per session")
|
||||
if args.session_expected_token and not args.session_input_wav:
|
||||
parser.error("--session-expected-token requires --session-input-wav")
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
result = asyncio.run(run_multi_session(parse_args()))
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
raise SystemExit(0 if result["ok"] else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,467 @@
|
||||
"""Validate MiniCPM-o Realtime duplex soft-interrupt delta streaming.
|
||||
|
||||
This E2E driver runs the public ``realtime_duplex_demo.py`` against a live
|
||||
duplex backend. Arbitrary audio defaults to the model-policy lifecycle contract.
|
||||
The stronger response-required mode binds the two-response contract to a known
|
||||
input checksum and expected second response.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import uuid
|
||||
import wave
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
DEMO_PATH = REPO_ROOT / "examples/online_serving/minicpmo/realtime_duplex_demo.py"
|
||||
AUDIO_DELTA_EVENTS = {"response.audio.delta", "response.output_audio.delta"}
|
||||
TRANSCRIPT_DELTA_EVENTS = {
|
||||
"response.audio_transcript.delta",
|
||||
"response.output_text.delta",
|
||||
}
|
||||
|
||||
|
||||
def _canonical_path(path: str) -> Path:
|
||||
return Path(path).expanduser().resolve(strict=False)
|
||||
|
||||
|
||||
def _validate_input_sha256(path: Path, expected: str) -> str:
|
||||
actual = hashlib.sha256(path.read_bytes()).hexdigest()
|
||||
if actual != expected.lower():
|
||||
raise ValueError(f"input WAV SHA256 mismatch: expected {expected.lower()}, got {actual}")
|
||||
return actual
|
||||
|
||||
|
||||
def _read_json(path: Path) -> dict[str, object]:
|
||||
if not path.is_file():
|
||||
return {}
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
|
||||
def _read_events(path: Path) -> list[dict[str, object]]:
|
||||
if not path.is_file():
|
||||
return []
|
||||
events: list[dict[str, object]] = []
|
||||
for line in path.read_text(encoding="utf-8").splitlines():
|
||||
if not line.strip():
|
||||
continue
|
||||
event = json.loads(line)
|
||||
if isinstance(event, dict):
|
||||
events.append(event)
|
||||
return events
|
||||
|
||||
|
||||
def _event_response_id(event: dict[str, object]) -> str | None:
|
||||
response_id = event.get("response_id")
|
||||
if isinstance(response_id, str):
|
||||
return response_id
|
||||
response = event.get("response")
|
||||
if isinstance(response, dict) and isinstance(response.get("id"), str):
|
||||
return response["id"]
|
||||
return None
|
||||
|
||||
|
||||
def _event_received_at_s(event: dict[str, object]) -> float | None:
|
||||
received_at = event.get("_client_received_at_s")
|
||||
if isinstance(received_at, int | float):
|
||||
return float(received_at)
|
||||
return None
|
||||
|
||||
|
||||
def _first_event_index(events: list[dict[str, object]], event_type: str) -> int | None:
|
||||
for index, event in enumerate(events):
|
||||
if event.get("type") == event_type:
|
||||
return index
|
||||
return None
|
||||
|
||||
|
||||
def _event_counts(events: list[dict[str, object]]) -> dict[str, int]:
|
||||
counts: dict[str, int] = {}
|
||||
for event in events:
|
||||
event_type = str(event.get("type"))
|
||||
counts[event_type] = counts.get(event_type, 0) + 1
|
||||
return counts
|
||||
|
||||
|
||||
def _response_created_ids(events: list[dict[str, object]]) -> list[str]:
|
||||
response_ids: list[str] = []
|
||||
for event in events:
|
||||
if event.get("type") != "response.created":
|
||||
continue
|
||||
response_id = _event_response_id(event)
|
||||
if response_id and response_id not in response_ids:
|
||||
response_ids.append(response_id)
|
||||
return response_ids
|
||||
|
||||
|
||||
def _audio_delta_duration_ms(event: dict[str, object]) -> float | None:
|
||||
encoded = event.get("delta") or event.get("audio")
|
||||
if not isinstance(encoded, str) or not encoded:
|
||||
return None
|
||||
try:
|
||||
raw = base64.b64decode(encoded)
|
||||
except ValueError:
|
||||
return None
|
||||
fmt = str(event.get("format") or event.get("audio_format") or "pcm16").lower()
|
||||
sample_rate = event.get("sample_rate_hz") or event.get("sample_rate") or 24_000
|
||||
if not isinstance(sample_rate, int | float) or sample_rate <= 0:
|
||||
return None
|
||||
if "wav" in fmt:
|
||||
try:
|
||||
with wave.open(io.BytesIO(raw), "rb") as wav_file:
|
||||
return round((wav_file.getnframes() * 1000.0) / wav_file.getframerate(), 3)
|
||||
except (EOFError, wave.Error):
|
||||
return None
|
||||
bytes_per_sample = 4 if "f32" in fmt else 2
|
||||
return round((len(raw) * 1000.0) / (bytes_per_sample * float(sample_rate)), 3)
|
||||
|
||||
|
||||
def _normalize_text(text: str) -> str:
|
||||
return "".join(text.split())
|
||||
|
||||
|
||||
def _compact_sequence(events: list[dict[str, object]]) -> list[str]:
|
||||
sequence: list[str] = []
|
||||
for event in events:
|
||||
event_type = event.get("type")
|
||||
response_id = _event_response_id(event)
|
||||
suffix = f" {response_id[-8:]}" if response_id else ""
|
||||
if event_type == "response.listen":
|
||||
sequence.append("listen")
|
||||
elif event_type == "response.created":
|
||||
sequence.append(f"response.created{suffix}")
|
||||
elif event_type == "response.speak":
|
||||
sequence.append(f"speak{suffix}")
|
||||
elif event_type in AUDIO_DELTA_EVENTS:
|
||||
sequence.append(f"audio.delta{suffix}")
|
||||
elif event_type in TRANSCRIPT_DELTA_EVENTS:
|
||||
delta = str(event.get("delta") or "")
|
||||
sequence.append(f"text.delta{suffix}> {delta}")
|
||||
elif event_type == "response.done":
|
||||
sequence.append(f"response.done{suffix}")
|
||||
elif event_type == "input_audio_buffer.committed":
|
||||
sequence.append("input.committed")
|
||||
return sequence
|
||||
|
||||
|
||||
def _response_summary(
|
||||
events: list[dict[str, object]],
|
||||
response_id: str,
|
||||
*,
|
||||
t0: float | None,
|
||||
) -> dict[str, object]:
|
||||
created_indices = [
|
||||
index
|
||||
for index, event in enumerate(events)
|
||||
if event.get("type") == "response.created" and _event_response_id(event) == response_id
|
||||
]
|
||||
done_indices = [
|
||||
index
|
||||
for index, event in enumerate(events)
|
||||
if event.get("type") == "response.done" and _event_response_id(event) == response_id
|
||||
]
|
||||
audio_indices = [
|
||||
index
|
||||
for index, event in enumerate(events)
|
||||
if event.get("type") in AUDIO_DELTA_EVENTS and _event_response_id(event) == response_id
|
||||
]
|
||||
transcript_indices = [
|
||||
index
|
||||
for index, event in enumerate(events)
|
||||
if event.get("type") in TRANSCRIPT_DELTA_EVENTS and _event_response_id(event) == response_id
|
||||
]
|
||||
done_index = done_indices[0] if done_indices else None
|
||||
audio_before_done_ok = done_index is not None and all(index < done_index for index in audio_indices)
|
||||
stale_audio_count = sum(index > done_index for index in audio_indices) if done_index is not None else 0
|
||||
|
||||
def offset_ms(index: int) -> float | None:
|
||||
if t0 is None:
|
||||
return None
|
||||
received_at = _event_received_at_s(events[index])
|
||||
return round((received_at - t0) * 1000.0, 3) if received_at is not None else None
|
||||
|
||||
return {
|
||||
"response_id": response_id,
|
||||
"created_index": created_indices[0] if created_indices else None,
|
||||
"done_indices": done_indices,
|
||||
"created_offset_ms": offset_ms(created_indices[0]) if created_indices else None,
|
||||
"done_offset_ms": offset_ms(done_indices[0]) if done_indices else None,
|
||||
"audio_delta_count": len(audio_indices),
|
||||
"transcript_delta_count": len(transcript_indices),
|
||||
"audio_delta_offsets_ms": [offset_ms(index) for index in audio_indices],
|
||||
"audio_duration_ms": [_audio_delta_duration_ms(events[index]) for index in audio_indices],
|
||||
"transcript": "".join(str(events[index].get("delta") or "") for index in transcript_indices),
|
||||
"one_done": len(done_indices) == 1,
|
||||
"audio_before_done_ok": audio_before_done_ok,
|
||||
"stale_audio_count": stale_audio_count,
|
||||
}
|
||||
|
||||
|
||||
def summarize_artifacts(
|
||||
*,
|
||||
output_dir: Path,
|
||||
validation_mode: str,
|
||||
min_responses: int,
|
||||
min_audio_deltas_per_response: int,
|
||||
expect_second_response_substring: str | None,
|
||||
) -> dict[str, object]:
|
||||
if validation_mode not in {"model-policy", "response-required"}:
|
||||
raise ValueError(f"unsupported validation mode: {validation_mode}")
|
||||
result = _read_json(output_dir / "result.json")
|
||||
events = _read_events(output_dir / "events.jsonl")
|
||||
response_ids = _response_created_ids(events)
|
||||
times = [_event_received_at_s(event) for event in events]
|
||||
clean_times = [received_at for received_at in times if received_at is not None]
|
||||
t0 = min(clean_times) if clean_times else None
|
||||
response_summaries = [_response_summary(events, response_id, t0=t0) for response_id in response_ids]
|
||||
commit_index = _first_event_index(events, "input_audio_buffer.committed")
|
||||
first_created_index = next(
|
||||
(summary["created_index"] for summary in response_summaries if isinstance(summary.get("created_index"), int)),
|
||||
None,
|
||||
)
|
||||
first_done_index = (
|
||||
response_summaries[0]["done_indices"][0]
|
||||
if response_summaries and response_summaries[0]["done_indices"]
|
||||
else None
|
||||
)
|
||||
second_created_index = (
|
||||
response_summaries[1]["created_index"]
|
||||
if len(response_summaries) >= 2 and isinstance(response_summaries[1].get("created_index"), int)
|
||||
else None
|
||||
)
|
||||
last_done_index = next(
|
||||
(summary["done_indices"][0] for summary in reversed(response_summaries) if summary.get("done_indices")),
|
||||
None,
|
||||
)
|
||||
listen_indices = [index for index, event in enumerate(events) if event.get("type") == "response.listen"]
|
||||
effective_min_responses = min_responses if validation_mode == "response-required" else 1
|
||||
enough_responses = len(response_summaries) >= effective_min_responses
|
||||
response_lifecycle_ok = enough_responses and all(bool(summary.get("one_done")) for summary in response_summaries)
|
||||
multi_delta_ok = enough_responses and all(
|
||||
int(summary.get("audio_delta_count", 0)) >= min_audio_deltas_per_response for summary in response_summaries
|
||||
)
|
||||
response_audio_contract_ok = enough_responses and all(
|
||||
summary.get("audio_before_done_ok") is True and int(summary.get("stale_audio_count", 0)) == 0
|
||||
for summary in response_summaries
|
||||
)
|
||||
response_before_final_commit = (
|
||||
first_created_index is not None and commit_index is not None and first_created_index < commit_index
|
||||
)
|
||||
second_response_before_final_commit = (
|
||||
second_created_index is not None and commit_index is not None and second_created_index < commit_index
|
||||
)
|
||||
listen_before_first_response = first_created_index is not None and any(
|
||||
index < first_created_index for index in listen_indices
|
||||
)
|
||||
listen_between_responses = (
|
||||
first_done_index is not None
|
||||
and second_created_index is not None
|
||||
and any(first_done_index < index < second_created_index for index in listen_indices)
|
||||
)
|
||||
listen_after_last_done = last_done_index is not None and any(index > last_done_index for index in listen_indices)
|
||||
listen_after_response_before_commit = (
|
||||
last_done_index is not None
|
||||
and commit_index is not None
|
||||
and any(last_done_index < index < commit_index for index in listen_indices)
|
||||
)
|
||||
final_listen_after_commit = commit_index is not None and any(index > commit_index for index in listen_indices)
|
||||
transcript = "".join(str(summary.get("transcript") or "") for summary in response_summaries)
|
||||
second_response_transcript = (
|
||||
str(response_summaries[1].get("transcript") or "") if len(response_summaries) >= 2 else ""
|
||||
)
|
||||
second_response_transcript_expectation_ok = not expect_second_response_substring or _normalize_text(
|
||||
expect_second_response_substring
|
||||
) in _normalize_text(second_response_transcript)
|
||||
error_events = [
|
||||
event
|
||||
for event in events
|
||||
if event.get("type") == "error"
|
||||
or str(event.get("type") or "").endswith(".cancelled")
|
||||
or str(event.get("type") or "").endswith(".truncated")
|
||||
]
|
||||
cancelled_count = sum(
|
||||
1
|
||||
for event in events
|
||||
if event.get("type") == "response.done"
|
||||
and isinstance(event.get("response"), dict)
|
||||
and event["response"].get("status") == "cancelled"
|
||||
)
|
||||
result_ok = result.get("ok") is True
|
||||
common_contract_ok = bool(
|
||||
result_ok
|
||||
and not error_events
|
||||
and cancelled_count == 0
|
||||
and enough_responses
|
||||
and response_lifecycle_ok
|
||||
and multi_delta_ok
|
||||
and response_audio_contract_ok
|
||||
and response_before_final_commit
|
||||
and listen_after_response_before_commit
|
||||
and final_listen_after_commit
|
||||
)
|
||||
mode_contract_ok = validation_mode == "model-policy" or (
|
||||
second_response_before_final_commit
|
||||
and listen_before_first_response
|
||||
and listen_between_responses
|
||||
and listen_after_last_done
|
||||
and second_response_transcript_expectation_ok
|
||||
)
|
||||
ok = common_contract_ok and mode_contract_ok
|
||||
return {
|
||||
"ok": ok,
|
||||
"result_ok": result_ok,
|
||||
"validation_mode": validation_mode,
|
||||
"event_counts": _event_counts(events),
|
||||
"response_ids": response_ids,
|
||||
"response_summaries": response_summaries,
|
||||
"min_responses": min_responses,
|
||||
"effective_min_responses": effective_min_responses,
|
||||
"min_audio_deltas_per_response": min_audio_deltas_per_response,
|
||||
"enough_responses": enough_responses,
|
||||
"response_lifecycle_ok": response_lifecycle_ok,
|
||||
"multi_delta_ok": multi_delta_ok,
|
||||
"response_audio_contract_ok": response_audio_contract_ok,
|
||||
"response_before_final_commit": response_before_final_commit,
|
||||
"second_response_before_final_commit": second_response_before_final_commit,
|
||||
"listen_before_first_response": listen_before_first_response,
|
||||
"listen_between_responses": listen_between_responses,
|
||||
"listen_after_last_done": listen_after_last_done,
|
||||
"listen_after_response_before_commit": listen_after_response_before_commit,
|
||||
"final_listen_after_commit": final_listen_after_commit,
|
||||
"transcript": transcript,
|
||||
"second_response_transcript": second_response_transcript,
|
||||
"expect_second_response_substring": expect_second_response_substring,
|
||||
"second_response_transcript_expectation_ok": second_response_transcript_expectation_ok,
|
||||
"error_count": len(error_events),
|
||||
"cancelled_count": cancelled_count,
|
||||
"compact_sequence": _compact_sequence(events),
|
||||
"output_dir": str(output_dir),
|
||||
}
|
||||
|
||||
|
||||
async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]:
|
||||
input_wav = _canonical_path(args.input_wav)
|
||||
output_dir = _canonical_path(args.output_dir)
|
||||
if not input_wav.is_file():
|
||||
raise ValueError(f"input WAV does not exist: {input_wav}")
|
||||
input_sha256 = hashlib.sha256(input_wav.read_bytes()).hexdigest()
|
||||
if args.input_sha256:
|
||||
input_sha256 = _validate_input_sha256(input_wav, args.input_sha256)
|
||||
command = [
|
||||
sys.executable,
|
||||
str(DEMO_PATH),
|
||||
"--url",
|
||||
args.url,
|
||||
"--model",
|
||||
args.model,
|
||||
"--input-wav",
|
||||
str(input_wav),
|
||||
"--output-dir",
|
||||
str(output_dir),
|
||||
"--chunk-ms",
|
||||
str(args.chunk_ms),
|
||||
"--timeout-s",
|
||||
str(args.timeout_s),
|
||||
"--session-id",
|
||||
f"duplex-soft-interrupt-{uuid.uuid4().hex}",
|
||||
]
|
||||
command.extend(["--ref-audio", str(_canonical_path(args.ref_audio))])
|
||||
if args.require_audio:
|
||||
command.append("--require-audio")
|
||||
if args.no_realtime_pacing:
|
||||
command.append("--no-realtime-pacing")
|
||||
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*command,
|
||||
cwd=str(REPO_ROOT),
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
try:
|
||||
stdout_bytes, stderr_bytes = await asyncio.wait_for(
|
||||
process.communicate(),
|
||||
timeout=args.timeout_s + 30.0,
|
||||
)
|
||||
except TimeoutError:
|
||||
process.kill()
|
||||
stdout_bytes, stderr_bytes = await process.communicate()
|
||||
return {
|
||||
"ok": False,
|
||||
"returncode": -1,
|
||||
"stdout_tail": stdout_bytes.decode(errors="replace")[-4000:],
|
||||
"stderr_tail": (stderr_bytes.decode(errors="replace") + "\nclient timeout")[-4000:],
|
||||
"output_dir": str(output_dir),
|
||||
}
|
||||
|
||||
summary = summarize_artifacts(
|
||||
output_dir=output_dir,
|
||||
validation_mode=args.validation_mode,
|
||||
min_responses=args.min_responses,
|
||||
min_audio_deltas_per_response=args.min_audio_deltas_per_response,
|
||||
expect_second_response_substring=args.expect_second_response_substring,
|
||||
)
|
||||
summary.update(
|
||||
{
|
||||
"returncode": process.returncode if process.returncode is not None else -1,
|
||||
"input_wav": str(input_wav),
|
||||
"input_sha256": input_sha256,
|
||||
"stdout_tail": stdout_bytes.decode(errors="replace")[-4000:],
|
||||
"stderr_tail": stderr_bytes.decode(errors="replace")[-4000:],
|
||||
}
|
||||
)
|
||||
summary["ok"] = bool(summary["ok"] and summary["returncode"] == 0)
|
||||
summary_output = Path(args.summary_output) if args.summary_output else output_dir / "summary.json"
|
||||
summary_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
summary_output.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return summary
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1")
|
||||
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
|
||||
parser.add_argument("--input-wav", required=True)
|
||||
parser.add_argument("--ref-audio", required=True)
|
||||
parser.add_argument("--output-dir", required=True)
|
||||
parser.add_argument("--summary-output")
|
||||
parser.add_argument("--chunk-ms", type=int, default=200)
|
||||
parser.add_argument("--timeout-s", type=float, default=120.0)
|
||||
parser.add_argument("--require-audio", action="store_true")
|
||||
parser.add_argument("--no-realtime-pacing", action="store_true")
|
||||
parser.add_argument(
|
||||
"--validation-mode",
|
||||
choices=("model-policy", "response-required"),
|
||||
default="model-policy",
|
||||
)
|
||||
parser.add_argument("--min-responses", type=int, default=2)
|
||||
parser.add_argument("--min-audio-deltas-per-response", type=int, default=2)
|
||||
parser.add_argument("--input-sha256")
|
||||
parser.add_argument("--expect-second-response-substring")
|
||||
args = parser.parse_args()
|
||||
if args.validation_mode == "response-required" and args.min_responses < 2:
|
||||
parser.error("--min-responses must be at least 2")
|
||||
if args.validation_mode == "response-required" and not args.input_sha256:
|
||||
parser.error("--input-sha256 is required in response-required mode")
|
||||
if args.validation_mode == "response-required" and not args.expect_second_response_substring:
|
||||
parser.error("--expect-second-response-substring is required in response-required mode")
|
||||
if args.min_audio_deltas_per_response < 1:
|
||||
parser.error("--min-audio-deltas-per-response must be positive")
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
result = asyncio.run(run_soft_interrupt(parse_args()))
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
raise SystemExit(0 if result["ok"] else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,4 +1,5 @@
|
||||
import pytest
|
||||
import torch
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
@@ -64,6 +65,48 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
|
||||
output_processor.add_request.assert_not_called()
|
||||
|
||||
|
||||
def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
params = SamplingParams(max_tokens=8)
|
||||
engine.default_sampling_params_list = [params]
|
||||
engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
|
||||
engine.supported_tasks = ("speech",)
|
||||
|
||||
input_processor = mocker.Mock()
|
||||
input_processor.process_inputs.return_value = _make_engine_core_request()
|
||||
engine.input_processor = input_processor
|
||||
|
||||
output_processor = mocker.Mock()
|
||||
engine.output_processors = [output_processor]
|
||||
|
||||
hidden = torch.arange(6, dtype=torch.float32).reshape(2, 3)
|
||||
prompt = {
|
||||
"prompt_token_ids": [1, 1, 1],
|
||||
"model_intermediate_buffer": {
|
||||
"ids": {"tts": [11, 12]},
|
||||
"hidden_states": {"tts": hidden},
|
||||
},
|
||||
}
|
||||
|
||||
msg = engine._build_add_request_message(
|
||||
request_id="req-1",
|
||||
prompt=prompt,
|
||||
sampling_params_list=[params],
|
||||
final_stage_id=0,
|
||||
arrival_time=0.0,
|
||||
)
|
||||
|
||||
request = msg.prompt
|
||||
assert isinstance(request, OmniEngineCoreRequest)
|
||||
assert request.additional_information is not None
|
||||
assert request.additional_information.entries["global_request_id"].list_data == ["req-1"]
|
||||
assert request.additional_information.entries["omni_final_stage_id"].scalar_data == 0
|
||||
assert isinstance(request.model_intermediate_buffer, dict)
|
||||
info = request.model_intermediate_buffer
|
||||
assert info["ids"]["tts"] == [11, 12]
|
||||
assert torch.equal(info["hidden_states"]["tts"], hidden)
|
||||
|
||||
|
||||
def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
params = SamplingParams(max_tokens=8)
|
||||
|
||||
@@ -4,13 +4,33 @@ Focuses on the critical behavior: when the orchestrator thread dies,
|
||||
subsequent attempts to collect output raise RuntimeError.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import queue
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from vllm_omni.engine.async_engine_utils import weak_shutdown_async_omni_engine
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
||||
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
|
||||
from vllm_omni.engine.messages import (
|
||||
CollectiveRPCResultMessage,
|
||||
ErrorMessage,
|
||||
OutputMessage,
|
||||
)
|
||||
from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_control_client import (
|
||||
DuplexControlClient,
|
||||
DuplexControlRequestError,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import (
|
||||
AppendDuplexInputMessage,
|
||||
DuplexControlResultMessage,
|
||||
DuplexFence,
|
||||
SignalDuplexTurnMessage,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -26,6 +46,33 @@ def _make_engine(output_queue, mocker: MockerFixture, *, thread_alive: bool = Tr
|
||||
return engine
|
||||
|
||||
|
||||
def test_weak_shutdown_closes_rpc_router_before_joining_orchestrator(mocker: MockerFixture):
|
||||
request_queue = mocker.MagicMock()
|
||||
output_queue = mocker.MagicMock()
|
||||
rpc_output_queue = mocker.MagicMock()
|
||||
router = mocker.MagicMock()
|
||||
orchestrator_thread = mocker.MagicMock()
|
||||
orchestrator_thread.is_alive.return_value = True
|
||||
|
||||
def assert_router_closed_before_join(*args, **kwargs):
|
||||
router.close.assert_called_once_with()
|
||||
|
||||
orchestrator_thread.join.side_effect = assert_router_closed_before_join
|
||||
|
||||
weak_shutdown_async_omni_engine(
|
||||
orchestrator_thread,
|
||||
request_queue,
|
||||
output_queue,
|
||||
rpc_output_queue,
|
||||
router,
|
||||
)
|
||||
|
||||
request_queue.sync_q.put.assert_called_once()
|
||||
assert request_queue.sync_q.put.call_args.kwargs["timeout"] > 0
|
||||
orchestrator_thread.join.assert_called_once()
|
||||
assert orchestrator_thread.join.call_args.kwargs["timeout"] > 0
|
||||
|
||||
|
||||
def test_try_get_output_raises_after_orchestrator_dies(mocker: MockerFixture):
|
||||
"""Draining remaining results then hitting an empty queue with a dead
|
||||
orchestrator must raise RuntimeError so callers know the pipeline is gone."""
|
||||
@@ -56,18 +103,17 @@ def test_try_get_output_raises_after_orchestrator_dies(mocker: MockerFixture):
|
||||
@pytest.mark.asyncio
|
||||
async def test_try_get_output_async_raises_after_orchestrator_dies(mocker: MockerFixture):
|
||||
"""Same scenario as above but for the async variant."""
|
||||
mock_queue = mocker.MagicMock()
|
||||
mock_queue.sync_q.get_nowait.side_effect = [
|
||||
raw_queue = queue.Queue()
|
||||
raw_queue.put_nowait(
|
||||
OutputMessage(
|
||||
request_id="r1",
|
||||
stage_id=0,
|
||||
engine_outputs=OmniRequestOutput(request_id="r1"),
|
||||
finished=False,
|
||||
),
|
||||
queue.Empty,
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
engine = _make_engine(mock_queue, mocker, thread_alive=True)
|
||||
engine = _make_engine(SimpleNamespace(sync_q=raw_queue), mocker, thread_alive=True)
|
||||
|
||||
assert (await engine.try_get_output_async()).request_id == "r1"
|
||||
|
||||
@@ -102,12 +148,495 @@ async def test_fatal_error_message_surfaces_through_try_get_output_async(mocker:
|
||||
"""Async variant of the fatal error message test."""
|
||||
fatal_msg = ErrorMessage(error="Orchestrator thread crashed", fatal=True)
|
||||
|
||||
mock_queue = mocker.MagicMock()
|
||||
mock_queue.sync_q.get_nowait.return_value = fatal_msg
|
||||
raw_queue = queue.Queue()
|
||||
raw_queue.put_nowait(fatal_msg)
|
||||
|
||||
engine = _make_engine(mock_queue, mocker, thread_alive=False)
|
||||
engine = _make_engine(SimpleNamespace(sync_q=raw_queue), mocker, thread_alive=False)
|
||||
|
||||
msg = await engine.try_get_output_async()
|
||||
assert msg is not None
|
||||
assert msg.type == "error"
|
||||
assert msg.fatal is True
|
||||
|
||||
|
||||
def test_output_remains_on_shared_output_path(mocker: MockerFixture):
|
||||
raw_queue = queue.Queue()
|
||||
raw_queue.put_nowait(
|
||||
OutputMessage(
|
||||
request_id="legacy-duplex-request",
|
||||
stage_id=0,
|
||||
engine_outputs=OmniRequestOutput(request_id="legacy-duplex-request"),
|
||||
finished=False,
|
||||
)
|
||||
)
|
||||
engine = _make_engine(SimpleNamespace(sync_q=raw_queue), mocker)
|
||||
|
||||
output = engine.try_get_output(timeout=0.01)
|
||||
|
||||
assert output.request_id == "legacy-duplex-request"
|
||||
|
||||
|
||||
def test_duplex_control_client_uses_engine_owned_correlated_transport():
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._duplex_control_client = None
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(queue.Queue(), queue.Queue())
|
||||
|
||||
try:
|
||||
client = engine._get_duplex_control_client()
|
||||
|
||||
assert isinstance(client, DuplexControlClient)
|
||||
assert client._transport is engine._correlated_rpc_client
|
||||
assert not hasattr(engine, "_submit_rpc_and_wait")
|
||||
finally:
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_duplex_control_api_requires_explicit_fence():
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
|
||||
with pytest.raises(TypeError, match="fence"):
|
||||
engine.open_duplex_session("sid")
|
||||
with pytest.raises(TypeError, match="fence"):
|
||||
engine.append_duplex_input("sid", mode="append_audio_chunk", payload={})
|
||||
with pytest.raises(TypeError, match="fence"):
|
||||
engine.signal_duplex_turn("sid", event="turn.end")
|
||||
with pytest.raises(TypeError, match="fence"):
|
||||
engine.close_duplex_session("sid")
|
||||
|
||||
|
||||
def test_open_duplex_session_waits_for_control_ack(mocker: MockerFixture):
|
||||
request_q = queue.Queue()
|
||||
rpc_q = queue.Queue()
|
||||
control_result = DuplexControlResultMessage(
|
||||
control_id="ctrl-1",
|
||||
fence=DuplexFence("sid"),
|
||||
operation="open",
|
||||
session_id="sid",
|
||||
ok=False,
|
||||
stage_results=[{"stage_id": 0, "replica_id": 0, "result": {"supported": False}}],
|
||||
unsupported_count=1,
|
||||
error_count=0,
|
||||
)
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
engine._duplex_control_client = None
|
||||
mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-1"))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
pending = executor.submit(engine.open_duplex_session, "sid", fence=DuplexFence("sid"), timeout=1)
|
||||
msg = request_q.get(timeout=1)
|
||||
assert msg.type == "open_duplex_session"
|
||||
assert msg.control_id == "ctrl-1"
|
||||
rpc_q.put(control_result)
|
||||
with pytest.raises(DuplexControlRequestError) as exc_info:
|
||||
pending.result(timeout=1)
|
||||
|
||||
assert exc_info.value.result["unsupported_count"] == 1
|
||||
assert exc_info.value.result["stage_results"][0]["result"]["supported"] is False
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_signal_duplex_turn_message_carries_next_fence(mocker: MockerFixture):
|
||||
request_q = queue.Queue()
|
||||
rpc_q = queue.Queue()
|
||||
cancelled_fence = DuplexFence("sid-cancel", epoch=4, turn_id=2)
|
||||
next_fence = DuplexFence("sid-cancel", epoch=5, turn_id=2)
|
||||
control_result = DuplexControlResultMessage(
|
||||
control_id="ctrl-signal",
|
||||
fence=cancelled_fence,
|
||||
operation="signal",
|
||||
session_id="sid-cancel",
|
||||
ok=True,
|
||||
stage_results=[],
|
||||
)
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
engine._duplex_control_client = None
|
||||
mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-signal"))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
pending = executor.submit(
|
||||
engine.signal_duplex_turn,
|
||||
"sid-cancel",
|
||||
event="input.cancel",
|
||||
fence=cancelled_fence,
|
||||
next_fence=next_fence,
|
||||
session_config={"temperature": 0.0},
|
||||
timeout=1,
|
||||
)
|
||||
msg = request_q.get(timeout=1)
|
||||
assert isinstance(msg, SignalDuplexTurnMessage)
|
||||
assert msg.fence == cancelled_fence
|
||||
assert msg.next_fence == next_fence
|
||||
assert msg.session_config == {"temperature": 0.0}
|
||||
rpc_q.put(control_result)
|
||||
assert pending.result(timeout=1)["ok"] is True
|
||||
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_async_append_does_not_stop_executor_callable(mocker: MockerFixture):
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
completed = threading.Event()
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
|
||||
def blocking_append(*args, **kwargs):
|
||||
del args, kwargs
|
||||
entered.set()
|
||||
release.wait(timeout=1)
|
||||
completed.set()
|
||||
return {}
|
||||
|
||||
mocker.patch.object(engine, "append_duplex_input", side_effect=blocking_append)
|
||||
task = asyncio.create_task(
|
||||
engine.append_duplex_input_async(
|
||||
"sid-executor-cancel",
|
||||
mode="append_audio_chunk",
|
||||
payload=b"audio",
|
||||
fence=DuplexFence("sid-executor-cancel"),
|
||||
)
|
||||
)
|
||||
assert await asyncio.to_thread(entered.wait, 1)
|
||||
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
release.set()
|
||||
|
||||
assert await asyncio.to_thread(completed.wait, 1)
|
||||
|
||||
|
||||
def test_duplex_controls_route_out_of_order_without_global_lock():
|
||||
request_q = queue.Queue()
|
||||
rpc_q = queue.Queue()
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
engine._duplex_control_client = None
|
||||
first_fence = DuplexFence("sid", epoch=1)
|
||||
second_fence = DuplexFence("sid", epoch=2)
|
||||
first = AppendDuplexInputMessage(
|
||||
control_id="first",
|
||||
fence=first_fence,
|
||||
session_id="sid",
|
||||
expected_epoch=1,
|
||||
mode="audio",
|
||||
payload=b"first",
|
||||
)
|
||||
second = AppendDuplexInputMessage(
|
||||
control_id="second",
|
||||
fence=second_fence,
|
||||
session_id="sid",
|
||||
expected_epoch=2,
|
||||
mode="audio",
|
||||
payload=b"second",
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
client = engine._get_duplex_control_client()
|
||||
first_result = executor.submit(client.execute, first, timeout=1)
|
||||
second_result = executor.submit(client.execute, second, timeout=1)
|
||||
submitted = {request_q.get(timeout=1).control_id, request_q.get(timeout=1).control_id}
|
||||
assert submitted == {"first", "second"}
|
||||
rpc_q.put(
|
||||
DuplexControlResultMessage(
|
||||
control_id="second",
|
||||
fence=second_fence,
|
||||
operation="append",
|
||||
session_id="sid",
|
||||
ok=True,
|
||||
stage_results=[],
|
||||
)
|
||||
)
|
||||
rpc_q.put(
|
||||
DuplexControlResultMessage(
|
||||
control_id="first",
|
||||
fence=first_fence,
|
||||
operation="append",
|
||||
session_id="sid",
|
||||
ok=True,
|
||||
stage_results=[],
|
||||
)
|
||||
)
|
||||
|
||||
assert first_result.result(timeout=1)["fence"] == first_fence
|
||||
assert second_result.result(timeout=1)["fence"] == second_fence
|
||||
|
||||
assert not hasattr(engine, "_rpc_lock")
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_terminal_rpc_router_rejects_without_enqueuing_new_control():
|
||||
request_q = queue.Queue()
|
||||
rpc_q = queue.Queue()
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
engine._duplex_control_client = None
|
||||
pending_message = AppendDuplexInputMessage(
|
||||
control_id="pending",
|
||||
fence=DuplexFence("sid"),
|
||||
session_id="sid",
|
||||
mode="audio",
|
||||
payload=b"audio",
|
||||
)
|
||||
message = AppendDuplexInputMessage(
|
||||
control_id="after-fatal",
|
||||
fence=DuplexFence("sid"),
|
||||
session_id="sid",
|
||||
mode="audio",
|
||||
payload=b"audio",
|
||||
)
|
||||
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
pending = executor.submit(engine._get_duplex_control_client().execute, pending_message, timeout=1)
|
||||
assert request_q.get(timeout=1).control_id == "pending"
|
||||
rpc_q.put(ErrorMessage(error="orchestrator failed", fatal=True))
|
||||
with pytest.raises(RuntimeError, match="orchestrator failed"):
|
||||
pending.result(timeout=1)
|
||||
|
||||
with pytest.raises(RuntimeError, match="orchestrator failed"):
|
||||
engine._get_duplex_control_client().execute(message, timeout=1)
|
||||
|
||||
assert request_q.empty()
|
||||
finally:
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_collective_and_duplex_results_share_router_without_swapping(mocker: MockerFixture):
|
||||
request_q = queue.Queue()
|
||||
rpc_q = queue.Queue()
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
engine._duplex_control_client = None
|
||||
fence = DuplexFence("sid")
|
||||
duplex = AppendDuplexInputMessage(
|
||||
control_id="duplex-control",
|
||||
fence=fence,
|
||||
session_id="sid",
|
||||
mode="audio",
|
||||
payload=b"audio",
|
||||
)
|
||||
mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="collective-rpc"))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
duplex_result = executor.submit(engine._get_duplex_control_client().execute, duplex, timeout=1)
|
||||
collective_result = executor.submit(engine.collective_rpc, "health", timeout=1)
|
||||
submitted_types = {request_q.get(timeout=1).type, request_q.get(timeout=1).type}
|
||||
assert submitted_types == {"append_duplex_input", "collective_rpc"}
|
||||
rpc_q.put(
|
||||
CollectiveRPCResultMessage(
|
||||
rpc_id="collective-rpc",
|
||||
method="health",
|
||||
stage_ids=[0],
|
||||
results=["healthy"],
|
||||
)
|
||||
)
|
||||
rpc_q.put(
|
||||
DuplexControlResultMessage(
|
||||
control_id="duplex-control",
|
||||
fence=fence,
|
||||
operation="append",
|
||||
session_id="sid",
|
||||
ok=True,
|
||||
stage_results=[],
|
||||
)
|
||||
)
|
||||
|
||||
assert collective_result.result(timeout=1) == ["healthy"]
|
||||
assert duplex_result.result(timeout=1)["fence"] == fence
|
||||
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixture):
|
||||
class SignallingQueue(queue.Queue):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(maxsize=1)
|
||||
self.put_attempted = threading.Event()
|
||||
|
||||
def put(self, item, block=True, timeout=None):
|
||||
self.put_attempted.set()
|
||||
return super().put(item, block=block, timeout=timeout)
|
||||
|
||||
request_q = SignallingQueue()
|
||||
request_q.put("queue-is-full")
|
||||
request_q.put_attempted.clear()
|
||||
rpc_q = queue.Queue()
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
engine._duplex_control_client = None
|
||||
mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-rpc"))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
pending = executor.submit(engine.collective_rpc, "health", timeout=1)
|
||||
assert request_q.put_attempted.wait(timeout=1)
|
||||
assert not pending.done()
|
||||
|
||||
assert request_q.get(timeout=1) == "queue-is-full"
|
||||
assert request_q.get(timeout=1).rpc_id == "blocked-rpc"
|
||||
rpc_q.put(
|
||||
CollectiveRPCResultMessage(
|
||||
rpc_id="blocked-rpc",
|
||||
method="health",
|
||||
stage_ids=[0],
|
||||
results=["healthy"],
|
||||
)
|
||||
)
|
||||
assert pending.result(timeout=1) == ["healthy"]
|
||||
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_open_duplex_session_raises_on_stage_control_error(mocker: MockerFixture):
|
||||
request_q = queue.Queue()
|
||||
rpc_q = queue.Queue()
|
||||
control_result = DuplexControlResultMessage(
|
||||
control_id="ctrl-error",
|
||||
fence=DuplexFence("sid"),
|
||||
operation="open",
|
||||
session_id="sid",
|
||||
ok=False,
|
||||
stage_results=[{"stage_id": 0, "replica_id": 0, "result": {"supported": False, "error": "boom"}}],
|
||||
unsupported_count=1,
|
||||
error_count=1,
|
||||
)
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
engine._duplex_control_client = None
|
||||
mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-error"))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
pending = executor.submit(engine.open_duplex_session, "sid", fence=DuplexFence("sid"), timeout=1)
|
||||
request_q.get(timeout=1)
|
||||
rpc_q.put(control_result)
|
||||
with pytest.raises(RuntimeError, match="duplex open failed"):
|
||||
pending.result(timeout=1)
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_shutdown_does_not_race_runtime_cleanup_with_live_orchestrator(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine._weak_finalizer = None
|
||||
engine.request_queue = mocker.MagicMock()
|
||||
engine.request_queue.sync_q.put.side_effect = queue.Full
|
||||
engine.output_queue = mocker.MagicMock()
|
||||
engine.rpc_output_queue = mocker.MagicMock()
|
||||
engine._correlated_rpc_client = mocker.MagicMock()
|
||||
engine.orchestrator_thread = mocker.MagicMock()
|
||||
engine.orchestrator_thread.is_alive.return_value = True
|
||||
engine._runtime = mocker.MagicMock()
|
||||
|
||||
engine.shutdown()
|
||||
|
||||
engine.request_queue.sync_q.put.assert_called_once()
|
||||
assert engine.request_queue.sync_q.put.call_args.kwargs["timeout"] > 0
|
||||
assert engine.orchestrator_thread.join.call_args_list[0].kwargs["timeout"] > 0
|
||||
engine._correlated_rpc_client.close.assert_called_once_with()
|
||||
engine.request_queue.close.assert_called_once_with()
|
||||
engine.output_queue.close.assert_called_once_with()
|
||||
engine.rpc_output_queue.close.assert_called_once_with()
|
||||
engine._runtime.shutdown.assert_not_called()
|
||||
|
||||
|
||||
def test_shutdown_releases_runtime_after_orchestrator_stops(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine._weak_finalizer = None
|
||||
engine.request_queue = mocker.MagicMock()
|
||||
engine.output_queue = mocker.MagicMock()
|
||||
engine.rpc_output_queue = mocker.MagicMock()
|
||||
engine._correlated_rpc_client = mocker.MagicMock()
|
||||
engine.orchestrator_thread = mocker.MagicMock()
|
||||
engine.orchestrator_thread.is_alive.side_effect = [True, False]
|
||||
engine._runtime = mocker.MagicMock()
|
||||
|
||||
engine.shutdown()
|
||||
|
||||
engine.orchestrator_thread.join.assert_called_once()
|
||||
engine._runtime.shutdown.assert_called_once_with()
|
||||
|
||||
|
||||
def test_shutdown_defers_runtime_release_until_live_orchestrator_stops(mocker: MockerFixture):
|
||||
orchestrator_stopped = threading.Event()
|
||||
runtime_released = threading.Event()
|
||||
|
||||
class ControllableOrchestratorThread:
|
||||
def __init__(self) -> None:
|
||||
self.join_timeouts: list[float | None] = []
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
return not orchestrator_stopped.is_set()
|
||||
|
||||
def join(self, timeout: float | None = None) -> None:
|
||||
self.join_timeouts.append(timeout)
|
||||
if timeout is None:
|
||||
orchestrator_stopped.wait()
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine._weak_finalizer = None
|
||||
engine.request_queue = mocker.MagicMock()
|
||||
engine.output_queue = mocker.MagicMock()
|
||||
engine.rpc_output_queue = mocker.MagicMock()
|
||||
engine._correlated_rpc_client = mocker.MagicMock()
|
||||
engine.orchestrator_thread = ControllableOrchestratorThread()
|
||||
engine._runtime = mocker.MagicMock()
|
||||
engine._runtime.shutdown.side_effect = runtime_released.set
|
||||
|
||||
try:
|
||||
engine.shutdown()
|
||||
|
||||
engine._runtime.shutdown.assert_not_called()
|
||||
orchestrator_stopped.set()
|
||||
assert runtime_released.wait(timeout=1)
|
||||
engine._runtime.shutdown.assert_called_once_with()
|
||||
assert engine.orchestrator_thread.join_timeouts[0] is not None
|
||||
assert engine.orchestrator_thread.join_timeouts[-1] is None
|
||||
finally:
|
||||
orchestrator_stopped.set()
|
||||
|
||||
|
||||
def test_weak_shutdown_is_bounded_when_request_queue_is_full(mocker: MockerFixture):
|
||||
request_queue = mocker.MagicMock()
|
||||
request_queue.sync_q.put.side_effect = queue.Full
|
||||
output_queue = mocker.MagicMock()
|
||||
rpc_output_queue = mocker.MagicMock()
|
||||
router = mocker.MagicMock()
|
||||
orchestrator_thread = mocker.MagicMock()
|
||||
orchestrator_thread.is_alive.return_value = True
|
||||
|
||||
weak_shutdown_async_omni_engine(
|
||||
orchestrator_thread,
|
||||
request_queue,
|
||||
output_queue,
|
||||
rpc_output_queue,
|
||||
router,
|
||||
)
|
||||
|
||||
request_queue.sync_q.put.assert_called_once()
|
||||
assert request_queue.sync_q.put.call_args.kwargs["timeout"] > 0
|
||||
orchestrator_thread.join.assert_called_once()
|
||||
assert orchestrator_thread.join.call_args.kwargs["timeout"] > 0
|
||||
request_queue.close.assert_called_once_with()
|
||||
output_queue.close.assert_called_once_with()
|
||||
rpc_output_queue.close.assert_called_once_with()
|
||||
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import queue
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.engine.messages import (
|
||||
CollectiveRPCRequestMessage,
|
||||
CollectiveRPCResultMessage,
|
||||
ErrorMessage,
|
||||
)
|
||||
from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _request(rpc_id: str) -> CollectiveRPCRequestMessage:
|
||||
return CollectiveRPCRequestMessage(
|
||||
rpc_id=rpc_id,
|
||||
method="health",
|
||||
timeout=1,
|
||||
args=(),
|
||||
kwargs={},
|
||||
stage_ids=None,
|
||||
)
|
||||
|
||||
|
||||
def test_correlated_rpc_client_unregisters_timeout_before_late_result() -> None:
|
||||
request_queue: queue.Queue = queue.Queue()
|
||||
result_queue: queue.Queue = queue.Queue()
|
||||
client = CorrelatedRpcClient(request_queue, result_queue)
|
||||
|
||||
try:
|
||||
with pytest.raises(TimeoutError, match="first timed out"):
|
||||
client.execute(
|
||||
("collective", "first"),
|
||||
_request("first"),
|
||||
timeout=0.01,
|
||||
timeout_message="first timed out",
|
||||
)
|
||||
assert request_queue.get_nowait().rpc_id == "first"
|
||||
|
||||
result_queue.put(
|
||||
CollectiveRPCResultMessage(
|
||||
rpc_id="first",
|
||||
method="health",
|
||||
stage_ids=[0],
|
||||
results=["late"],
|
||||
)
|
||||
)
|
||||
result_queue.put(
|
||||
CollectiveRPCResultMessage(
|
||||
rpc_id="second",
|
||||
method="health",
|
||||
stage_ids=[0],
|
||||
results=["current"],
|
||||
)
|
||||
)
|
||||
current = client.execute(
|
||||
("collective", "second"),
|
||||
_request("second"),
|
||||
timeout=1,
|
||||
timeout_message="second timed out",
|
||||
)
|
||||
|
||||
assert isinstance(current, CollectiveRPCResultMessage)
|
||||
assert current.results == ["current"]
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def test_correlated_rpc_client_rejects_after_fatal_without_enqueuing() -> None:
|
||||
request_queue: queue.Queue = queue.Queue()
|
||||
result_queue: queue.Queue = queue.Queue()
|
||||
client = CorrelatedRpcClient(request_queue, result_queue)
|
||||
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
pending = executor.submit(
|
||||
client.execute,
|
||||
("collective", "pending"),
|
||||
_request("pending"),
|
||||
timeout=1,
|
||||
timeout_message="unexpected timeout",
|
||||
)
|
||||
assert request_queue.get(timeout=1).rpc_id == "pending"
|
||||
result_queue.put(ErrorMessage(error="orchestrator failed", fatal=True))
|
||||
with pytest.raises(RuntimeError, match="orchestrator failed"):
|
||||
pending.result(timeout=1)
|
||||
|
||||
with pytest.raises(RuntimeError, match="orchestrator failed"):
|
||||
client.execute(
|
||||
("collective", "after-fatal"),
|
||||
_request("after-fatal"),
|
||||
timeout=1,
|
||||
timeout_message="unexpected timeout",
|
||||
)
|
||||
assert request_queue.empty()
|
||||
finally:
|
||||
client.close()
|
||||
@@ -0,0 +1,102 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def _assert_isolated_import_succeeds(script: str) -> None:
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
cwd=REPO_ROOT,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_stable_engine_imports_do_not_load_experimental_duplex() -> None:
|
||||
_assert_isolated_import_succeeds("""
|
||||
import sys
|
||||
|
||||
import vllm_omni.engine.async_omni_engine
|
||||
import vllm_omni.engine.orchestrator
|
||||
import vllm_omni.entrypoints.async_omni
|
||||
|
||||
loaded = sorted(
|
||||
name
|
||||
for name in sys.modules
|
||||
if name == "vllm_omni.experimental.fullduplex"
|
||||
or name.startswith("vllm_omni.experimental.fullduplex.")
|
||||
)
|
||||
if loaded:
|
||||
raise SystemExit("stable imports loaded experimental duplex: " + ", ".join(loaded))
|
||||
""")
|
||||
|
||||
|
||||
def test_stable_engine_does_not_expose_duplex_contract_modules() -> None:
|
||||
_assert_isolated_import_succeeds("""
|
||||
import importlib.util
|
||||
|
||||
from vllm_omni.engine import messages
|
||||
|
||||
legacy_modules = (
|
||||
"vllm_omni.engine.duplex_contracts",
|
||||
"vllm_omni.engine.duplex_lease",
|
||||
"vllm_omni.engine.resumable",
|
||||
)
|
||||
present = [name for name in legacy_modules if importlib.util.find_spec(name) is not None]
|
||||
if present:
|
||||
raise SystemExit("stable duplex modules still exist: " + ", ".join(present))
|
||||
|
||||
duplex_exports = sorted(name for name in vars(messages) if name.startswith("Duplex"))
|
||||
if duplex_exports:
|
||||
raise SystemExit("stable messages still expose duplex contracts: " + ", ".join(duplex_exports))
|
||||
""")
|
||||
|
||||
|
||||
def test_stable_outputs_do_not_declare_duplex_decision_field() -> None:
|
||||
_assert_isolated_import_succeeds("""
|
||||
import dataclasses
|
||||
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
fields = {field.name for field in dataclasses.fields(OmniRequestOutput)}
|
||||
if "duplex_output_decision" in fields:
|
||||
raise SystemExit("stable output declares duplex_output_decision")
|
||||
""")
|
||||
|
||||
|
||||
def test_stable_model_executor_does_not_expose_duplex_helper_module() -> None:
|
||||
_assert_isolated_import_succeeds("""
|
||||
import importlib.util
|
||||
|
||||
if importlib.util.find_spec("vllm_omni.model_executor.duplex") is not None:
|
||||
raise SystemExit("stable model_executor still exposes duplex helper module")
|
||||
""")
|
||||
|
||||
|
||||
def test_runtime_package_does_not_bundle_the_browser_demo() -> None:
|
||||
assert not (REPO_ROOT / "vllm_omni" / "experimental" / "fullduplex" / "web").exists()
|
||||
|
||||
|
||||
def test_experimental_engine_uses_canonical_contract_module_names() -> None:
|
||||
engine_dir = REPO_ROOT / "vllm_omni" / "experimental" / "fullduplex" / "engine"
|
||||
core_dir = REPO_ROOT / "vllm_omni" / "experimental" / "fullduplex" / "core"
|
||||
|
||||
assert (engine_dir / "contracts.py").is_file()
|
||||
assert (engine_dir / "lease.py").is_file()
|
||||
assert (engine_dir / "messages.py").is_file()
|
||||
assert not (engine_dir / "duplex_lease.py").exists()
|
||||
assert not (engine_dir / "duplex_types.py").exists()
|
||||
assert not (core_dir / "identity.py").exists()
|
||||
@@ -11,6 +11,7 @@ import pytest
|
||||
|
||||
from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus
|
||||
from vllm_omni.engine.membership_controller import MembershipController
|
||||
from vllm_omni.engine.stage_pool import StagePool
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -45,10 +46,10 @@ class FakePool:
|
||||
def attach_load_balancer(self, lb):
|
||||
self.lb = lb
|
||||
|
||||
def add_client(self, input_addr, client):
|
||||
self.added.append((input_addr, client))
|
||||
def add_client(self, input_addr, client, *, replica_id=None):
|
||||
self.added.append((input_addr, client, replica_id))
|
||||
self.clients.append(client)
|
||||
return len(self.clients) - 1
|
||||
return len(self.clients) - 1 if replica_id is None else replica_id
|
||||
|
||||
def invalidate_addr(self, input_addr):
|
||||
self.invalidated.append(input_addr)
|
||||
@@ -155,4 +156,35 @@ async def test_do_register_offloads_remote_factory(monkeypatch):
|
||||
|
||||
assert factory_thread_ids
|
||||
assert factory_thread_ids[0] != threading.get_ident()
|
||||
assert pool.added[0][0] == "tcp://stage-0-replica-1"
|
||||
assert pool.added[0] == ("tcp://stage-0-replica-1", pool.clients[0], 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unregister_then_register_restores_coordinator_replica_slot(monkeypatch):
|
||||
class Client:
|
||||
def __init__(self, input_addr: str) -> None:
|
||||
self.client_addresses = {"input_address": input_addr}
|
||||
self.shutdown_calls = 0
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self.shutdown_calls += 1
|
||||
|
||||
old_clients = [Client("tcp://old-0"), Client("tcp://old-1")]
|
||||
replacement = Client("tcp://new-1")
|
||||
pool = StagePool(0, old_clients)
|
||||
|
||||
def factory(stage_id: int, replica_id: int):
|
||||
assert (stage_id, replica_id) == (0, 1)
|
||||
return replacement
|
||||
|
||||
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=factory)
|
||||
await controller.handle_unregister(0, "tcp://old-0")
|
||||
await controller.handle_unregister(0, "tcp://old-1")
|
||||
|
||||
await controller.handle_register(0, 1)
|
||||
await controller.drain_tasks(timeout=1)
|
||||
|
||||
assert pool.clients == [None, replacement]
|
||||
assert pool.get_replica_id_by_addr("tcp://new-1") == 1
|
||||
assert pool.get_replica_id_by_addr("tcp://old-1") is None
|
||||
assert pool.available_replica_ids() == [1]
|
||||
@@ -0,0 +1,46 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.outputs.mm_outputs import MultimodalPayload
|
||||
from vllm_omni.outputs.multimodal_accumulation import (
|
||||
drain_delta_payload,
|
||||
is_non_final_delta_audio_chunk,
|
||||
replace_snapshot_keys,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state():
|
||||
accumulated = MultimodalPayload.from_dict(
|
||||
{
|
||||
"audio": torch.tensor([1.0]),
|
||||
"meta.tts_is_last_chunk": torch.tensor([0]),
|
||||
"meta.turn_end": torch.tensor([0]),
|
||||
"meta.stable_request_value": "keep",
|
||||
}
|
||||
)
|
||||
incoming = MultimodalPayload.from_dict(
|
||||
{
|
||||
"audio": torch.tensor([2.0]),
|
||||
"meta.tts_is_last_chunk": torch.tensor([1]),
|
||||
"meta.turn_end": torch.tensor([1]),
|
||||
}
|
||||
)
|
||||
assert accumulated is not None
|
||||
assert incoming is not None
|
||||
|
||||
replace_snapshot_keys(accumulated, incoming)
|
||||
merged = accumulated.merged_with(incoming)
|
||||
|
||||
assert not is_non_final_delta_audio_chunk(merged, "audio")
|
||||
|
||||
drain_delta_payload(merged)
|
||||
|
||||
assert "audio" not in merged
|
||||
assert "meta.tts_is_last_chunk" not in merged
|
||||
assert "meta.turn_end" not in merged
|
||||
assert merged.metadata["meta.stable_request_value"] == "keep"
|
||||
File diff suppressed because it is too large.
Load diff
@@ -38,6 +38,7 @@ class FakeStageClient:
|
||||
self.custom_process_input_func = None
|
||||
self.next_inputs = list(next_inputs or [])
|
||||
self.add_request_calls: list[tuple[Any, ...]] = []
|
||||
self.decoded_source_tokens: str | None = None
|
||||
self._engine_core_outputs = queue.Queue()
|
||||
|
||||
async def add_request_async(self, *args, **_kwargs) -> None:
|
||||
@@ -50,6 +51,9 @@ class FakeStageClient:
|
||||
return SimpleNamespace(outputs=[])
|
||||
|
||||
def process_engine_inputs(self, _source_outputs, prompt=None, streaming_context=None):
|
||||
decoder = getattr(streaming_context, "source_token_decoder", None)
|
||||
if callable(decoder):
|
||||
self.decoded_source_tokens = decoder([11, 12], skip_special_tokens=True)
|
||||
return list(self.next_inputs)
|
||||
|
||||
async def abort_requests_async(self, _request_ids: list[str]) -> None:
|
||||
@@ -66,6 +70,9 @@ class FakeStageClient:
|
||||
|
||||
|
||||
class FakeOutputProcessor:
|
||||
def __init__(self, tokenizer=None) -> None:
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def add_request(self, *args, **kwargs) -> None:
|
||||
return None
|
||||
|
||||
@@ -108,6 +115,11 @@ def _request_output(request_id: str) -> RequestOutput:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_forward_text_prompt_uses_target_stage_input_processor() -> None:
|
||||
class SourceTokenizer:
|
||||
def decode(self, token_ids, *, skip_special_tokens):
|
||||
assert skip_special_tokens is True
|
||||
return ":".join(str(token_id) for token_id in token_ids)
|
||||
|
||||
stage0 = FakeStageClient(final_output=True)
|
||||
stage1 = FakeStageClient(
|
||||
final_output=True,
|
||||
@@ -117,7 +129,7 @@ async def test_forward_text_prompt_uses_target_stage_input_processor() -> None:
|
||||
StagePool(
|
||||
0,
|
||||
[stage0],
|
||||
output_processor=FakeOutputProcessor(),
|
||||
output_processor=FakeOutputProcessor(tokenizer=SourceTokenizer()),
|
||||
stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
|
||||
),
|
||||
StagePool(
|
||||
@@ -150,6 +162,8 @@ async def test_forward_text_prompt_uses_target_stage_input_processor() -> None:
|
||||
|
||||
assert input_processor.calls
|
||||
assert input_processor.calls[0]["prompt"] == {"prompt": "hello", "multi_modal_data": {"video": ["frame"]}}
|
||||
assert stage1.decoded_source_tokens == "11:12"
|
||||
assert req_state.streaming.source_token_decoder is None
|
||||
assert stage1.add_request_calls
|
||||
submitted_request = stage1.add_request_calls[0][0]
|
||||
assert submitted_request.prompt_token_ids == [101, 102]
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
@@ -10,9 +11,11 @@ import torch
|
||||
from vllm.outputs import PoolingRequestOutput
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.v1.engine import FinishReason
|
||||
from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
|
||||
|
||||
from vllm_omni.outputs import output_processor
|
||||
from vllm_omni.outputs.output_modality import OutputModalityNames
|
||||
from vllm_omni.outputs.output_processor import OmniRequestState
|
||||
from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -62,6 +65,65 @@ def test_init_empty_dict():
|
||||
assert _make_state(RequestOutputKind.DELTA).mm_accumulated == {}
|
||||
|
||||
|
||||
def test_streaming_update_resets_native_text_metrics_for_next_segment():
|
||||
state = _make_state(RequestOutputKind.DELTA)
|
||||
state.native_text_stats.num_generation_tokens = 12
|
||||
state.native_text_stats.first_token_ts = 10.0
|
||||
state.native_text_stats.last_token_ts = 10.2
|
||||
|
||||
state.apply_streaming_update(
|
||||
SimpleNamespace(
|
||||
final=False,
|
||||
prompt=None,
|
||||
prompt_token_ids=[],
|
||||
arrival_time=20.0,
|
||||
)
|
||||
)
|
||||
|
||||
assert state.native_text_stats.arrival_time == 20.0
|
||||
assert state.native_text_stats.num_generation_tokens == 0
|
||||
assert state.native_text_stats.first_token_ts == 0.0
|
||||
assert state.native_text_stats.last_token_ts == 0.0
|
||||
|
||||
|
||||
def test_native_text_tpot_snapshot_matches_vllm_finished_metric_definition():
|
||||
calculate = getattr(output_processor, "_mean_time_per_output_token_ms", None)
|
||||
assert calculate is not None
|
||||
stats = SimpleNamespace(
|
||||
num_generation_tokens=4,
|
||||
first_token_ts=10.0,
|
||||
last_token_ts=10.045,
|
||||
)
|
||||
|
||||
assert calculate(stats) == pytest.approx(15.0)
|
||||
|
||||
|
||||
def test_native_text_metrics_include_segment_generation_token_count(monkeypatch):
|
||||
monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
|
||||
processor = object.__new__(MultimodalOutputProcessor)
|
||||
processor._native_text_metrics_by_request = {}
|
||||
processor.lora_states = {}
|
||||
state = _make_state(RequestOutputKind.DELTA)
|
||||
state.is_prefilling = False
|
||||
iteration_stats = MagicMock()
|
||||
|
||||
def update_native_stats(_output, _timestamp, _was_prefilling, native_stats, *_args):
|
||||
native_stats.num_generation_tokens = 27
|
||||
native_stats.first_token_ts = 10.0
|
||||
native_stats.last_token_ts = 10.26
|
||||
|
||||
iteration_stats.update_from_output.side_effect = update_native_stats
|
||||
|
||||
processor._update_stats_from_output(
|
||||
state,
|
||||
SimpleNamespace(),
|
||||
10.26,
|
||||
iteration_stats,
|
||||
)
|
||||
|
||||
assert processor.pop_native_text_metrics("r")["num_generation_tokens"] == 27
|
||||
|
||||
|
||||
def test_delta_drains_output_modality_per_step():
|
||||
"""DELTA drains the mm_type key (output modality) but preserves hidden-state keys."""
|
||||
s = _make_state(RequestOutputKind.DELTA)
|
||||
@@ -172,6 +234,145 @@ def test_finish_consolidation_drains_mm_delta():
|
||||
assert AUDIO not in s.mm_accumulated # drained
|
||||
|
||||
|
||||
def test_delta_drains_audio_chunk_metadata_per_step():
|
||||
s = _make_state(RequestOutputKind.DELTA)
|
||||
segment_text = "你好"
|
||||
segment_utf8 = torch.tensor(list(segment_text.encode("utf-8")), dtype=torch.uint8)
|
||||
|
||||
s.add_multimodal_tensor(
|
||||
{
|
||||
"model_outputs": torch.ones(5),
|
||||
"meta.duplex_epoch": torch.tensor([7], dtype=torch.int32),
|
||||
"meta.duplex_turn_id": torch.tensor([1], dtype=torch.int32),
|
||||
"meta.llm_output_text_utf8": segment_utf8,
|
||||
"meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
|
||||
"meta.tts_is_last_chunk": torch.tensor([0], dtype=torch.int32),
|
||||
},
|
||||
mm_type=AUDIO,
|
||||
)
|
||||
result1 = s.make_request_output([1], None, FinishReason.STOP, None)
|
||||
assert result1 is not None and not isinstance(result1, PoolingRequestOutput)
|
||||
meta1 = result1.outputs[0].multimodal_output["meta"]
|
||||
assert bytes(meta1["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text
|
||||
|
||||
s.add_multimodal_tensor(
|
||||
{
|
||||
"model_outputs": torch.ones(3),
|
||||
"meta.duplex_epoch": torch.tensor([8], dtype=torch.int32),
|
||||
"meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32),
|
||||
"meta.llm_output_text_utf8": segment_utf8,
|
||||
"meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
|
||||
"meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32),
|
||||
},
|
||||
mm_type=AUDIO,
|
||||
)
|
||||
result2 = s.make_request_output([2], None, FinishReason.STOP, None)
|
||||
|
||||
assert result2 is not None and not isinstance(result2, PoolingRequestOutput)
|
||||
meta2 = result2.outputs[0].multimodal_output["meta"]
|
||||
assert bytes(meta2["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text
|
||||
assert meta2["audio_text_total_chars"].tolist() == [len(segment_text)]
|
||||
assert meta2["duplex_epoch"].tolist() == [8]
|
||||
assert meta2["duplex_turn_id"].tolist() == [2]
|
||||
|
||||
|
||||
def test_cumulative_audio_replaces_chunk_metadata_per_step():
|
||||
s = _make_state(RequestOutputKind.CUMULATIVE)
|
||||
segment_text = "你好"
|
||||
segment_utf8 = torch.tensor(list(segment_text.encode("utf-8")), dtype=torch.uint8)
|
||||
|
||||
s.add_multimodal_tensor(
|
||||
{
|
||||
"model_outputs": torch.ones(5),
|
||||
"meta.duplex_epoch": torch.tensor([7], dtype=torch.int32),
|
||||
"meta.duplex_turn_id": torch.tensor([1], dtype=torch.int32),
|
||||
"meta.llm_output_text_utf8": segment_utf8,
|
||||
"meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
|
||||
"meta.tts_is_last_chunk": torch.tensor([0], dtype=torch.int32),
|
||||
},
|
||||
mm_type=AUDIO,
|
||||
)
|
||||
result1 = s.make_request_output([1], None, None, None)
|
||||
assert result1 is not None and not isinstance(result1, PoolingRequestOutput)
|
||||
assert result1.outputs[0].multimodal_output["audio"].shape[0] == 5
|
||||
|
||||
s.add_multimodal_tensor(
|
||||
{
|
||||
"model_outputs": torch.ones(3),
|
||||
"meta.duplex_epoch": torch.tensor([8], dtype=torch.int32),
|
||||
"meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32),
|
||||
"meta.llm_output_text_utf8": segment_utf8,
|
||||
"meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
|
||||
"meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32),
|
||||
},
|
||||
mm_type=AUDIO,
|
||||
)
|
||||
result2 = s.make_request_output([2], None, None, None)
|
||||
|
||||
assert result2 is not None and not isinstance(result2, PoolingRequestOutput)
|
||||
output = result2.outputs[0].multimodal_output
|
||||
assert output["audio"].shape[0] == 8
|
||||
meta = output["meta"]
|
||||
assert bytes(meta["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text
|
||||
assert meta["audio_text_total_chars"].tolist() == [len(segment_text)]
|
||||
assert meta["duplex_epoch"].tolist() == [8]
|
||||
assert meta["duplex_turn_id"].tolist() == [2]
|
||||
assert meta["tts_is_last_chunk"].tolist() == [1]
|
||||
|
||||
|
||||
def test_delta_audio_non_final_tts_chunk_overrides_spurious_finish():
|
||||
s = _make_state(RequestOutputKind.DELTA)
|
||||
s.add_multimodal_tensor(
|
||||
{
|
||||
"model_outputs": torch.ones(5),
|
||||
"meta.tts_is_last_chunk": torch.tensor([0]),
|
||||
},
|
||||
mm_type=AUDIO,
|
||||
)
|
||||
|
||||
result = s.make_request_output([1], None, FinishReason.STOP, None)
|
||||
|
||||
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
||||
assert result.finished is False
|
||||
assert result.outputs[0].finish_reason is None
|
||||
assert result.outputs[0].multimodal_output["meta"]["tts_is_last_chunk"].tolist() == [0]
|
||||
assert torch.equal(result.outputs[0].multimodal_output[AUDIO], torch.ones(5))
|
||||
|
||||
|
||||
def test_delta_audio_final_tts_chunk_keeps_finish():
|
||||
s = _make_state(RequestOutputKind.DELTA)
|
||||
s.add_multimodal_tensor(
|
||||
{
|
||||
"model_outputs": torch.ones(5),
|
||||
"meta.tts_is_last_chunk": torch.tensor([1]),
|
||||
},
|
||||
mm_type=AUDIO,
|
||||
)
|
||||
|
||||
result = s.make_request_output([151645], None, FinishReason.STOP, None)
|
||||
|
||||
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
||||
assert result.finished is True
|
||||
assert result.outputs[0].finish_reason == "stop"
|
||||
|
||||
|
||||
def test_delta_audio_non_final_tts_chunk_accepts_nested_meta():
|
||||
s = _make_state(RequestOutputKind.DELTA)
|
||||
s.add_multimodal_tensor(
|
||||
{
|
||||
AUDIO: torch.ones(5),
|
||||
"meta": {"tts_is_last_chunk": torch.tensor([0])},
|
||||
},
|
||||
mm_type=AUDIO,
|
||||
)
|
||||
|
||||
result = s.make_request_output([1], None, FinishReason.STOP, None)
|
||||
|
||||
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
||||
assert result.finished is False
|
||||
assert result.outputs[0].multimodal_output["meta"]["tts_is_last_chunk"].tolist() == [0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mm_type", [AUDIO, "hidden"])
|
||||
def test_final_only_consolidates_drainable_keys(mm_type):
|
||||
"""FINAL_ONLY never drains per-step, so modality keys and hidden state
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
import queue
|
||||
from typing import Literal
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.engine.messages import (
|
||||
CollectiveRPCResultMessage,
|
||||
EngineQueueMessage,
|
||||
ErrorMessage,
|
||||
)
|
||||
from vllm_omni.engine.rpc_result_router import RpcResultRouter
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class GenericCorrelatedResultMessage(EngineQueueMessage, kw_only=True):
|
||||
type: Literal["generic_correlated_result"] = "generic_correlated_result"
|
||||
namespace: str
|
||||
correlation_id: str
|
||||
|
||||
@property
|
||||
def rpc_correlation_key(self) -> tuple[str, str]:
|
||||
return (self.namespace, self.correlation_id)
|
||||
|
||||
|
||||
def _generic_result(correlation_id: str) -> GenericCorrelatedResultMessage:
|
||||
return GenericCorrelatedResultMessage(
|
||||
namespace="plugin",
|
||||
correlation_id=correlation_id,
|
||||
)
|
||||
|
||||
|
||||
def test_rpc_result_router_routes_out_of_order_results_by_correlation_id():
|
||||
source: queue.Queue = queue.Queue()
|
||||
router = RpcResultRouter(source)
|
||||
first = router.register(("plugin", "first"))
|
||||
second = router.register(("plugin", "second"))
|
||||
|
||||
source.put(_generic_result("second"))
|
||||
source.put(_generic_result("first"))
|
||||
|
||||
assert first.get(timeout=1).correlation_id == "first"
|
||||
assert second.get(timeout=1).correlation_id == "second"
|
||||
router.close()
|
||||
|
||||
|
||||
def test_rpc_result_router_drops_only_the_late_result_after_unregister():
|
||||
source: queue.Queue = queue.Queue()
|
||||
router = RpcResultRouter(source)
|
||||
expired = router.register(("plugin", "expired"))
|
||||
active = router.register(("collective", "active"))
|
||||
router.unregister(("plugin", "expired"), expired)
|
||||
|
||||
source.put(_generic_result("expired"))
|
||||
source.put(
|
||||
CollectiveRPCResultMessage(
|
||||
rpc_id="active",
|
||||
method="health",
|
||||
stage_ids=[0],
|
||||
results=["ok"],
|
||||
)
|
||||
)
|
||||
|
||||
assert active.get(timeout=1).rpc_id == "active"
|
||||
assert expired.empty()
|
||||
router.close()
|
||||
|
||||
|
||||
def test_rpc_result_router_broadcasts_fatal_errors_to_pending_waiters():
|
||||
source: queue.Queue = queue.Queue()
|
||||
router = RpcResultRouter(source)
|
||||
plugin = router.register(("plugin", "one"))
|
||||
collective = router.register(("collective", "two"))
|
||||
|
||||
source.put(ErrorMessage(error="orchestrator failed", fatal=True))
|
||||
|
||||
assert plugin.get(timeout=1).error == "orchestrator failed"
|
||||
assert collective.get(timeout=1).error == "orchestrator failed"
|
||||
with pytest.raises(RuntimeError, match="orchestrator failed"):
|
||||
router.register(("plugin", "after-failure"))
|
||||
router.close()
|
||||
|
||||
|
||||
def test_rpc_result_router_does_not_broadcast_uncorrelated_nonfatal_errors():
|
||||
source: queue.Queue = queue.Queue()
|
||||
router = RpcResultRouter(source)
|
||||
waiter = router.register(("plugin", "active"))
|
||||
|
||||
source.put(ErrorMessage(error="request failed", fatal=False, request_id="other"))
|
||||
source.put(_generic_result("active"))
|
||||
|
||||
assert waiter.get(timeout=1).correlation_id == "active"
|
||||
router.close()
|
||||
|
||||
|
||||
def test_rpc_result_router_close_unblocks_waiters_and_stops_consumer():
|
||||
source: queue.Queue = queue.Queue()
|
||||
router = RpcResultRouter(source)
|
||||
waiter = router.register(("plugin", "pending"))
|
||||
|
||||
router.close()
|
||||
|
||||
result = waiter.get(timeout=1)
|
||||
assert isinstance(result, ErrorMessage)
|
||||
assert result.fatal is True
|
||||
assert result.error == "RPC result router closed"
|
||||
assert not router._thread.is_alive()
|
||||
with pytest.raises(RuntimeError, match="router is closed"):
|
||||
router.register(("plugin", "after-close"))
|
||||
@@ -0,0 +1,474 @@
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.openai.protocol import (
|
||||
DuplexCapabilities,
|
||||
DuplexOverlapPolicy,
|
||||
DuplexSession,
|
||||
DuplexSessionConfig,
|
||||
DuplexSessionRegistry,
|
||||
DuplexTurnController,
|
||||
DuplexTurnEventType,
|
||||
DuplexTurnState,
|
||||
ResponseCreateOptions,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.openai.realtime_session import (
|
||||
NativeRealtimeSessionProtocol,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.openai.realtime_state import RealtimeStateOwner
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_realtime_state_owner_uses_explicit_bindings_not_dynamic_attribute_proxy():
|
||||
assert "__getattr__" not in RealtimeStateOwner.__dict__
|
||||
assert "__setattr__" not in RealtimeStateOwner.__dict__
|
||||
|
||||
|
||||
class _ProtocolWebSocket:
|
||||
def __init__(self, *events: dict[str, object]) -> None:
|
||||
self._events = asyncio.Queue()
|
||||
for event in events:
|
||||
self._events.put_nowait(json.dumps(event))
|
||||
|
||||
async def receive_text(self) -> str:
|
||||
return await self._events.get()
|
||||
|
||||
|
||||
def test_duplex_session_commits_text_and_audio_as_one_turn():
|
||||
registry = DuplexSessionRegistry()
|
||||
session = registry.create(DuplexSessionConfig(model="test-model"))
|
||||
|
||||
session.append_text("hello")
|
||||
session.append_audio("YWJj", fmt="wav", sample_rate_hz=16000)
|
||||
committed = session.commit_user_input()
|
||||
|
||||
assert committed is not None
|
||||
assert committed.turn_id == 0
|
||||
assert committed.input_commit_seq == 1
|
||||
assert committed.epoch == 0
|
||||
assert len(session.history) == 1
|
||||
content = session.history[0]["content"]
|
||||
assert isinstance(content, list)
|
||||
assert content[0] == {"type": "text", "text": "hello"}
|
||||
assert content[1]["type"] == "audio_url"
|
||||
assert content[1]["audio_url"]["url"] == "data:audio/wav;base64,YWJj"
|
||||
assert session.turn_state == DuplexTurnState.USER_COMMITTED
|
||||
|
||||
|
||||
def test_response_options_apply_to_one_response_without_mutating_session_defaults():
|
||||
session = DuplexSession(
|
||||
session_id="sid-response-options",
|
||||
config=DuplexSessionConfig(instructions="base", voice="base-voice", max_tokens=64),
|
||||
)
|
||||
session.reserve_response_options(
|
||||
ResponseCreateOptions(
|
||||
instructions="one response",
|
||||
voice="override-voice",
|
||||
max_tokens=8,
|
||||
)
|
||||
)
|
||||
|
||||
assert session.config.instructions == "base"
|
||||
assert session.config.voice == "base-voice"
|
||||
session.begin_response()
|
||||
assert session.config.instructions == "base"
|
||||
assert session.config.voice == "base-voice"
|
||||
assert session.response_config.instructions == "one response"
|
||||
assert session.response_config.voice == "override-voice"
|
||||
assert session.response_config.max_tokens == 8
|
||||
|
||||
session.end_response()
|
||||
assert session.config.instructions == "base"
|
||||
assert session.config.voice == "base-voice"
|
||||
assert session.config.max_tokens == 64
|
||||
|
||||
|
||||
def test_response_options_cannot_overwrite_an_unconsumed_reservation():
|
||||
session = DuplexSession(
|
||||
session_id="sid-response-options-pending",
|
||||
config=DuplexSessionConfig(instructions="base"),
|
||||
)
|
||||
session.reserve_response_options(ResponseCreateOptions(instructions="first"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="already reserved"):
|
||||
session.reserve_response_options(ResponseCreateOptions(instructions="second"))
|
||||
|
||||
session.begin_response()
|
||||
assert session.response_config.instructions == "first"
|
||||
|
||||
|
||||
def test_response_options_cannot_be_reserved_while_response_is_active():
|
||||
session = DuplexSession(
|
||||
session_id="sid-response-options-active",
|
||||
config=DuplexSessionConfig(instructions="base"),
|
||||
)
|
||||
session.begin_response()
|
||||
|
||||
with pytest.raises(RuntimeError, match="active"):
|
||||
session.reserve_response_options(ResponseCreateOptions(instructions="too late"))
|
||||
|
||||
|
||||
def test_realtime_model_name_does_not_implicitly_enable_native_duplex():
|
||||
protocol = NativeRealtimeSessionProtocol({})
|
||||
|
||||
event = protocol._session_create_from_realtime({"model": "openbmb/MiniCPM-o-4_5"})
|
||||
|
||||
assert "minicpmo45_native_duplex" not in event["session"]["extra_body"]
|
||||
|
||||
|
||||
def test_realtime_explicit_native_duplex_flag_is_preserved():
|
||||
protocol = NativeRealtimeSessionProtocol({})
|
||||
|
||||
event = protocol._session_create_from_realtime(
|
||||
{
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
}
|
||||
)
|
||||
|
||||
assert event["session"]["extra_body"]["minicpmo45_native_duplex"] is True
|
||||
|
||||
|
||||
def test_realtime_explicit_query_native_duplex_flag_is_available_before_autostart():
|
||||
protocol = NativeRealtimeSessionProtocol(
|
||||
{
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"minicpmo45_native_duplex": "1",
|
||||
}
|
||||
)
|
||||
|
||||
event = json.loads(asyncio.run(protocol.receive_internal_event_text(None)))
|
||||
|
||||
assert event["type"] == "session.create"
|
||||
assert event["session"]["extra_body"]["minicpmo45_native_duplex"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_resume_heartbeat_and_event_ack_translate_without_session_autostart():
|
||||
protocol = NativeRealtimeSessionProtocol(
|
||||
{
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"minicpmo45_native_duplex": "1",
|
||||
}
|
||||
)
|
||||
|
||||
resume = await protocol._to_duplex_event(
|
||||
{
|
||||
"type": "session.resume",
|
||||
"session_id": "sid-resume",
|
||||
"incarnation": 3,
|
||||
"resume_token": "opaque-token",
|
||||
"last_received_server_event_seq": 17,
|
||||
}
|
||||
)
|
||||
heartbeat = await protocol._to_duplex_event(
|
||||
{
|
||||
"type": "session.heartbeat",
|
||||
"event_id": "event-heartbeat",
|
||||
}
|
||||
)
|
||||
event_ack = await protocol._to_duplex_event(
|
||||
{
|
||||
"type": "session.event_ack",
|
||||
"server_event_seq": 19,
|
||||
}
|
||||
)
|
||||
|
||||
assert resume == {
|
||||
"type": "session.resume",
|
||||
"session_id": "sid-resume",
|
||||
"incarnation": 3,
|
||||
"resume_token": "opaque-token",
|
||||
"last_received_server_event_seq": 17,
|
||||
}
|
||||
assert heartbeat == {
|
||||
"type": "session.heartbeat",
|
||||
"event_id": "event-heartbeat",
|
||||
}
|
||||
assert event_ack == {
|
||||
"type": "session.event_ack",
|
||||
"server_event_seq": 19,
|
||||
}
|
||||
assert protocol._opened is False
|
||||
|
||||
wire_protocol = NativeRealtimeSessionProtocol({})
|
||||
first_wire_event = json.loads(
|
||||
await wire_protocol.receive_internal_event_text(
|
||||
_ProtocolWebSocket(
|
||||
{
|
||||
"type": "session.resume",
|
||||
"session_id": "sid-resume",
|
||||
"incarnation": 3,
|
||||
"resume_token": "opaque-token",
|
||||
"last_received_server_event_seq": 17,
|
||||
}
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
assert first_wire_event["type"] == "session.resume"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_resume_query_suppresses_model_autostart_on_same_url():
|
||||
protocol = NativeRealtimeSessionProtocol(
|
||||
{
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"minicpmo45_native_duplex": "1",
|
||||
"resume": "1",
|
||||
}
|
||||
)
|
||||
resume = {
|
||||
"type": "session.resume",
|
||||
"session_id": "sid-resume",
|
||||
"incarnation": 3,
|
||||
"resume_token": "opaque-token",
|
||||
"last_received_server_event_seq": 17,
|
||||
}
|
||||
|
||||
first_wire_event = json.loads(await protocol.receive_internal_event_text(_ProtocolWebSocket(resume)))
|
||||
|
||||
assert first_wire_event == resume
|
||||
assert protocol._autostarted_default_session is False
|
||||
|
||||
|
||||
def test_realtime_projects_resume_lifecycle_events_without_duplex_prefix():
|
||||
protocol = NativeRealtimeSessionProtocol({})
|
||||
|
||||
created = protocol.encode_outbound_event(
|
||||
{
|
||||
"type": "session.created",
|
||||
"session": {"id": "sid-resume"},
|
||||
"incarnation": 3,
|
||||
"attachment_generation": 1,
|
||||
"resume_token": "first-token",
|
||||
}
|
||||
)[0]
|
||||
resumed = protocol.encode_outbound_event(
|
||||
{
|
||||
"type": "session.resumed",
|
||||
"session_id": "sid-resume",
|
||||
"incarnation": 3,
|
||||
"attachment_generation": 2,
|
||||
"resume_token": "second-token",
|
||||
}
|
||||
)[0]
|
||||
heartbeat = protocol.encode_outbound_event(
|
||||
{
|
||||
"type": "session.heartbeat_ack",
|
||||
"session_id": "sid-resume",
|
||||
}
|
||||
)[0]
|
||||
replaced = protocol.encode_outbound_event(
|
||||
{
|
||||
"type": "session.replaced",
|
||||
"session_id": "sid-resume",
|
||||
"attachment_generation": 1,
|
||||
}
|
||||
)[0]
|
||||
resync = protocol.encode_outbound_event(
|
||||
{
|
||||
"type": "session.resync_required",
|
||||
"session_id": "sid-resume",
|
||||
"reason": "journal_gap",
|
||||
}
|
||||
)[0]
|
||||
|
||||
assert created["type"] == "session.created"
|
||||
assert created["incarnation"] == 3
|
||||
assert created["attachment_generation"] == 1
|
||||
assert created["resume_token"] == "first-token"
|
||||
assert resumed["type"] == "session.resumed"
|
||||
assert resumed["resume_token"] == "second-token"
|
||||
assert heartbeat["type"] == "session.heartbeat_ack"
|
||||
assert replaced["type"] == "session.replaced"
|
||||
assert resync == {
|
||||
"type": "session.resync_required",
|
||||
"session_id": "sid-resume",
|
||||
"reason": "journal_gap",
|
||||
"event_id": resync["event_id"],
|
||||
}
|
||||
|
||||
|
||||
def test_duplex_session_registry_advances_incarnation_when_id_is_reused():
|
||||
registry = DuplexSessionRegistry()
|
||||
|
||||
first = registry.create(session_id="reused-session")
|
||||
registry.close(first.session_id)
|
||||
second = registry.create(session_id="reused-session")
|
||||
|
||||
assert first.incarnation == 0
|
||||
assert second.incarnation == 1
|
||||
|
||||
|
||||
def test_native_input_commit_does_not_advance_model_turn_identity():
|
||||
session = DuplexSession(
|
||||
session_id="sid-model-owned-turn",
|
||||
config=DuplexSessionConfig(extra_body={"auto_response": True}),
|
||||
)
|
||||
|
||||
first = session.commit_native_audio_input(transcript="first chunk")
|
||||
second = session.commit_native_audio_input(transcript="second chunk")
|
||||
|
||||
assert session.input_commit_seq == 2
|
||||
assert first.input_commit_seq == 1
|
||||
assert second.input_commit_seq == 2
|
||||
assert first.turn_id == second.turn_id == session.turn_id == 0
|
||||
|
||||
|
||||
def test_duplex_session_owns_response_and_overlap_identity():
|
||||
session = DuplexSession(session_id="sid-state-owner", config=DuplexSessionConfig())
|
||||
|
||||
response_id = session.begin_response()
|
||||
session.accumulate_overlap_speech(320)
|
||||
session.accumulate_overlap_speech(180)
|
||||
session.end_response()
|
||||
|
||||
assert session.active_response_id is None
|
||||
assert session.last_response_id == response_id
|
||||
assert session.overlap_speech_ms == 500
|
||||
assert session.reset_overlap_speech() == 500
|
||||
assert session.overlap_speech_ms == 0
|
||||
|
||||
|
||||
def test_duplex_session_composes_single_owner_ledgers():
|
||||
session = DuplexSession(session_id="sid-ledgers", config=DuplexSessionConfig())
|
||||
|
||||
session.append_text("hello")
|
||||
session.bind_request("req-1")
|
||||
response_id = session.begin_response(turn_id=3)
|
||||
session.mark_audio_sent(duration_ms=240)
|
||||
|
||||
assert not hasattr(session, "input")
|
||||
assert not hasattr(session, "response")
|
||||
assert not hasattr(session, "playback_ledger")
|
||||
assert not hasattr(session, "conversation")
|
||||
assert session.pending_text == ("hello",)
|
||||
assert session.active_request_id == "req-1"
|
||||
assert session.active_response_id == response_id
|
||||
assert session.active_response_turn_id == 3
|
||||
assert session.playback.sent_ms == 240
|
||||
|
||||
try:
|
||||
session.playback.sent_ms = 480
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("playback view must be immutable")
|
||||
|
||||
|
||||
def test_duplex_barge_in_advances_epoch_and_drops_uncommitted_assistant_text():
|
||||
registry = DuplexSessionRegistry()
|
||||
session = registry.create()
|
||||
response_id = session.begin_response()
|
||||
session.bind_request("chatcmpl-duplex-test")
|
||||
session.append_assistant_text("unplayed answer")
|
||||
|
||||
new_epoch = session.barge_in()
|
||||
|
||||
assert response_id is not None
|
||||
assert new_epoch == 1
|
||||
assert session.epoch == 1
|
||||
assert session.active_request_id is None
|
||||
assert session.active_response_id is None
|
||||
assert session.assistant_text_buffer == ()
|
||||
assert session.history == ()
|
||||
assert session.turn_state == DuplexTurnState.BARGE_IN
|
||||
|
||||
|
||||
def test_duplex_playback_ack_tracks_committed_cursor_separately():
|
||||
registry = DuplexSessionRegistry()
|
||||
session = registry.create()
|
||||
session.mark_audio_sent(duration_ms=10_000)
|
||||
|
||||
session.acknowledge_playback(played_ms=2_000)
|
||||
|
||||
assert session.playback.generated_ms == 10_000
|
||||
assert session.playback.sent_ms == 10_000
|
||||
assert session.playback.played_ms == 2_000
|
||||
assert session.playback.committed_ms == 2_000
|
||||
|
||||
|
||||
def test_duplex_history_commit_uses_audio_text_alignment_marks():
|
||||
registry = DuplexSessionRegistry()
|
||||
session = registry.create()
|
||||
session.begin_response()
|
||||
|
||||
session.append_assistant_text("hello ")
|
||||
session.mark_audio_sent(duration_ms=1_000, text_chars=6)
|
||||
session.append_assistant_text("world")
|
||||
session.mark_audio_sent(duration_ms=2_000, text_chars=11)
|
||||
session.acknowledge_playback(played_ms=1_200, committed_ms=1_200)
|
||||
|
||||
committed = session.end_response(commit_text=True)
|
||||
|
||||
assert committed == {"role": "assistant", "content": "hello w"}
|
||||
|
||||
|
||||
def test_turn_controller_accepts_external_signal_source():
|
||||
registry = DuplexSessionRegistry()
|
||||
session = registry.create()
|
||||
controller = DuplexTurnController()
|
||||
|
||||
event = controller.signal(session, DuplexTurnEventType.USER_STARTED.value)
|
||||
|
||||
assert event["type"] == "turn.event"
|
||||
assert event["event"] == "user_started"
|
||||
assert event["turn_state"] == "user_speaking"
|
||||
assert session.turn_state == DuplexTurnState.USER_SPEAKING
|
||||
|
||||
|
||||
def test_duplex_capabilities_do_not_claim_core_kv_or_input_append():
|
||||
registry = DuplexSessionRegistry()
|
||||
session = registry.create()
|
||||
caps = session.capabilities.as_dict()
|
||||
|
||||
assert caps["implementation_level"] == "serving_session_adapter"
|
||||
assert caps["supports_kv_lease"] is False
|
||||
assert caps["supports_input_append"] is False
|
||||
assert caps["supports_reencode_context"] is True
|
||||
assert caps["adapter_patterns"] == ["chunk_group_append"]
|
||||
assert "turn_commit_only" in caps["input_modes"]
|
||||
assert "client_event" in caps["signal_sources"]
|
||||
|
||||
|
||||
def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
|
||||
caps = DuplexCapabilities.minicpmo45_native(max_sessions=2).as_dict()
|
||||
|
||||
assert caps["implementation_level"] == "model_native_duplex"
|
||||
assert caps["supports_input_append"] is True
|
||||
assert caps["input_modes"] == ["append_audio_chunk"]
|
||||
assert caps["adapter_patterns"] == ["scheduler_data_plane"]
|
||||
assert caps["supports_model_internal_state"] is True
|
||||
assert caps["requires_model_runner_kv"] is True
|
||||
assert caps["requires_native_stage_role"] is True
|
||||
assert caps["supports_kv_lease"] is False
|
||||
assert caps["supports_core_kv_lease"] is False
|
||||
assert caps["supports_stage_resumption"] is True
|
||||
assert caps["supports_scheduler_native_append"] is False
|
||||
assert caps["supports_core_resumable_request"] is True
|
||||
assert caps["supports_stage_connector_handoff"] is True
|
||||
assert caps["supports_audio_truncate"] is True
|
||||
assert caps["supports_barge_in"] is False
|
||||
assert caps["target_barge_in_latency_ms"] is None
|
||||
assert caps["supports_multi_session"] is True
|
||||
assert caps["supports_multi_session_same_replica"] is True
|
||||
assert caps["supports_session_lease"] is True
|
||||
assert caps["supports_session_resume"] is True
|
||||
assert caps["session_admission_mode"] == "engine_managed"
|
||||
assert caps["stage_handoff_transport"] == "scheduler_data_plane"
|
||||
|
||||
|
||||
def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment():
|
||||
caps = DuplexCapabilities.minicpmo45_native(max_sessions=1).as_dict()
|
||||
|
||||
assert caps["supports_multi_session"] is False
|
||||
assert caps["supports_multi_session_same_replica"] is False
|
||||
|
||||
|
||||
def test_duplex_overlap_policy_defaults_and_invalid_values_to_listen_only():
|
||||
assert DuplexSessionConfig().overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
|
||||
assert DuplexSessionConfig._normalize_overlap_policy("not-a-policy") == DuplexOverlapPolicy.LISTEN_ONLY.value
|
||||
@@ -0,0 +1,387 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.openai.session_attachment import (
|
||||
DuplexEventJournal,
|
||||
DuplexJournalGapError,
|
||||
DuplexJournalOverflowError,
|
||||
DuplexResumeCredential,
|
||||
DuplexSessionAttachmentRegistry,
|
||||
InvalidResumeTokenError,
|
||||
ResumeToken,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _Clock:
|
||||
def __init__(self) -> None:
|
||||
self.value = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.value
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.value += seconds
|
||||
|
||||
|
||||
def _decoded_token_bytes(token: ResumeToken) -> bytes:
|
||||
padding = "=" * (-len(token.plaintext) % 4)
|
||||
return base64.urlsafe_b64decode(token.plaintext + padding)
|
||||
|
||||
|
||||
def test_resume_credential_uses_256_bit_token_digest_and_redacted_repr(monkeypatch) -> None:
|
||||
token = ResumeToken.generate()
|
||||
credential = DuplexResumeCredential.from_token(token)
|
||||
calls = []
|
||||
|
||||
import vllm_omni.experimental.fullduplex.openai.session_attachment as attachment_module
|
||||
|
||||
original_compare = attachment_module.hmac.compare_digest
|
||||
|
||||
def recording_compare(left, right):
|
||||
calls.append((left, right))
|
||||
return original_compare(left, right)
|
||||
|
||||
monkeypatch.setattr(attachment_module.hmac, "compare_digest", recording_compare)
|
||||
|
||||
assert len(_decoded_token_bytes(token)) == 32
|
||||
assert credential.verify(token.plaintext) is True
|
||||
assert credential.verify("invalid-token") is False
|
||||
assert len(calls) == 2
|
||||
assert token.plaintext not in repr(token)
|
||||
assert token.plaintext not in repr(credential)
|
||||
assert not hasattr(credential, "plaintext")
|
||||
assert len(credential.token_digest) == 32
|
||||
|
||||
|
||||
def test_resume_credential_rotation_revokes_old_token() -> None:
|
||||
first = ResumeToken.generate()
|
||||
credential = DuplexResumeCredential.from_token(first)
|
||||
|
||||
second = credential.rotate()
|
||||
|
||||
assert credential.verify(first.plaintext) is False
|
||||
assert credential.verify(second.plaintext) is True
|
||||
assert second.plaintext != first.plaintext
|
||||
|
||||
|
||||
def test_event_journal_sequences_acknowledges_and_replays_exact_payloads() -> None:
|
||||
clock = _Clock()
|
||||
journal = DuplexEventJournal(max_bytes=4096, ttl_s=60.0, clock=clock)
|
||||
|
||||
first = journal.record({"type": "response.audio.delta", "delta": "AAAA"})
|
||||
second = journal.record({"type": "response.done", "response_id": "resp-1"})
|
||||
|
||||
assert first.sequence == 1
|
||||
assert second.sequence == 2
|
||||
assert first.payload["server_event_seq"] == 1
|
||||
assert second.payload["server_event_seq"] == 2
|
||||
assert first.encoded_bytes == len(
|
||||
json.dumps(dict(first.payload), separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||||
)
|
||||
assert [entry.sequence for entry in journal.replay_after(0)] == [1, 2]
|
||||
assert journal.acknowledge(1) == 1
|
||||
assert [entry.sequence for entry in journal.replay_after(1)] == [2]
|
||||
assert journal.acknowledge(1) == 0
|
||||
|
||||
|
||||
def test_event_journal_ttl_pruning_reports_replay_gap() -> None:
|
||||
clock = _Clock()
|
||||
journal = DuplexEventJournal(max_bytes=4096, ttl_s=5.0, clock=clock)
|
||||
journal.record({"type": "one"})
|
||||
clock.advance(3.0)
|
||||
journal.record({"type": "two"})
|
||||
clock.advance(3.0)
|
||||
|
||||
assert journal.prune() == 1
|
||||
with pytest.raises(DuplexJournalGapError, match="older than retained"):
|
||||
journal.replay_after(0)
|
||||
assert [entry.payload["type"] for entry in journal.replay_after(1)] == ["two"]
|
||||
|
||||
|
||||
def test_event_journal_overflow_is_explicit_and_does_not_evict_silently() -> None:
|
||||
clock = _Clock()
|
||||
journal = DuplexEventJournal(max_bytes=100, ttl_s=60.0, clock=clock)
|
||||
first = journal.record({"type": "small", "value": "x"})
|
||||
|
||||
with pytest.raises(DuplexJournalOverflowError, match="byte limit"):
|
||||
journal.record({"type": "large", "value": "x" * 200})
|
||||
|
||||
assert journal.overflowed is True
|
||||
assert journal.retained_bytes == first.encoded_bytes
|
||||
with pytest.raises(DuplexJournalGapError, match="overflowed"):
|
||||
journal.replay_after(0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_registry_resume_rotates_token_replays_and_atomically_replaces_attachment() -> None:
|
||||
clock = _Clock()
|
||||
registry = DuplexSessionAttachmentRegistry(
|
||||
replay_ttl_s=60.0,
|
||||
replay_max_bytes_per_session=4096,
|
||||
clock=clock,
|
||||
)
|
||||
sends_a = []
|
||||
closes_a = []
|
||||
sends_b = []
|
||||
closes_b = []
|
||||
|
||||
async def send_a(payload):
|
||||
sends_a.append(payload)
|
||||
|
||||
async def close_a(reason):
|
||||
closes_a.append(reason)
|
||||
|
||||
async def send_b(payload):
|
||||
sends_b.append(payload)
|
||||
|
||||
async def close_b(reason):
|
||||
closes_b.append(reason)
|
||||
|
||||
created = await registry.create("sid", incarnation=2, send=send_a, close=close_a)
|
||||
await registry.send_event("sid", {"type": "event-1"})
|
||||
await registry.send_event("sid", {"type": "event-2"})
|
||||
|
||||
resumed = await registry.resume(
|
||||
"sid",
|
||||
incarnation=2,
|
||||
resume_token=created.resume_token.plaintext,
|
||||
last_received_server_event_seq=1,
|
||||
send=send_b,
|
||||
close=close_b,
|
||||
)
|
||||
|
||||
assert resumed.attachment_generation == 2
|
||||
assert resumed.resume_token.plaintext != created.resume_token.plaintext
|
||||
assert [entry.sequence for entry in resumed.replay_entries] == [2]
|
||||
assert resumed.replaced_attachment is not None
|
||||
assert resumed.replaced_attachment.generation == 1
|
||||
assert await registry.is_current_attachment("sid", created.attachment_generation) is False
|
||||
assert await registry.is_current_attachment("sid", resumed.attachment_generation) is True
|
||||
assert await registry.detach("sid", attachment_generation=1) is False
|
||||
assert await registry.detach("sid", attachment_generation=2) is True
|
||||
assert [(payload["type"], payload["server_event_seq"]) for payload in sends_a] == [
|
||||
("event-1", 1),
|
||||
("event-2", 2),
|
||||
]
|
||||
assert closes_a == [] and sends_b == [] and closes_b == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_registry_resume_sends_activation_then_replay_before_new_live_events() -> None:
|
||||
registry = DuplexSessionAttachmentRegistry(
|
||||
replay_ttl_s=60.0,
|
||||
replay_max_bytes_per_session=4096,
|
||||
)
|
||||
wire = []
|
||||
activation_started = asyncio.Event()
|
||||
release_activation = asyncio.Event()
|
||||
|
||||
async def old_send(payload):
|
||||
del payload
|
||||
|
||||
async def close(reason):
|
||||
del reason
|
||||
|
||||
async def new_send(payload):
|
||||
wire.append(payload["type"])
|
||||
if payload["type"] == "session.resumed":
|
||||
activation_started.set()
|
||||
await release_activation.wait()
|
||||
|
||||
created = await registry.create("sid-order", incarnation=0, send=old_send, close=close)
|
||||
await registry.detach("sid-order", attachment_generation=1)
|
||||
await registry.send_event("sid-order", {"type": "replayed"})
|
||||
resume_task = asyncio.create_task(
|
||||
registry.resume(
|
||||
"sid-order",
|
||||
incarnation=0,
|
||||
resume_token=created.resume_token.plaintext,
|
||||
last_received_server_event_seq=0,
|
||||
send=new_send,
|
||||
close=close,
|
||||
activation_payload_factory=lambda _token, _generation: {"type": "session.resumed"},
|
||||
)
|
||||
)
|
||||
await activation_started.wait()
|
||||
live_task = asyncio.create_task(registry.send_event("sid-order", {"type": "live"}))
|
||||
|
||||
release_activation.set()
|
||||
await asyncio.gather(resume_task, live_task)
|
||||
|
||||
assert wire == ["session.resumed", "replayed", "live"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recovery() -> None:
|
||||
registry = DuplexSessionAttachmentRegistry(
|
||||
replay_ttl_s=60.0,
|
||||
replay_max_bytes_per_session=4096,
|
||||
)
|
||||
|
||||
async def send(payload):
|
||||
del payload
|
||||
|
||||
async def failing_send(payload):
|
||||
del payload
|
||||
raise RuntimeError("transport lost before rotated token arrived")
|
||||
|
||||
async def close(reason):
|
||||
del reason
|
||||
|
||||
created = await registry.create("sid-recovery", incarnation=0, send=send, close=close)
|
||||
await registry.detach("sid-recovery", attachment_generation=1)
|
||||
|
||||
with pytest.raises(RuntimeError, match="transport lost"):
|
||||
await registry.resume(
|
||||
"sid-recovery",
|
||||
incarnation=0,
|
||||
resume_token=created.resume_token.plaintext,
|
||||
last_received_server_event_seq=0,
|
||||
send=failing_send,
|
||||
close=close,
|
||||
activation_payload_factory=lambda token, generation: {
|
||||
"type": "session.resumed",
|
||||
"resume_token": token.plaintext,
|
||||
"attachment_generation": generation,
|
||||
},
|
||||
)
|
||||
|
||||
recovered = await registry.resume(
|
||||
"sid-recovery",
|
||||
incarnation=0,
|
||||
resume_token=created.resume_token.plaintext,
|
||||
last_received_server_event_seq=0,
|
||||
send=send,
|
||||
close=close,
|
||||
)
|
||||
|
||||
assert recovered.attachment_generation == 3
|
||||
with pytest.raises(InvalidResumeTokenError):
|
||||
await registry.resume(
|
||||
"sid-recovery",
|
||||
incarnation=0,
|
||||
resume_token=created.resume_token.plaintext,
|
||||
last_received_server_event_seq=0,
|
||||
send=send,
|
||||
close=close,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_registry_concurrent_resume_allows_exactly_one_rotated_token_winner() -> None:
|
||||
registry = DuplexSessionAttachmentRegistry(
|
||||
replay_ttl_s=60.0,
|
||||
replay_max_bytes_per_session=4096,
|
||||
)
|
||||
|
||||
async def send(payload):
|
||||
del payload
|
||||
|
||||
async def close(reason):
|
||||
del reason
|
||||
|
||||
created = await registry.create("sid-race", incarnation=0, send=send, close=close)
|
||||
|
||||
async def attempt():
|
||||
try:
|
||||
return await registry.resume(
|
||||
"sid-race",
|
||||
incarnation=0,
|
||||
resume_token=created.resume_token.plaintext,
|
||||
last_received_server_event_seq=0,
|
||||
send=send,
|
||||
close=close,
|
||||
)
|
||||
except InvalidResumeTokenError as exc:
|
||||
return exc
|
||||
|
||||
results = await asyncio.gather(attempt(), attempt())
|
||||
|
||||
assert sum(not isinstance(result, Exception) for result in results) == 1
|
||||
assert sum(isinstance(result, InvalidResumeTokenError) for result in results) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> None:
|
||||
registry = DuplexSessionAttachmentRegistry(
|
||||
replay_ttl_s=60.0,
|
||||
replay_max_bytes_per_session=4096,
|
||||
)
|
||||
|
||||
async def send(payload):
|
||||
del payload
|
||||
|
||||
async def close(reason):
|
||||
del reason
|
||||
|
||||
created_a = await registry.create("sid-a", incarnation=1, send=send, close=close)
|
||||
created_b = await registry.create("sid-b", incarnation=4, send=send, close=close)
|
||||
await registry.detach("sid-a", attachment_generation=1)
|
||||
await registry.detach("sid-b", attachment_generation=1)
|
||||
await registry.send_event("sid-a", {"type": "a-only"})
|
||||
await registry.send_event("sid-b", {"type": "b-only"})
|
||||
|
||||
resumed_a = await registry.resume(
|
||||
"sid-a",
|
||||
incarnation=1,
|
||||
resume_token=created_a.resume_token.plaintext,
|
||||
last_received_server_event_seq=0,
|
||||
send=send,
|
||||
close=close,
|
||||
)
|
||||
resumed_b = await registry.resume(
|
||||
"sid-b",
|
||||
incarnation=4,
|
||||
resume_token=created_b.resume_token.plaintext,
|
||||
last_received_server_event_seq=0,
|
||||
send=send,
|
||||
close=close,
|
||||
)
|
||||
|
||||
assert [entry.payload["type"] for entry in resumed_a.replay_entries] == ["a-only"]
|
||||
assert [entry.payload["type"] for entry in resumed_b.replay_entries] == ["b-only"]
|
||||
with pytest.raises(ValueError, match="incarnation mismatch"):
|
||||
await registry.resume(
|
||||
"sid-a",
|
||||
incarnation=2,
|
||||
resume_token=resumed_a.resume_token.plaintext,
|
||||
last_received_server_event_seq=1,
|
||||
send=send,
|
||||
close=close,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_registry_repr_never_contains_plaintext_tokens() -> None:
|
||||
registry = DuplexSessionAttachmentRegistry(
|
||||
replay_ttl_s=60.0,
|
||||
replay_max_bytes_per_session=4096,
|
||||
)
|
||||
|
||||
async def send(payload):
|
||||
del payload
|
||||
|
||||
async def close(reason):
|
||||
del reason
|
||||
|
||||
created = await registry.create("sid-repr", incarnation=0, send=send, close=close)
|
||||
await registry.detach("sid-repr", attachment_generation=1)
|
||||
entry = await registry.send_event(
|
||||
"sid-repr",
|
||||
{"type": "session.resumed", "resume_token": created.resume_token.plaintext},
|
||||
)
|
||||
|
||||
assert entry is not None
|
||||
assert created.resume_token.plaintext not in repr(created)
|
||||
assert created.resume_token.plaintext not in repr(entry)
|
||||
assert created.resume_token.plaintext not in repr(registry)
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_api_server_does_not_import_experimental_duplex_at_module_load() -> None:
|
||||
api_server = Path(__file__).parents[3] / "vllm_omni/entrypoints/openai/api_server.py"
|
||||
module = ast.parse(api_server.read_text())
|
||||
|
||||
top_level_imports = {
|
||||
node.module for node in module.body if isinstance(node, ast.ImportFrom) and isinstance(node.module, str)
|
||||
}
|
||||
|
||||
assert not any(module.startswith("vllm_omni.experimental.fullduplex") for module in top_level_imports)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -138,6 +138,112 @@ async def test_receive_config_accepts_client_legacy_aliases():
|
||||
assert config.frame_filter_threshold == 0.87
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_frame_ack_reports_receiver_buffer_state():
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.config",
|
||||
"model": "test",
|
||||
"enable_frame_filter": False,
|
||||
}
|
||||
),
|
||||
json.dumps(
|
||||
{
|
||||
"type": "video.frame",
|
||||
"data": _b64(_make_jpeg()),
|
||||
"frame_id": "frame-7",
|
||||
"pts_ms": 700,
|
||||
"capture_ts_ms": 1234.5,
|
||||
}
|
||||
),
|
||||
json.dumps({"type": "video.done"}),
|
||||
]
|
||||
)
|
||||
handler = QwenOmniStreamingVideoHandler(chat_service=object())
|
||||
|
||||
await handler.handle_session(ws)
|
||||
|
||||
ack = next(message for message in ws.sent if message.get("type") == "video.frame.ack")
|
||||
assert ack["frame_id"] == "frame-7"
|
||||
assert ack["pts_ms"] == 700
|
||||
assert ack["capture_ts_ms"] == 1234.5
|
||||
assert ack["accepted"] is True
|
||||
assert ack["buffered_frames"] == 1
|
||||
assert ack["server_receive_ts_ms"] > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_frames_consumed_is_emitted_after_engine_uses_frame_prompt():
|
||||
class OneOutputEngine:
|
||||
def generate(self, **_kwargs):
|
||||
async def _gen():
|
||||
await asyncio.sleep(0.05)
|
||||
yield _text_result("visible")
|
||||
|
||||
return _gen()
|
||||
|
||||
class CapturingHandler(QwenOmniStreamingVideoHandler):
|
||||
async def _preprocess_to_engine_prompt(self, request):
|
||||
return {"prompt": "with-video"}
|
||||
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.config",
|
||||
"model": "test",
|
||||
"modalities": ["text"],
|
||||
"num_frames": 1,
|
||||
"enable_frame_filter": False,
|
||||
}
|
||||
),
|
||||
json.dumps(
|
||||
{
|
||||
"type": "video.frame",
|
||||
"data": _b64(_make_jpeg()),
|
||||
"frame_id": "frame-9",
|
||||
"pts_ms": 900,
|
||||
"source_pts_ms": 880,
|
||||
"quality_profile": "balanced",
|
||||
}
|
||||
),
|
||||
json.dumps({"type": "video.query", "text": "describe"}),
|
||||
json.dumps({"type": "video.done"}),
|
||||
]
|
||||
)
|
||||
handler = CapturingHandler(
|
||||
chat_service=object(),
|
||||
engine_client=OneOutputEngine(),
|
||||
idle_timeout=2.0,
|
||||
)
|
||||
|
||||
await handler.handle_session(ws)
|
||||
|
||||
consumed = next(message for message in ws.sent if message.get("type") == "video.frames.consumed")
|
||||
assert consumed["frame_ids"] == ["frame-9"]
|
||||
assert consumed["latest_pts_ms"] == 900
|
||||
assert consumed["request_id"].startswith("video-")
|
||||
assert consumed["model_selected_ts_ms"] > 0
|
||||
assert consumed["frames"] == [
|
||||
{
|
||||
"frame_id": "frame-9",
|
||||
"pts_ms": 900,
|
||||
"source_pts_ms": 880,
|
||||
"quality_profile": "balanced",
|
||||
"receiver_received_ts_ms": consumed["frames"][0]["receiver_received_ts_ms"],
|
||||
"decoded_ready_ts_ms": consumed["frames"][0]["decoded_ready_ts_ms"],
|
||||
}
|
||||
]
|
||||
assert consumed["frames"][0]["receiver_received_ts_ms"] > 0
|
||||
assert consumed["frames"][0]["decoded_ready_ts_ms"] >= consumed["frames"][0]["receiver_received_ts_ms"]
|
||||
assert consumed["model_selected_ts_ms"] >= consumed["frames"][0]["decoded_ready_ts_ms"]
|
||||
assert ws.sent.index(consumed) < next(
|
||||
index for index, message in enumerate(ws.sent) if message.get("type") == "response.text.delta"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audio_in_video_sets_mm_processor_kwargs():
|
||||
captured_requests = []
|
||||
|
||||
@@ -0,0 +1,600 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.engine.messages import OutputMessage
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.client_request_state import ClientRequestState
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
|
||||
DuplexOutputAction,
|
||||
DuplexOutputDecision,
|
||||
duplex_resource_request_id,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision
|
||||
from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient
|
||||
from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats, StageStats
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
|
||||
def test_async_omni_uses_extracted_duplex_request_client():
|
||||
app = object.__new__(AsyncOmni)
|
||||
app.engine = SimpleNamespace(num_stages=2)
|
||||
app.request_states = {}
|
||||
app._duplex_request_client = None
|
||||
|
||||
client = app._get_duplex_request_client()
|
||||
|
||||
assert isinstance(client, DuplexRequestClient)
|
||||
assert not hasattr(client, "_async_omni")
|
||||
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_open_duplex_forwards_session_config_and_timeout():
|
||||
calls = []
|
||||
output_handler_starts = 0
|
||||
|
||||
async def open_duplex_session_async(session_id, **kwargs):
|
||||
calls.append((session_id, kwargs))
|
||||
return {"ok": True}
|
||||
|
||||
app = object.__new__(AsyncOmni)
|
||||
app.engine = SimpleNamespace(open_duplex_session_async=open_duplex_session_async)
|
||||
|
||||
def start_output_handler():
|
||||
nonlocal output_handler_starts
|
||||
output_handler_starts += 1
|
||||
|
||||
app._final_output_handler = start_output_handler
|
||||
|
||||
result = await app.open_duplex_session_async(
|
||||
"sid",
|
||||
fence=DuplexFence("sid"),
|
||||
capabilities={"implementation_level": "model_native_duplex"},
|
||||
session_config={"instructions": "Be brief.", "idle_timeout_s": 30},
|
||||
timeout=7.5,
|
||||
)
|
||||
|
||||
assert result == {"ok": True}
|
||||
assert output_handler_starts == 1
|
||||
assert calls == [
|
||||
(
|
||||
"sid",
|
||||
{
|
||||
"session_mode": "duplex",
|
||||
"capabilities": {"implementation_level": "model_native_duplex"},
|
||||
"session_config": {"instructions": "Be brief.", "idle_timeout_s": 30},
|
||||
"fence": DuplexFence("sid"),
|
||||
"timeout": 7.5,
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_duplex_runtime_controls_forward_timeout():
|
||||
calls = []
|
||||
cancelled_fence = DuplexFence("sid", epoch=2, turn_id=3)
|
||||
next_fence = DuplexFence("sid", epoch=3, turn_id=3)
|
||||
|
||||
async def append_duplex_input_async(session_id, **kwargs):
|
||||
calls.append(("append", session_id, kwargs))
|
||||
return {"ok": True}
|
||||
|
||||
async def signal_duplex_turn_async(session_id, **kwargs):
|
||||
calls.append(("signal", session_id, kwargs))
|
||||
return {"ok": True}
|
||||
|
||||
async def close_duplex_session_async(session_id, **kwargs):
|
||||
calls.append(("close", session_id, kwargs))
|
||||
return {"ok": True}
|
||||
|
||||
async def touch_duplex_session_async(session_id, **kwargs):
|
||||
calls.append(("touch", session_id, kwargs))
|
||||
return {"ok": True}
|
||||
|
||||
async def resume_duplex_session_async(session_id, **kwargs):
|
||||
calls.append(("resume", session_id, kwargs))
|
||||
return {"ok": True}
|
||||
|
||||
app = object.__new__(AsyncOmni)
|
||||
app.request_states = {}
|
||||
app._final_output_handler = lambda: None
|
||||
app.engine = SimpleNamespace(
|
||||
append_duplex_input_async=append_duplex_input_async,
|
||||
signal_duplex_turn_async=signal_duplex_turn_async,
|
||||
close_duplex_session_async=close_duplex_session_async,
|
||||
touch_duplex_session_async=touch_duplex_session_async,
|
||||
resume_duplex_session_async=resume_duplex_session_async,
|
||||
)
|
||||
|
||||
await app.append_duplex_input_async(
|
||||
"sid",
|
||||
mode="append_audio_chunk",
|
||||
payload={},
|
||||
fence=cancelled_fence,
|
||||
timeout=12.5,
|
||||
)
|
||||
await app.signal_duplex_turn_async(
|
||||
"sid",
|
||||
event="barge_in",
|
||||
fence=cancelled_fence,
|
||||
next_fence=next_fence,
|
||||
timeout=13.5,
|
||||
)
|
||||
await app.close_duplex_session_async("sid", reason="done", fence=next_fence, timeout=14.5)
|
||||
await app.touch_duplex_session_async(
|
||||
"sid",
|
||||
fence=next_fence,
|
||||
activity=DuplexLeaseActivity.HEARTBEAT,
|
||||
timeout=15.5,
|
||||
)
|
||||
await app.resume_duplex_session_async(
|
||||
"sid",
|
||||
fence=next_fence,
|
||||
expected_lease_generation=3,
|
||||
timeout=16.5,
|
||||
)
|
||||
|
||||
assert calls == [
|
||||
(
|
||||
"append",
|
||||
"sid",
|
||||
{
|
||||
"mode": "append_audio_chunk",
|
||||
"payload": {},
|
||||
"final": False,
|
||||
"expected_epoch": None,
|
||||
"fence": cancelled_fence,
|
||||
"timeout": 12.5,
|
||||
},
|
||||
),
|
||||
(
|
||||
"signal",
|
||||
"sid",
|
||||
{
|
||||
"event": "barge_in",
|
||||
"fence": cancelled_fence,
|
||||
"next_fence": next_fence,
|
||||
"timeout": 13.5,
|
||||
},
|
||||
),
|
||||
(
|
||||
"close",
|
||||
"sid",
|
||||
{
|
||||
"reason": "done",
|
||||
"fence": next_fence,
|
||||
"timeout": 14.5,
|
||||
},
|
||||
),
|
||||
(
|
||||
"touch",
|
||||
"sid",
|
||||
{
|
||||
"fence": next_fence,
|
||||
"activity": DuplexLeaseActivity.HEARTBEAT,
|
||||
"timeout": 15.5,
|
||||
},
|
||||
),
|
||||
(
|
||||
"resume",
|
||||
"sid",
|
||||
{
|
||||
"fence": next_fence,
|
||||
"expected_lease_generation": 3,
|
||||
"timeout": 16.5,
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_duplex_append_can_defer_data_plane_output_collection():
|
||||
app = object.__new__(AsyncOmni)
|
||||
request_id = duplex_resource_request_id(DuplexFence("sid"), "stage0")
|
||||
|
||||
async def append_duplex_input_async(session_id, **kwargs):
|
||||
del session_id, kwargs
|
||||
request_state = app.request_states.get(request_id)
|
||||
assert request_state is not None
|
||||
assert request_state.metrics is not None
|
||||
return {
|
||||
"stage_results": [
|
||||
{
|
||||
"result": {
|
||||
"data_plane_append": True,
|
||||
"request_id": request_id,
|
||||
"response_stage_id": 1,
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
app.engine = SimpleNamespace(append_duplex_input_async=append_duplex_input_async)
|
||||
app.request_states = {}
|
||||
app._final_output_handler = lambda: None
|
||||
|
||||
result = await app.append_duplex_input_async(
|
||||
"sid",
|
||||
mode="append_audio_chunk",
|
||||
payload={},
|
||||
fence=DuplexFence("sid"),
|
||||
collect_outputs=False,
|
||||
)
|
||||
|
||||
assert "data_plane_outputs" not in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_duplex_append_forwards_fence():
|
||||
calls = []
|
||||
|
||||
async def append_duplex_input_async(session_id, **kwargs):
|
||||
calls.append((session_id, kwargs))
|
||||
return {"ok": True}
|
||||
|
||||
app = object.__new__(AsyncOmni)
|
||||
app.request_states = {}
|
||||
app._final_output_handler = lambda: None
|
||||
app.engine = SimpleNamespace(append_duplex_input_async=append_duplex_input_async)
|
||||
fence = DuplexFence("sid", epoch=1, turn_id=2)
|
||||
|
||||
await app.append_duplex_input_async(
|
||||
"sid",
|
||||
mode="append_audio_chunk",
|
||||
payload={},
|
||||
fence=fence,
|
||||
collect_outputs=False,
|
||||
)
|
||||
|
||||
assert calls[0][1]["fence"] is fence
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_duplex_collect_waits_for_response_stage():
|
||||
app = object.__new__(AsyncOmni)
|
||||
req_state = ClientRequestState("duplex-sid-e0-stage0")
|
||||
stage0_output = OmniRequestOutput(
|
||||
request_id="duplex-sid-e0-stage0",
|
||||
stage_id=0,
|
||||
finished=False,
|
||||
)
|
||||
stage1_output = OmniRequestOutput(
|
||||
request_id="duplex-sid-e0-stage0",
|
||||
stage_id=1,
|
||||
final_output_type="audio",
|
||||
finished=False,
|
||||
)
|
||||
await req_state.queue.put(
|
||||
OutputMessage(
|
||||
request_id="duplex-sid-e0-stage0",
|
||||
stage_id=0,
|
||||
engine_outputs=stage0_output,
|
||||
finished=False,
|
||||
)
|
||||
)
|
||||
await req_state.queue.put(
|
||||
OutputMessage(
|
||||
request_id="duplex-sid-e0-stage0",
|
||||
stage_id=1,
|
||||
engine_outputs=stage1_output,
|
||||
finished=False,
|
||||
)
|
||||
)
|
||||
|
||||
outputs = await app._collect_duplex_data_plane_outputs(
|
||||
"duplex-sid-e0-stage0",
|
||||
req_state,
|
||||
response_stage_id=1,
|
||||
timeout=1.0,
|
||||
)
|
||||
|
||||
assert outputs == [stage1_output]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplex_request_client_retains_output_route_at_segment_terminal():
|
||||
request_id = "duplex-sid-e0-stage0"
|
||||
request_state = ClientRequestState(request_id)
|
||||
request_states = {request_id: request_state}
|
||||
output = OmniRequestOutput(
|
||||
request_id=request_id,
|
||||
stage_id=1,
|
||||
final_output_type="audio",
|
||||
finished=True,
|
||||
)
|
||||
await request_state.queue.put(
|
||||
OutputMessage(
|
||||
request_id=request_id,
|
||||
stage_id=1,
|
||||
engine_outputs=output,
|
||||
finished=True,
|
||||
)
|
||||
)
|
||||
client = DuplexRequestClient(
|
||||
SimpleNamespace(),
|
||||
SimpleNamespace(
|
||||
request_states=request_states,
|
||||
num_stages=2,
|
||||
log_stats=False,
|
||||
),
|
||||
)
|
||||
|
||||
outputs = await client.collect_outputs(
|
||||
request_id,
|
||||
request_state,
|
||||
response_stage_id=1,
|
||||
timeout=1.0,
|
||||
)
|
||||
|
||||
assert outputs == [output]
|
||||
assert request_states[request_id] is request_state
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_duplex_direct_output_keeps_stage_metrics():
|
||||
app = object.__new__(AsyncOmni)
|
||||
app._duplex_request_client = None
|
||||
app._final_output_handler = lambda: None
|
||||
request_id = "duplex-sid-e0-stage0"
|
||||
req_state = ClientRequestState(request_id)
|
||||
decision = DuplexOutputDecision(
|
||||
action=DuplexOutputAction.DIRECT_RESPONSE,
|
||||
metadata={
|
||||
"duplex_direct_response": True,
|
||||
"duplex_native_decision": "speak",
|
||||
},
|
||||
)
|
||||
direct_output = attach_duplex_output_decision(
|
||||
OmniRequestOutput(
|
||||
request_id=request_id,
|
||||
stage_id=0,
|
||||
final_output_type="text",
|
||||
finished=True,
|
||||
),
|
||||
decision,
|
||||
)
|
||||
await req_state.queue.put(
|
||||
OutputMessage(
|
||||
request_id=request_id,
|
||||
stage_id=0,
|
||||
engine_outputs=direct_output,
|
||||
metrics=StageRequestStats(
|
||||
batch_id=1,
|
||||
batch_size=1,
|
||||
num_tokens_in=160,
|
||||
num_tokens_out=8,
|
||||
stage_gen_time_ms=80.0,
|
||||
rx_transfer_bytes=0,
|
||||
rx_decode_time_ms=0.0,
|
||||
rx_in_flight_time_ms=0.0,
|
||||
stage_stats=StageStats(total_token=8, total_gen_time_ms=80.0),
|
||||
vllm_ttft_ms=12.5,
|
||||
vllm_tpot_ms=9.0,
|
||||
vllm_itl_ms=8.75,
|
||||
vllm_itls_ms=[8.0, 9.5],
|
||||
),
|
||||
finished=True,
|
||||
)
|
||||
)
|
||||
|
||||
outputs = await app._collect_duplex_data_plane_outputs(
|
||||
request_id,
|
||||
req_state,
|
||||
response_stage_id=1,
|
||||
timeout=1.0,
|
||||
)
|
||||
|
||||
assert outputs == [direct_output]
|
||||
assert direct_output.metrics["stage_metrics"]["0"]["vllm_ttft_ms"] == 12.5
|
||||
assert direct_output.metrics["stage_metrics"]["0"]["vllm_tpot_ms"] == 9.0
|
||||
assert direct_output.metrics["stage_metrics"]["0"]["vllm_itls_ms"] == [8.0, 9.5]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplex_metrics_cursor_isolates_resumable_collect_windows():
|
||||
app = object.__new__(AsyncOmni)
|
||||
app._duplex_request_client = None
|
||||
app._final_output_handler = lambda: None
|
||||
request_id = "duplex-sid-e0-stage0"
|
||||
req_state = ClientRequestState(request_id)
|
||||
req_state.metrics = OrchestratorAggregator(
|
||||
num_stages=2,
|
||||
log_stats=False,
|
||||
wall_start_ts=0.0,
|
||||
final_stage_id_for_e2e=1,
|
||||
)
|
||||
|
||||
def stage_message(*, batch_id: int, tokens: int, ttft_ms: float, itls: list[float]) -> OutputMessage:
|
||||
decision = DuplexOutputDecision(
|
||||
action=DuplexOutputAction.DIRECT_RESPONSE,
|
||||
metadata={"duplex_direct_response": True, "duplex_native_decision": "speak"},
|
||||
)
|
||||
output = attach_duplex_output_decision(
|
||||
OmniRequestOutput(
|
||||
request_id=request_id,
|
||||
stage_id=0,
|
||||
final_output_type="text",
|
||||
finished=True,
|
||||
),
|
||||
decision,
|
||||
)
|
||||
return OutputMessage(
|
||||
request_id=request_id,
|
||||
stage_id=0,
|
||||
engine_outputs=output,
|
||||
metrics=StageRequestStats(
|
||||
batch_id=batch_id,
|
||||
batch_size=1,
|
||||
num_tokens_in=160,
|
||||
num_tokens_out=tokens,
|
||||
stage_gen_time_ms=80.0,
|
||||
rx_transfer_bytes=0,
|
||||
rx_decode_time_ms=0.0,
|
||||
rx_in_flight_time_ms=0.0,
|
||||
stage_stats=StageStats(total_token=tokens, total_gen_time_ms=80.0),
|
||||
vllm_ttft_ms=ttft_ms,
|
||||
vllm_tpot_ms=9.0,
|
||||
vllm_itl_ms=sum(itls) / len(itls),
|
||||
vllm_itls_ms=itls,
|
||||
),
|
||||
finished=True,
|
||||
)
|
||||
|
||||
await req_state.queue.put(stage_message(batch_id=1, tokens=71, ttft_ms=62.5, itls=[10.0, 11.0]))
|
||||
first = await app._collect_duplex_data_plane_outputs(
|
||||
request_id,
|
||||
req_state,
|
||||
response_stage_id=1,
|
||||
timeout=1.0,
|
||||
)
|
||||
await req_state.queue.put(stage_message(batch_id=2, tokens=155, ttft_ms=81.0, itls=[12.0, 13.0, 14.0]))
|
||||
second = await app._collect_duplex_data_plane_outputs(
|
||||
request_id,
|
||||
req_state,
|
||||
response_stage_id=1,
|
||||
timeout=1.0,
|
||||
)
|
||||
|
||||
assert first[0].metrics["stage_metrics"]["0"]["num_tokens_out"] == 71
|
||||
second_stage0 = second[0].metrics["stage_metrics"]["0"]
|
||||
assert second_stage0["num_tokens_out"] == 155
|
||||
assert second_stage0["vllm_ttft_ms"] == 81.0
|
||||
assert second_stage0["vllm_itls_ms"] == [12.0, 13.0, 14.0]
|
||||
|
||||
|
||||
def test_async_omni_duplex_direct_output_prefers_outer_control_metadata():
|
||||
decision = DuplexOutputDecision(
|
||||
action=DuplexOutputAction.DIRECT_RESPONSE,
|
||||
metadata={
|
||||
"duplex_direct_response": True,
|
||||
"duplex_native_decision": "listen",
|
||||
},
|
||||
)
|
||||
output = attach_duplex_output_decision(
|
||||
OmniRequestOutput(
|
||||
request_id="duplex-sid-e0-stage0",
|
||||
stage_id=0,
|
||||
request_output=SimpleNamespace(
|
||||
outputs=[],
|
||||
multimodal_output=SimpleNamespace(kind="processed-payload"),
|
||||
_custom_output={"inner": "completion"},
|
||||
),
|
||||
),
|
||||
decision,
|
||||
)
|
||||
|
||||
assert AsyncOmni._is_direct_duplex_data_plane_response(output)
|
||||
|
||||
|
||||
def test_async_omni_drops_unregistered_duplex_prefixed_output():
|
||||
app = object.__new__(AsyncOmni)
|
||||
app.request_states = {}
|
||||
output = OmniRequestOutput(request_id="duplex-unregistered-e0-stage0", stage_id=0)
|
||||
|
||||
keep_processing, request_id, stage_id, request_state = app._handle_output_message(
|
||||
OutputMessage(
|
||||
request_id=output.request_id,
|
||||
stage_id=0,
|
||||
engine_outputs=output,
|
||||
finished=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert keep_processing is True
|
||||
assert request_id is stage_id is request_state is None
|
||||
assert app.request_states == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_duplex_collect_wraps_raw_response_stage_output():
|
||||
app = object.__new__(AsyncOmni)
|
||||
app.log_stats = False
|
||||
app._enable_ar_profiler = False
|
||||
app.request_states = {}
|
||||
app.prom_metrics = SimpleNamespace(
|
||||
set_running=lambda value: None,
|
||||
set_waiting=lambda value: None,
|
||||
)
|
||||
app.engine = SimpleNamespace(
|
||||
num_stages=2,
|
||||
get_stage_metadata=lambda stage_id: SimpleNamespace(
|
||||
final_output=stage_id == 1,
|
||||
final_output_type="audio",
|
||||
),
|
||||
)
|
||||
req_state = ClientRequestState("duplex-sid-e0-stage0")
|
||||
raw_stage1_output = SimpleNamespace(
|
||||
finished=False,
|
||||
outputs=[],
|
||||
stage_durations={},
|
||||
peak_memory_mb=0.0,
|
||||
final_output_type="audio",
|
||||
)
|
||||
await req_state.queue.put(
|
||||
OutputMessage(
|
||||
request_id="duplex-sid-e0-stage0",
|
||||
stage_id=1,
|
||||
engine_outputs=raw_stage1_output,
|
||||
finished=False,
|
||||
)
|
||||
)
|
||||
|
||||
outputs = await app._collect_duplex_data_plane_outputs(
|
||||
"duplex-sid-e0-stage0",
|
||||
req_state,
|
||||
response_stage_id=1,
|
||||
timeout=1.0,
|
||||
)
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert outputs[0].request_id == "duplex-sid-e0-stage0"
|
||||
assert outputs[0].finished is False
|
||||
assert outputs[0].stage_id == 1
|
||||
assert outputs[0].final_output_type == "audio"
|
||||
assert outputs[0].request_output is raw_stage1_output
|
||||
|
||||
|
||||
def test_async_omni_duplex_request_info_includes_response_stage():
|
||||
request_id, response_stage_id = AsyncOmni._duplex_data_plane_request_info(
|
||||
{
|
||||
"stage_results": [
|
||||
{
|
||||
"result": {
|
||||
"data_plane_append": True,
|
||||
"request_id": "duplex-sid-e0-stage0",
|
||||
"response_stage_id": 1,
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
assert request_id == "duplex-sid-e0-stage0"
|
||||
assert response_stage_id == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplex_request_client_retains_output_route_when_close_fails():
|
||||
fence = DuplexFence("sid-close-route")
|
||||
request_id = duplex_resource_request_id(fence, "stage0")
|
||||
request_states = {request_id: ClientRequestState(request_id)}
|
||||
|
||||
async def close_duplex_session_async(session_id, **kwargs):
|
||||
del session_id, kwargs
|
||||
raise RuntimeError("close failed")
|
||||
|
||||
client = DuplexRequestClient(
|
||||
SimpleNamespace(close_duplex_session_async=close_duplex_session_async),
|
||||
SimpleNamespace(request_states=request_states),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="close failed"):
|
||||
await client.close(fence.session_id, reason="test", fence=fence, timeout=1.0)
|
||||
|
||||
assert request_id in request_states
|
||||
@@ -0,0 +1,123 @@
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai_api.test_duplex_handler import FakeChatService, TimedWebSocket
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
from vllm_omni.experimental.fullduplex.openai.protocol import DuplexSession, DuplexSessionConfig
|
||||
from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FenceRecordingEngine:
|
||||
output_modalities = ["text", "audio"]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, DuplexFence | None]] = []
|
||||
self.signal_next_fences: list[DuplexFence | None] = []
|
||||
|
||||
async def open_duplex_session_async(self, session_id, *, fence=None, **kwargs):
|
||||
del session_id, kwargs
|
||||
self.calls.append(("open", fence))
|
||||
return {"ok": True}
|
||||
|
||||
async def append_duplex_input_async(self, session_id, *, fence=None, **kwargs):
|
||||
del session_id, kwargs
|
||||
self.calls.append(("append", fence))
|
||||
return {"ok": True}
|
||||
|
||||
async def signal_duplex_turn_async(
|
||||
self,
|
||||
session_id,
|
||||
*,
|
||||
event,
|
||||
fence=None,
|
||||
next_fence=None,
|
||||
timeout=None,
|
||||
):
|
||||
del session_id, event, timeout
|
||||
self.calls.append(("signal", fence))
|
||||
self.signal_next_fences.append(next_fence)
|
||||
return {"ok": True}
|
||||
|
||||
async def close_duplex_session_async(self, session_id, *, fence=None, **kwargs):
|
||||
del session_id, kwargs
|
||||
self.calls.append(("close", fence))
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_forwards_fence_for_all_runtime_operations():
|
||||
engine = FenceRecordingEngine()
|
||||
app = object.__new__(AsyncOmni)
|
||||
app.engine = engine
|
||||
app.request_states = {}
|
||||
app._final_output_handler = lambda: None
|
||||
fence = DuplexFence("sid", epoch=1, turn_id=2)
|
||||
|
||||
await app.open_duplex_session_async("sid", fence=fence)
|
||||
await app.append_duplex_input_async(
|
||||
"sid", mode="append_audio_chunk", payload={}, fence=fence, collect_outputs=False
|
||||
)
|
||||
await app.signal_duplex_turn_async("sid", event="turn.end", fence=fence)
|
||||
await app.close_duplex_session_async("sid", reason="done", fence=fence)
|
||||
|
||||
assert engine.calls == [(operation, fence) for operation in ("open", "append", "signal", "close")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_handler_passes_current_session_fence_for_runtime_controls():
|
||||
engine = FenceRecordingEngine()
|
||||
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
|
||||
session = DuplexSession(
|
||||
session_id="sid",
|
||||
config=DuplexSessionConfig(),
|
||||
epoch=1,
|
||||
turn_id=2,
|
||||
)
|
||||
ws = TimedWebSocket()
|
||||
fence = DuplexFence("sid", epoch=1, turn_id=2)
|
||||
|
||||
await handler._open_runtime_session(session, ws.send_json)
|
||||
await handler._signal_runtime_session(session, "turn.end", send_json=ws.send_json)
|
||||
await handler._close_runtime_session(session, reason="done", send_json=ws.send_json)
|
||||
|
||||
assert engine.calls == [(operation, fence) for operation in ("open", "signal", "close")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_signal_uses_fence_captured_before_session_turn_advances():
|
||||
engine = FenceRecordingEngine()
|
||||
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
|
||||
session = DuplexSession(session_id="sid", config=DuplexSessionConfig())
|
||||
captured_fence = DuplexFence("sid", epoch=0, turn_id=0)
|
||||
session.turn_id = 1
|
||||
|
||||
await handler._signal_runtime_session(
|
||||
session,
|
||||
"input.commit",
|
||||
fence=captured_fence,
|
||||
)
|
||||
|
||||
assert engine.calls == [("signal", captured_fence)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_handler_forwards_cancel_fence_through_async_omni_facade():
|
||||
engine = FenceRecordingEngine()
|
||||
app = object.__new__(AsyncOmni)
|
||||
app.engine = engine
|
||||
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(app))
|
||||
session = DuplexSession(session_id="sid-cancel", config=DuplexSessionConfig())
|
||||
cancelled_fence = DuplexFence("sid-cancel", epoch=2, turn_id=3)
|
||||
next_fence = DuplexFence("sid-cancel", epoch=3, turn_id=3)
|
||||
|
||||
assert await handler._signal_runtime_session(
|
||||
session,
|
||||
"input.cancel",
|
||||
fence=cancelled_fence,
|
||||
next_fence=next_fence,
|
||||
)
|
||||
|
||||
assert engine.calls == [("signal", cancelled_fence)]
|
||||
assert engine.signal_next_fences == [next_fence]
|
||||
@@ -108,6 +108,8 @@ class FakeAsyncOmniEngine:
|
||||
self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)]
|
||||
self.input_processor = None
|
||||
self.endpoint_restrictions = ()
|
||||
self.duplex_session_config = None
|
||||
self.duplex_serving_adapter_path = None
|
||||
|
||||
self.output_q: queue.Queue[Any] = queue.Queue()
|
||||
self.submitted: list[dict[str, Any]] = []
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
import base64
|
||||
import importlib.util
|
||||
import sys
|
||||
import wave
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
DEMO_PATH = Path(__file__).resolve().parents[2] / "examples/online_serving/minicpmo/realtime_duplex_demo.py"
|
||||
|
||||
|
||||
def _load_demo_module():
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"minicpmo_realtime_duplex_simple_demo_test",
|
||||
DEMO_PATH,
|
||||
)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_post_commit_decision_ignores_streaming_listens_before_commit_ack():
|
||||
demo = _load_demo_module()
|
||||
events = [
|
||||
{"type": "session.created"},
|
||||
{"type": "response.listen"},
|
||||
{"type": "response.listen"},
|
||||
{"type": "input_audio_buffer.committed"},
|
||||
]
|
||||
|
||||
commit_index = demo._input_committed_index(events, after_index=0)
|
||||
|
||||
assert commit_index == 3
|
||||
assert demo._post_commit_model_decision(events, commit_index) is None
|
||||
|
||||
|
||||
def test_post_commit_decision_accepts_listen_after_commit_ack():
|
||||
demo = _load_demo_module()
|
||||
events = [
|
||||
{"type": "response.listen"},
|
||||
{"type": "input_audio_buffer.committed"},
|
||||
{"type": "response.listen"},
|
||||
]
|
||||
|
||||
commit_index = demo._input_committed_index(events, after_index=0)
|
||||
|
||||
assert demo._post_commit_model_decision(events, commit_index) == "listen"
|
||||
|
||||
|
||||
def test_post_commit_decision_accepts_drained_speak_after_commit_ack():
|
||||
demo = _load_demo_module()
|
||||
events = [
|
||||
{"type": "input_audio_buffer.committed"},
|
||||
{"type": "response.created", "response": {"id": "resp-a"}},
|
||||
{"type": "response.audio.delta", "response_id": "resp-a", "delta": "AAAA"},
|
||||
{"type": "response.done", "response_id": "resp-a"},
|
||||
]
|
||||
|
||||
commit_index = demo._input_committed_index(events, after_index=0)
|
||||
|
||||
assert demo._post_commit_model_decision(events, commit_index) == "speak"
|
||||
|
||||
|
||||
def test_post_commit_decision_ignores_cancelled_response():
|
||||
demo = _load_demo_module()
|
||||
events = [
|
||||
{"type": "input_audio_buffer.committed"},
|
||||
{"type": "response.done", "response": {"status": "cancelled"}},
|
||||
]
|
||||
|
||||
commit_index = demo._input_committed_index(events, after_index=0)
|
||||
|
||||
assert demo._post_commit_model_decision(events, commit_index) is None
|
||||
|
||||
|
||||
def test_commit_ack_search_starts_after_stream_cursor():
|
||||
demo = _load_demo_module()
|
||||
events = [
|
||||
{"type": "input_audio_buffer.committed"},
|
||||
{"type": "response.listen"},
|
||||
{"type": "input_audio_buffer.committed"},
|
||||
]
|
||||
|
||||
assert demo._input_committed_index(events, after_index=1) == 2
|
||||
|
||||
|
||||
def test_exact_model_unit_boundary_does_not_require_an_extra_decision():
|
||||
demo = _load_demo_module()
|
||||
unit_bytes = demo.PCM16_SAMPLE_RATE * demo.PCM16_BYTES_PER_SAMPLE
|
||||
|
||||
assert not demo._has_residual_model_unit(b"\0" * (unit_bytes * 6), chunk_period_ms=1000)
|
||||
|
||||
|
||||
def test_partial_model_unit_requires_a_post_commit_decision():
|
||||
demo = _load_demo_module()
|
||||
unit_bytes = demo.PCM16_SAMPLE_RATE * demo.PCM16_BYTES_PER_SAMPLE
|
||||
|
||||
assert demo._has_residual_model_unit(b"\0" * (unit_bytes * 6 + 512), chunk_period_ms=1000)
|
||||
|
||||
|
||||
def test_latest_streaming_decision_is_used_at_exact_boundary():
|
||||
demo = _load_demo_module()
|
||||
events = [
|
||||
{"type": "response.listen"},
|
||||
{"type": "response.listen"},
|
||||
{"type": "input_audio_buffer.committed"},
|
||||
]
|
||||
|
||||
assert demo._latest_model_decision(events, after_index=0) == "listen"
|
||||
|
||||
|
||||
def test_ref_audio_data_url_encodes_explicit_wav(tmp_path):
|
||||
demo = _load_demo_module()
|
||||
ref = tmp_path / "ref.wav"
|
||||
ref.write_bytes(b"RIFFfake")
|
||||
|
||||
assert demo._ref_audio_data_url(str(ref)) == "data:audio/wav;base64,UklGRmZha2U="
|
||||
assert demo._ref_audio_data_url(None) is None
|
||||
|
||||
|
||||
def test_realtime_duplex_demo_requires_ref_audio(monkeypatch):
|
||||
demo = _load_demo_module()
|
||||
monkeypatch.setattr(
|
||||
demo.sys,
|
||||
"argv",
|
||||
[
|
||||
"realtime_duplex_demo.py",
|
||||
"--input-wav",
|
||||
"input.wav",
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(SystemExit):
|
||||
demo.parse_args()
|
||||
|
||||
|
||||
def test_realtime_duplex_demo_accepts_explicit_ref_audio(monkeypatch):
|
||||
demo = _load_demo_module()
|
||||
monkeypatch.setattr(
|
||||
demo.sys,
|
||||
"argv",
|
||||
[
|
||||
"realtime_duplex_demo.py",
|
||||
"--input-wav",
|
||||
"input.wav",
|
||||
"--ref-audio",
|
||||
"ref.wav",
|
||||
],
|
||||
)
|
||||
|
||||
args = demo.parse_args()
|
||||
|
||||
assert args.ref_audio == "ref.wav"
|
||||
|
||||
|
||||
def test_open_streaming_response_requires_post_commit_drain():
|
||||
demo = _load_demo_module()
|
||||
|
||||
assert demo._response_in_progress([{"type": "response.created"}])
|
||||
assert not demo._response_in_progress(
|
||||
[
|
||||
{"type": "response.created"},
|
||||
{"type": "response.done"},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_streaming_output_writer_persists_audio_deltas_as_they_arrive(tmp_path, capsys):
|
||||
demo = _load_demo_module()
|
||||
writer = demo._StreamingOutputWriter(tmp_path)
|
||||
first_pcm = b"\x01\x00\x02\x00"
|
||||
second_pcm = b"\x03\x00"
|
||||
|
||||
writer.handle(
|
||||
{
|
||||
"type": "response.audio.delta",
|
||||
"response_id": "resp-a",
|
||||
"delta": base64.b64encode(first_pcm).decode("ascii"),
|
||||
"sample_rate_hz": 24_000,
|
||||
}
|
||||
)
|
||||
writer.handle(
|
||||
{
|
||||
"type": "response.audio_transcript.delta",
|
||||
"response_id": "resp-a",
|
||||
"delta": "你好",
|
||||
}
|
||||
)
|
||||
writer.handle(
|
||||
{
|
||||
"type": "response.audio.delta",
|
||||
"response_id": "resp-a",
|
||||
"delta": base64.b64encode(second_pcm).decode("ascii"),
|
||||
"sample_rate_hz": 24_000,
|
||||
}
|
||||
)
|
||||
|
||||
assert (tmp_path / "output.pcm").read_bytes() == first_pcm + second_pcm
|
||||
chunk_paths = sorted((tmp_path / "audio_chunks").glob("*.wav"))
|
||||
assert [path.name for path in chunk_paths] == ["chunk_0001.wav", "chunk_0002.wav"]
|
||||
with wave.open(str(chunk_paths[0]), "rb") as chunk_wav:
|
||||
assert chunk_wav.getframerate() == 24_000
|
||||
assert chunk_wav.readframes(chunk_wav.getnframes()) == first_pcm
|
||||
assert writer.audio_chunk_paths == chunk_paths
|
||||
stderr = capsys.readouterr().err
|
||||
assert "audio chunk 1" in stderr
|
||||
assert "你好" in stderr
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,99 @@
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
SERVER_PATH = Path(__file__).resolve().parents[2] / "examples/online_serving/minicpmo/realtime_web/server.py"
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"minicpmo_realtime_web_server_test",
|
||||
SERVER_PATH,
|
||||
)
|
||||
assert spec is not None and spec.loader is not None
|
||||
server = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = server
|
||||
spec.loader.exec_module(server)
|
||||
|
||||
APP_DIR = server.APP_DIR
|
||||
STATIC_DIR = server.STATIC_DIR
|
||||
_join_ws_url = server._join_ws_url
|
||||
build_app = server.build_app
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_join_ws_url_preserves_realtime_query():
|
||||
assert (
|
||||
_join_ws_url(
|
||||
"ws://127.0.0.1:8099/",
|
||||
"/v1/realtime",
|
||||
"duplex=1&model=openbmb%2FMiniCPM-o-4_5",
|
||||
)
|
||||
== "ws://127.0.0.1:8099/v1/realtime?duplex=1&model=openbmb%2FMiniCPM-o-4_5"
|
||||
)
|
||||
|
||||
|
||||
def test_build_app_serves_health_and_injected_client_config():
|
||||
app = build_app(
|
||||
ws_backend="ws://127.0.0.1:9001",
|
||||
model="local/MiniCPM-o-4_5",
|
||||
)
|
||||
client = TestClient(app)
|
||||
|
||||
health = client.get("/healthz")
|
||||
assert health.status_code == 200
|
||||
assert health.text == "ok"
|
||||
|
||||
index = client.get("/")
|
||||
assert index.status_code == 200
|
||||
assert '"model": "local/MiniCPM-o-4_5"' in index.text
|
||||
assert '"realtimePath": "v1/realtime"' in index.text
|
||||
assert "__FULL_DUPLEX_CONFIG__" not in index.text
|
||||
|
||||
|
||||
def test_build_app_injects_public_realtime_url():
|
||||
public_realtime_url = "wss://proxy.example.test/backend/v1/realtime"
|
||||
app = build_app(
|
||||
ws_backend="ws://127.0.0.1:9001",
|
||||
public_realtime_url=public_realtime_url,
|
||||
)
|
||||
client = TestClient(app)
|
||||
|
||||
index = client.get("/")
|
||||
|
||||
assert index.status_code == 200
|
||||
assert f'"realtimePath": "{public_realtime_url}"' in index.text
|
||||
|
||||
|
||||
def test_build_app_versions_client_bundle_and_disables_index_cache():
|
||||
client = TestClient(build_app())
|
||||
expected_version_hash = hashlib.sha256()
|
||||
for asset_name in ("app.js", "pcm_worklet.js", "playback_worklet.js"):
|
||||
expected_version_hash.update((STATIC_DIR / asset_name).read_bytes())
|
||||
expected_version = expected_version_hash.hexdigest()[:12]
|
||||
|
||||
assert (APP_DIR / "index.html").is_file()
|
||||
assert (STATIC_DIR / "app.js").is_file()
|
||||
assert not (APP_DIR.parent / "static").exists()
|
||||
assert not (STATIC_DIR / "index.html").exists()
|
||||
|
||||
index = client.get("/")
|
||||
|
||||
assert index.status_code == 200
|
||||
assert f'src="static/app.js?v={expected_version}"' in index.text
|
||||
assert index.headers["cache-control"] == "no-store"
|
||||
|
||||
|
||||
def test_build_app_exposes_realtime_websocket_and_static_assets():
|
||||
app = build_app()
|
||||
paths = {route.path for route in app.routes}
|
||||
|
||||
assert "/v1/realtime" in paths
|
||||
assert "/static" in paths
|
||||
|
||||
client = TestClient(app)
|
||||
assert client.get("/static/app.js").status_code == 200
|
||||
assert client.get("/static/pcm_worklet.js").status_code == 200
|
||||
assert client.get("/static/playback_worklet.js").status_code == 200
|
||||
@@ -0,0 +1,209 @@
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import textwrap
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2] / "examples" / "online_serving" / "minicpmo" / "realtime_web"
|
||||
APP_ROOT = ROOT / "app"
|
||||
STATIC_ROOT = APP_ROOT / "static"
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_page_exposes_focused_call_conversation_and_log_surfaces():
|
||||
html = (APP_ROOT / "index.html").read_text(encoding="utf-8")
|
||||
|
||||
assert 'id="callButton"' in html
|
||||
assert 'id="muteButton"' in html
|
||||
assert 'id="connectionState"' in html
|
||||
assert 'id="modelState"' in html
|
||||
assert 'id="conversation"' in html
|
||||
assert 'id="eventLog"' in html
|
||||
assert "<details" in html
|
||||
assert "Automatic barge-in" not in html
|
||||
assert "Server VAD" not in html
|
||||
|
||||
|
||||
def test_client_uses_proxy_relative_realtime_url_and_model_policy_session():
|
||||
source = (STATIC_ROOT / "app.js").read_text(encoding="utf-8")
|
||||
|
||||
assert "new URL(config.realtimePath, window.location.href)" in source
|
||||
assert "url.protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'" in source
|
||||
assert "url.searchParams.set('autostart', '0')" in source
|
||||
assert "url.searchParams.set('minicpmo45_native_duplex', '1')" in source
|
||||
assert "auto_response: true" in source
|
||||
assert "input_audio_buffer.append" in source
|
||||
assert "input_audio_buffer.commit" not in source
|
||||
assert "playback.ack" in source
|
||||
assert "event.event || event" in source
|
||||
assert "response.audio.delta" in source
|
||||
assert "response.audio_transcript.delta" in source
|
||||
assert "conversation.item.input_audio_transcription" in source
|
||||
assert "force_barge_in" not in source
|
||||
assert "server_vad" not in source
|
||||
assert "type: 'response.create'" not in source
|
||||
assert 'type: "response.create"' not in source
|
||||
|
||||
|
||||
def test_web_server_requires_ref_audio_for_audio_output_session():
|
||||
source = (ROOT / "server.py").read_text(encoding="utf-8")
|
||||
|
||||
ref_audio_arg = re.search(
|
||||
r"parser\.add_argument\(\s*\"--ref-audio\",(?P<body>.*?)\n\s*\)",
|
||||
source,
|
||||
re.DOTALL,
|
||||
)
|
||||
assert ref_audio_arg is not None
|
||||
assert "required=True" in ref_audio_arg.group("body")
|
||||
assert "Optional reference voice" not in ref_audio_arg.group("body")
|
||||
|
||||
|
||||
def test_client_sends_ref_audio_in_realtime_session_contract():
|
||||
source = (STATIC_ROOT / "app.js").read_text(encoding="utf-8")
|
||||
|
||||
assert "if (config.refAudio) session.ref_audio = config.refAudio;" in source
|
||||
assert "extraBody.ref_audio" not in source
|
||||
|
||||
|
||||
def test_client_has_transactional_cleanup_and_visible_event_logging():
|
||||
source = (STATIC_ROOT / "app.js").read_text(encoding="utf-8")
|
||||
|
||||
assert "async function stopSession" in source
|
||||
assert "type: 'session.close'" in source
|
||||
assert "waitForSessionClosed" in source
|
||||
assert "SESSION_CLOSE_TIMEOUT_MS" in source
|
||||
assert "case 'session.closed':" in source
|
||||
assert "track.stop()" in source
|
||||
assert "clearInterval(sendTimer)" in source
|
||||
assert "appendEventLog(event)" in source
|
||||
assert "stopSession({ terminal: false })" in source
|
||||
|
||||
|
||||
def test_client_keeps_microphone_upload_active_during_assistant_playback():
|
||||
source = (STATIC_ROOT / "app.js").read_text(encoding="utf-8")
|
||||
upload_gate = re.search(r"function microphoneUploadEnabled\(\) \{(?P<body>.*?)\n \}", source, re.DOTALL)
|
||||
begin_assistant = re.search(r"function beginAssistant\(responseId\) \{(?P<body>.*?)\n \}", source, re.DOTALL)
|
||||
|
||||
assert upload_gate is not None
|
||||
assert "return running && !muted;" in upload_gate.group("body")
|
||||
assert "assistantActive" not in upload_gate.group("body")
|
||||
assert begin_assistant is not None
|
||||
assert "pendingCapture = []" not in begin_assistant.group("body")
|
||||
|
||||
|
||||
def test_audio_worklets_define_capture_and_playback_processors():
|
||||
capture = (STATIC_ROOT / "pcm_worklet.js").read_text(encoding="utf-8")
|
||||
playback = (STATIC_ROOT / "playback_worklet.js").read_text(encoding="utf-8")
|
||||
|
||||
assert "registerProcessor('fullduplex-pcm-capture'" in capture
|
||||
assert "Int16Array" in capture
|
||||
assert "registerProcessor('fullduplex-pcm-playback'" in playback
|
||||
assert "playback-drained" in playback
|
||||
assert "clear" in playback
|
||||
|
||||
|
||||
def test_audio_worklet_urls_use_the_static_asset_version():
|
||||
server = (ROOT / "server.py").read_text(encoding="utf-8")
|
||||
app = (STATIC_ROOT / "app.js").read_text(encoding="utf-8")
|
||||
|
||||
assert '"appVersion": app_version' in server
|
||||
assert 'STATIC_DIR / "playback_worklet.js"' in server
|
||||
assert 'STATIC_DIR / "pcm_worklet.js"' in server
|
||||
assert "staticAssetUrl('static/playback_worklet.js')" in app
|
||||
assert "staticAssetUrl('static/pcm_worklet.js')" in app
|
||||
|
||||
|
||||
def test_playback_worklet_buffers_first_400ms_and_reports_underruns():
|
||||
app = (STATIC_ROOT / "app.js").read_text(encoding="utf-8")
|
||||
playback = (STATIC_ROOT / "playback_worklet.js").read_text(encoding="utf-8")
|
||||
|
||||
assert "INITIAL_PLAYBACK_BUFFER_MS = 400" in app
|
||||
assert "initialBufferMs" in app
|
||||
assert "responseId" in app
|
||||
assert "playback-underrun" in app
|
||||
assert "underrunMs" in app
|
||||
assert "initialBufferFrames" in playback
|
||||
assert "playback-underrun" in playback
|
||||
assert "underrunFrames" in playback
|
||||
assert "underrunMs" in playback
|
||||
|
||||
|
||||
def test_playback_worklet_waits_before_playing_and_rebuffers_after_underrun():
|
||||
node = shutil.which("node")
|
||||
if node is None:
|
||||
pytest.skip("node is required for the AudioWorklet regression test")
|
||||
|
||||
script = textwrap.dedent(
|
||||
"""
|
||||
const fs = require('fs');
|
||||
const vm = require('vm');
|
||||
|
||||
global.sampleRate = 1000;
|
||||
global.AudioWorkletProcessor = class {
|
||||
constructor() {
|
||||
this.port = { onmessage: null, postMessage: () => {} };
|
||||
}
|
||||
};
|
||||
let Processor = null;
|
||||
global.registerProcessor = (_name, processor) => { Processor = processor; };
|
||||
vm.runInThisContext(fs.readFileSync(process.argv[1], 'utf8'));
|
||||
|
||||
const processor = new Processor();
|
||||
const render = () => {
|
||||
const output = new Float32Array(100);
|
||||
processor.process([], [[output]]);
|
||||
return output;
|
||||
};
|
||||
const assert = (condition, message) => {
|
||||
if (!condition) throw new Error(message);
|
||||
};
|
||||
|
||||
const first = new Int16Array(150);
|
||||
first.fill(16384);
|
||||
processor.handleMessage({
|
||||
type: 'audio',
|
||||
pcm: first,
|
||||
responseId: 'response-1',
|
||||
initialBufferMs: 400,
|
||||
});
|
||||
assert(!processor.started, 'large first delta must not bypass wall-clock prebuffer');
|
||||
assert(render().every((sample) => sample === 0), 'first prebuffer render must be silent');
|
||||
assert(render().every((sample) => sample === 0), 'second prebuffer render must be silent');
|
||||
assert(render().every((sample) => sample === 0), 'third prebuffer render must be silent');
|
||||
assert(render().every((sample) => sample === 0), 'fourth prebuffer render must be silent');
|
||||
const firstPlayback = render();
|
||||
assert(firstPlayback.some((sample) => sample !== 0), 'playback must start after prebuffer');
|
||||
|
||||
const underrun = render();
|
||||
assert(!processor.started, 'an empty queue must return to buffering');
|
||||
assert(underrun[underrun.length - 1] === 0, 'underrun boundary must fade to zero');
|
||||
|
||||
const resumed = new Int16Array(300);
|
||||
resumed.fill(8192);
|
||||
processor.handleMessage({
|
||||
type: 'audio',
|
||||
pcm: resumed,
|
||||
responseId: 'response-1',
|
||||
initialBufferMs: 400,
|
||||
});
|
||||
assert(render().every((sample) => sample === 0), 'resume must rebuild jitter buffer');
|
||||
assert(render().every((sample) => sample === 0), 'resume must keep rebuilding jitter buffer');
|
||||
assert(render().every((sample) => sample === 0), 'resume must keep waiting');
|
||||
assert(render().every((sample) => sample === 0), 'resume must wait the full buffer interval');
|
||||
const resumedPlayback = render();
|
||||
assert(resumedPlayback.some((sample) => sample !== 0), 'playback must resume after rebuffer');
|
||||
assert(
|
||||
Math.abs(resumedPlayback[0]) < Math.abs(resumedPlayback[50]),
|
||||
'resumed playback must fade in instead of jumping from silence',
|
||||
);
|
||||
"""
|
||||
)
|
||||
subprocess.run(
|
||||
[node, "-e", script, str(STATIC_ROOT / "playback_worklet.js")],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
@@ -0,0 +1,110 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_control_client import (
|
||||
DuplexControlClient,
|
||||
DuplexControlRequestError,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import (
|
||||
DuplexControlError,
|
||||
DuplexControlResultMessage,
|
||||
DuplexFence,
|
||||
ResumeDuplexSessionMessage,
|
||||
TouchDuplexSessionMessage,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _Transport:
|
||||
def __init__(self) -> None:
|
||||
self.calls = []
|
||||
|
||||
def execute(self, key, message, **kwargs):
|
||||
self.calls.append((key, message, kwargs))
|
||||
return DuplexControlResultMessage(
|
||||
control_id=message.control_id,
|
||||
fence=message.fence,
|
||||
operation="touch" if isinstance(message, TouchDuplexSessionMessage) else "resume",
|
||||
session_id=message.session_id,
|
||||
ok=True,
|
||||
stage_results=[],
|
||||
)
|
||||
|
||||
|
||||
class _UnsupportedTransport:
|
||||
def execute(self, key, message, **kwargs):
|
||||
del key, kwargs
|
||||
return DuplexControlResultMessage(
|
||||
control_id=message.control_id,
|
||||
fence=message.fence,
|
||||
operation="touch",
|
||||
session_id=message.session_id,
|
||||
ok=False,
|
||||
stage_results=[{"result": {"supported": False}}],
|
||||
unsupported_count=1,
|
||||
error=DuplexControlError(code="invalid_capability", message="unsupported", retryable=False),
|
||||
accepted_fence=DuplexFence(message.session_id, epoch=2),
|
||||
lease_generation=7,
|
||||
)
|
||||
|
||||
|
||||
def test_control_client_routes_touch_and_resume_by_control_id() -> None:
|
||||
transport = _Transport()
|
||||
control_ids = iter(("touch-id", "resume-id"))
|
||||
client = DuplexControlClient(transport, control_id_factory=lambda: next(control_ids))
|
||||
fence = DuplexFence("sid-client")
|
||||
|
||||
assert (
|
||||
client.touch(
|
||||
fence.session_id,
|
||||
fence=fence,
|
||||
activity=DuplexLeaseActivity.PLAYBACK_ACK,
|
||||
timeout=2.0,
|
||||
)["ok"]
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
client.resume(
|
||||
fence.session_id,
|
||||
fence=fence,
|
||||
expected_lease_generation=7,
|
||||
timeout=3.0,
|
||||
)["ok"]
|
||||
is True
|
||||
)
|
||||
|
||||
touch_key, touch_message, touch_kwargs = transport.calls[0]
|
||||
assert touch_key == ("duplex", "touch-id")
|
||||
assert isinstance(touch_message, TouchDuplexSessionMessage)
|
||||
assert touch_message.activity == DuplexLeaseActivity.PLAYBACK_ACK.value
|
||||
assert touch_kwargs["timeout"] == 2.0
|
||||
|
||||
resume_key, resume_message, resume_kwargs = transport.calls[1]
|
||||
assert resume_key == ("duplex", "resume-id")
|
||||
assert isinstance(resume_message, ResumeDuplexSessionMessage)
|
||||
assert resume_message.expected_lease_generation == 7
|
||||
assert resume_kwargs["timeout"] == 3.0
|
||||
|
||||
|
||||
def test_control_client_treats_ok_false_as_authoritative_failure() -> None:
|
||||
client = DuplexControlClient(_UnsupportedTransport(), control_id_factory=lambda: "unsupported")
|
||||
fence = DuplexFence("sid-unsupported")
|
||||
|
||||
with pytest.raises(DuplexControlRequestError, match="duplex touch failed") as exc_info:
|
||||
client.touch(
|
||||
fence.session_id,
|
||||
fence=fence,
|
||||
activity=DuplexLeaseActivity.HEARTBEAT,
|
||||
timeout=1.0,
|
||||
)
|
||||
|
||||
assert exc_info.value.code == "invalid_capability"
|
||||
assert exc_info.value.retryable is False
|
||||
assert exc_info.value.accepted_fence == DuplexFence(fence.session_id, epoch=2)
|
||||
assert exc_info.value.lease_generation == 7
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,68 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Focused tests for typed deployment-only runtime configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import FrozenInstanceError
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.stage_config import (
|
||||
DuplexSessionRuntimeConfig,
|
||||
load_deploy_config,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_duplex_session_runtime_defaults_are_typed_and_immutable(tmp_path) -> None:
|
||||
deploy_path = tmp_path / "duplex.yaml"
|
||||
deploy_path.write_text("duplex_session: {}\nstages: []\n", encoding="utf-8")
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
|
||||
assert isinstance(deploy.duplex_session, DuplexSessionRuntimeConfig)
|
||||
assert deploy.duplex_session.idle_ttl_s == 300.0
|
||||
assert deploy.duplex_session.disconnect_grace_s == 30.0
|
||||
assert deploy.duplex_session.reaper_interval_s == 5.0
|
||||
assert deploy.duplex_session.resume_replay_ttl_s == 60.0
|
||||
assert deploy.duplex_session.resume_replay_max_bytes_per_session == 8 * 1024 * 1024
|
||||
assert deploy.duplex_session.max_pending_input_bytes_per_session == 16 * 1024 * 1024
|
||||
assert deploy.duplex_session.max_pending_turns_per_session == 4
|
||||
assert deploy.duplex_session.max_sessions == 1
|
||||
assert deploy.duplex_session.completed_append_cache_size == 256
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
deploy.duplex_session.idle_ttl_s = 1.0 # type: ignore[misc]
|
||||
|
||||
|
||||
def test_duplex_session_runtime_accepts_disabled_idle_expiry(tmp_path) -> None:
|
||||
deploy_path = tmp_path / "duplex.yaml"
|
||||
deploy_path.write_text("duplex_session:\n idle_ttl_s: null\nstages: []\n", encoding="utf-8")
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
|
||||
assert deploy.duplex_session.idle_ttl_s is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("name", "value"),
|
||||
[
|
||||
("idle_ttl_s", 0),
|
||||
("idle_ttl_s", -1),
|
||||
("disconnect_grace_s", 0),
|
||||
("reaper_interval_s", -1),
|
||||
("resume_replay_ttl_s", 0),
|
||||
("resume_replay_max_bytes_per_session", -1),
|
||||
("max_pending_input_bytes_per_session", 0),
|
||||
("max_pending_turns_per_session", -1),
|
||||
("max_sessions", 0),
|
||||
("completed_append_cache_size", 0),
|
||||
],
|
||||
)
|
||||
def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str, value: int) -> None:
|
||||
deploy_path = tmp_path / "duplex.yaml"
|
||||
deploy_path.write_text(f"duplex_session:\n {name}: {value}\nstages: []\n", encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match=rf"duplex_session\.{name} must be positive"):
|
||||
load_deploy_config(deploy_path)
|
||||
@@ -0,0 +1,14 @@
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.intermediate import get_stream_request_key
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_get_stream_request_key_requires_stable_identifier():
|
||||
with pytest.raises(ValueError, match="stable request id"):
|
||||
get_stream_request_key({"ids": {"tts": [1, 2, 3]}})
|
||||
|
||||
|
||||
def test_get_stream_request_key_accepts_global_request_id():
|
||||
assert get_stream_request_key({"global_request_id": ["duplex-sid-stage0"]}) == "duplex-sid-stage0"
|
||||
@@ -0,0 +1,224 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import FrozenInstanceError
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_runtime import DuplexInputMode
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_session import DuplexSessionRuntimeManager
|
||||
from vllm_omni.experimental.fullduplex.engine.lease import (
|
||||
DuplexLeaseActivity,
|
||||
DuplexLeaseConfig,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeMonotonicClock:
|
||||
def __init__(self, value: float) -> None:
|
||||
self.value = value
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.value
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.value += seconds
|
||||
|
||||
|
||||
def _lease_config(*, idle_ttl_s: float | None = 300.0) -> DuplexLeaseConfig:
|
||||
return DuplexLeaseConfig(
|
||||
idle_ttl_s=idle_ttl_s,
|
||||
disconnect_grace_s=30.0,
|
||||
)
|
||||
|
||||
|
||||
def test_open_touch_detach_and_idle_expiry_use_monotonic_time() -> None:
|
||||
clock = FakeMonotonicClock(100.0)
|
||||
manager = DuplexSessionRuntimeManager(clock=clock)
|
||||
fence = DuplexFence("sid-expiry")
|
||||
session = manager.open_session(fence, lease_config=_lease_config())
|
||||
|
||||
assert session.lease.last_activity == 100.0
|
||||
clock.advance(10.0)
|
||||
session.touch(fence, DuplexLeaseActivity.HEARTBEAT)
|
||||
session.detach(fence)
|
||||
clock.advance(29.0)
|
||||
assert session.lease.disconnect_grace_expired(clock()) is False
|
||||
assert manager.collect_expired() == []
|
||||
|
||||
clock.advance(272.0)
|
||||
expired = manager.collect_expired()
|
||||
|
||||
assert [item.session_id for item in expired] == ["sid-expiry"]
|
||||
assert expired[0].reason == "disconnect_grace_expired"
|
||||
assert manager.get("sid-expiry") is session
|
||||
manager.finalize_close_session(session)
|
||||
assert manager.get("sid-expiry") is None
|
||||
|
||||
|
||||
def test_detach_grace_and_resume_advance_token_independent_lease_generation() -> None:
|
||||
clock = FakeMonotonicClock(10.0)
|
||||
manager = DuplexSessionRuntimeManager(clock=clock)
|
||||
fence = DuplexFence("sid-resume")
|
||||
session = manager.open_session(fence, lease_config=_lease_config())
|
||||
|
||||
session.detach(fence)
|
||||
clock.advance(29.0)
|
||||
assert session.lease.disconnect_grace_expired(clock()) is False
|
||||
clock.advance(2.0)
|
||||
assert session.lease.disconnect_grace_expired(clock()) is True
|
||||
|
||||
generation = session.resume(fence, expected_lease_generation=0)
|
||||
|
||||
assert generation == 1
|
||||
assert session.lease.detached_at is None
|
||||
assert session.lease.disconnect_grace_expired(clock()) is False
|
||||
with pytest.raises(ValueError, match="lease generation mismatch"):
|
||||
session.resume(fence, expected_lease_generation=0)
|
||||
|
||||
|
||||
def test_active_operation_prevents_mid_transaction_expiry() -> None:
|
||||
clock = FakeMonotonicClock(0.0)
|
||||
manager = DuplexSessionRuntimeManager(clock=clock)
|
||||
fence = DuplexFence("sid-operation")
|
||||
session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=5.0))
|
||||
|
||||
session.begin_operation(fence, "append-1")
|
||||
clock.advance(10.0)
|
||||
assert manager.collect_expired() == []
|
||||
|
||||
session.end_operation(fence, "append-1")
|
||||
assert manager.collect_expired() == []
|
||||
clock.advance(6.0)
|
||||
assert [item.session_id for item in manager.collect_expired()] == ["sid-operation"]
|
||||
|
||||
|
||||
def test_sessions_have_independent_activity_deadlines_and_resources() -> None:
|
||||
clock = FakeMonotonicClock(0.0)
|
||||
manager = DuplexSessionRuntimeManager(clock=clock)
|
||||
fence_a = DuplexFence("sid-a")
|
||||
fence_b = DuplexFence("sid-b")
|
||||
session_a = manager.open_session(fence_a, lease_config=_lease_config(idle_ttl_s=10.0))
|
||||
session_b = manager.open_session(fence_b, lease_config=_lease_config(idle_ttl_s=10.0))
|
||||
session_a.reserve_stage_request(0, "req-a-reserved", fence=fence_a)
|
||||
session_a.bind_stage_request(1, "req-a-submitted", fence=fence_a)
|
||||
session_b.bind_stage_request(0, "req-b", fence=fence_b)
|
||||
|
||||
clock.advance(6.0)
|
||||
session_b.touch(fence_b, DuplexLeaseActivity.MODEL_OUTPUT)
|
||||
clock.advance(5.0)
|
||||
expired = manager.collect_expired()
|
||||
|
||||
assert len(expired) == 1
|
||||
assert expired[0].session_id == "sid-a"
|
||||
assert expired[0].reserved_request_ids == ("req-a-reserved",)
|
||||
assert expired[0].submitted_request_ids == ("req-a-submitted",)
|
||||
assert manager.get("sid-b") is session_b
|
||||
assert session_b.resource_request_ids() == ["req-b"]
|
||||
|
||||
|
||||
def test_close_and_reaper_select_exactly_one_terminal_transition() -> None:
|
||||
clock = FakeMonotonicClock(0.0)
|
||||
manager = DuplexSessionRuntimeManager(clock=clock)
|
||||
fence = DuplexFence("sid-race")
|
||||
manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=1.0))
|
||||
clock.advance(2.0)
|
||||
|
||||
assert len(manager.collect_expired()) == 1
|
||||
assert manager.close_session(fence, reason="explicit_close") is None
|
||||
assert manager.collect_expired() == []
|
||||
|
||||
fence_2 = DuplexFence("sid-race-2")
|
||||
manager.open_session(fence_2, lease_config=_lease_config(idle_ttl_s=1.0))
|
||||
assert manager.close_session(fence_2, reason="explicit_close") is not None
|
||||
clock.advance(2.0)
|
||||
assert manager.collect_expired() == []
|
||||
|
||||
|
||||
def test_stale_fence_cannot_touch_detach_or_resume_lease() -> None:
|
||||
clock = FakeMonotonicClock(0.0)
|
||||
manager = DuplexSessionRuntimeManager(clock=clock)
|
||||
current = DuplexFence("sid-fence", epoch=1)
|
||||
stale = DuplexFence("sid-fence", epoch=0)
|
||||
session = manager.open_session(current, lease_config=_lease_config())
|
||||
|
||||
with pytest.raises(RuntimeError, match="fence mismatch"):
|
||||
session.touch(stale, DuplexLeaseActivity.APPEND)
|
||||
with pytest.raises(RuntimeError, match="fence mismatch"):
|
||||
session.detach(stale)
|
||||
with pytest.raises(RuntimeError, match="fence mismatch"):
|
||||
session.resume(stale, expected_lease_generation=0)
|
||||
|
||||
|
||||
def test_disabled_idle_expiry_never_collects_session() -> None:
|
||||
clock = FakeMonotonicClock(0.0)
|
||||
manager = DuplexSessionRuntimeManager(clock=clock)
|
||||
manager.open_session(DuplexFence("sid-no-expiry"), lease_config=_lease_config(idle_ttl_s=None))
|
||||
|
||||
clock.advance(1_000_000.0)
|
||||
|
||||
assert manager.collect_expired() == []
|
||||
|
||||
|
||||
def test_detached_session_expires_at_disconnect_grace_when_idle_ttl_is_disabled() -> None:
|
||||
clock = FakeMonotonicClock(0.0)
|
||||
manager = DuplexSessionRuntimeManager(clock=clock)
|
||||
fence = DuplexFence("sid-disconnect-grace")
|
||||
session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=None))
|
||||
|
||||
session.detach(fence)
|
||||
clock.advance(29.0)
|
||||
assert manager.collect_expired() == []
|
||||
|
||||
clock.advance(1.0)
|
||||
expired = manager.collect_expired()
|
||||
|
||||
assert [item.session_id for item in expired] == [fence.session_id]
|
||||
assert expired[0].reason == "disconnect_grace_expired"
|
||||
|
||||
|
||||
def test_runtime_manager_enforces_server_owned_session_admission() -> None:
|
||||
manager = DuplexSessionRuntimeManager(max_sessions=2)
|
||||
manager.open_session(DuplexFence("sid-admission-a"))
|
||||
manager.open_session(DuplexFence("sid-admission-b"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="duplex_session_capacity_exhausted"):
|
||||
manager.open_session(DuplexFence("sid-admission-c"))
|
||||
|
||||
|
||||
def test_completed_append_cache_is_bounded() -> None:
|
||||
manager = DuplexSessionRuntimeManager(completed_append_limit=2)
|
||||
fence = DuplexFence("sid-completed-cache")
|
||||
session = manager.open_session(fence)
|
||||
|
||||
for index in range(3):
|
||||
session.record_completed_append(
|
||||
f"operation-{index}",
|
||||
fence=fence,
|
||||
mode=DuplexInputMode.TURN_COMMIT_ONLY,
|
||||
final=True,
|
||||
stage_results=[{"index": index}],
|
||||
)
|
||||
|
||||
assert list(session.completed_appends) == ["operation-1", "operation-2"]
|
||||
|
||||
session.accept_fence(DuplexFence(fence.session_id, epoch=1))
|
||||
assert session.completed_appends == {}
|
||||
|
||||
|
||||
def test_lease_config_and_expiry_record_are_immutable() -> None:
|
||||
config = _lease_config()
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
config.idle_ttl_s = 1.0 # type: ignore[misc]
|
||||
|
||||
clock = FakeMonotonicClock(0.0)
|
||||
manager = DuplexSessionRuntimeManager(clock=clock)
|
||||
manager.open_session(DuplexFence("sid-record"), lease_config=_lease_config(idle_ttl_s=1.0))
|
||||
clock.advance(2.0)
|
||||
record = manager.collect_expired()[0]
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
record.reason = "changed" # type: ignore[misc]
|
||||
@@ -0,0 +1,480 @@
|
||||
from dataclasses import FrozenInstanceError
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine import duplex_runtime
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import (
|
||||
duplex_data_plane_request_info,
|
||||
duplex_resource_request_belongs_to_session,
|
||||
duplex_resource_request_id,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
|
||||
DuplexInputMode,
|
||||
DuplexOutputAction,
|
||||
DuplexOutputDecision,
|
||||
DuplexRuntimeCapabilities,
|
||||
DuplexSessionRuntimeManager,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
from vllm_omni.experimental.fullduplex.minicpmo45.runtime import (
|
||||
MiniCPMO45DuplexRuntimeExtension,
|
||||
build_duplex_data_plane_prompt,
|
||||
duplex_scheduler_token_budget,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_duplex_fence_is_immutable():
|
||||
fence = DuplexFence("session")
|
||||
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
fence.epoch = 1 # type: ignore[misc]
|
||||
|
||||
assert not hasattr(fence, "__dict__")
|
||||
|
||||
|
||||
def test_duplex_runtime_tracks_stage_bindings_and_barge_in_epoch():
|
||||
manager = DuplexSessionRuntimeManager()
|
||||
session = manager.open_session(
|
||||
DuplexFence("sid-1"),
|
||||
capabilities=DuplexRuntimeCapabilities(
|
||||
input_modes={DuplexInputMode.APPEND_TOKENS},
|
||||
),
|
||||
)
|
||||
session.bind_stage_request(stage_id=0, request_id="req-stage0", fence=session.fence)
|
||||
session.bind_stage_request(stage_id=1, request_id="req-stage1", fence=session.fence)
|
||||
|
||||
update = session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence)
|
||||
next_fence = DuplexFence("sid-1", epoch=1)
|
||||
stale_request_ids = session.release_fence(session.fence)
|
||||
session.accept_fence(next_fence)
|
||||
|
||||
assert update.seq == 1
|
||||
assert session.fence == next_fence
|
||||
assert stale_request_ids == ["req-stage0", "req-stage1"]
|
||||
assert session.stage_bindings == {}
|
||||
|
||||
|
||||
def test_duplex_runtime_tracks_same_request_id_for_each_pipeline_stage():
|
||||
manager = DuplexSessionRuntimeManager()
|
||||
fence = DuplexFence("sid-shared-pipeline-request")
|
||||
session = manager.open_session(fence)
|
||||
|
||||
session.reserve_stage_request(0, "req-shared", fence=fence)
|
||||
session.bind_stage_request(0, "req-shared", fence=fence)
|
||||
session.bind_stage_request(1, "req-shared", fence=fence)
|
||||
|
||||
assert session.stage_bindings[0].request_id == "req-shared"
|
||||
assert session.stage_bindings[1].request_id == "req-shared"
|
||||
assert session.resource_request_ids() == ["req-shared"]
|
||||
assert session.input_seq == 0
|
||||
|
||||
|
||||
def test_duplex_runtime_extension_validation_rejects_missing_methods():
|
||||
class IncompleteExtension:
|
||||
def configure_sampling_params(self, *, runtime_config, defaults):
|
||||
del runtime_config
|
||||
return defaults
|
||||
|
||||
with pytest.raises(TypeError, match="plan_append"):
|
||||
duplex_runtime.validate_duplex_runtime_extension(IncompleteExtension())
|
||||
|
||||
|
||||
def test_duplex_runtime_extension_validation_rejects_stage_count_mismatch():
|
||||
class WrongStageCountExtension(MiniCPMO45DuplexRuntimeExtension):
|
||||
def configure_sampling_params(self, *, runtime_config, defaults):
|
||||
del runtime_config
|
||||
return defaults[:1]
|
||||
|
||||
with pytest.raises(ValueError, match="one sampling parameter per stage"):
|
||||
duplex_runtime.validate_duplex_runtime_extension(
|
||||
WrongStageCountExtension(),
|
||||
sampling_defaults=(SamplingParams(), SamplingParams()),
|
||||
)
|
||||
|
||||
|
||||
def test_duplex_runtime_extension_validation_rejects_sampling_type_mismatch():
|
||||
class WrongSamplingTypeExtension(MiniCPMO45DuplexRuntimeExtension):
|
||||
def configure_sampling_params(self, *, runtime_config, defaults):
|
||||
del runtime_config
|
||||
return tuple(object() for _ in defaults)
|
||||
|
||||
with pytest.raises(TypeError, match="sampling parameter type"):
|
||||
duplex_runtime.validate_duplex_runtime_extension(
|
||||
WrongSamplingTypeExtension(),
|
||||
sampling_defaults=(SamplingParams(), SamplingParams()),
|
||||
)
|
||||
|
||||
|
||||
def test_duplex_output_decision_metadata_is_immutable():
|
||||
decision = DuplexOutputDecision(
|
||||
action=DuplexOutputAction.DIRECT_RESPONSE,
|
||||
metadata={"model_listen": True},
|
||||
)
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
decision.metadata["model_listen"] = False
|
||||
|
||||
|
||||
def _decide_minicpmo_output(
|
||||
output: object,
|
||||
*,
|
||||
segment_token_ids: tuple[int, ...] = (),
|
||||
segment_output_metadata: dict | None = None,
|
||||
):
|
||||
return MiniCPMO45DuplexRuntimeExtension().decide_output(
|
||||
stage_id=0,
|
||||
final_stage_id=1,
|
||||
segment_finished=True,
|
||||
segment_token_ids=segment_token_ids,
|
||||
segment_output_metadata=segment_output_metadata or {},
|
||||
output=output,
|
||||
)
|
||||
|
||||
|
||||
def test_minicpmo_extension_owns_stage_sampling_overrides():
|
||||
defaults = (
|
||||
SamplingParams(max_tokens=4),
|
||||
SamplingParams(max_tokens=8),
|
||||
)
|
||||
|
||||
configured = MiniCPMO45DuplexRuntimeExtension().configure_sampling_params(
|
||||
runtime_config={
|
||||
"duplex_stage_max_tokens": {"0": 20},
|
||||
"duplex_stage_sampling_params": {"1": {"stop_token_ids": [151645]}},
|
||||
},
|
||||
defaults=defaults,
|
||||
)
|
||||
|
||||
assert configured[0].max_tokens == 20
|
||||
assert configured[1].stop_token_ids == [151645]
|
||||
assert defaults[0].max_tokens == 4
|
||||
assert 151645 not in (defaults[1].stop_token_ids or [])
|
||||
|
||||
|
||||
def test_minicpmo_output_decision_uses_raw_streaming_token_snapshot():
|
||||
decision = _decide_minicpmo_output(
|
||||
SimpleNamespace(outputs=[SimpleNamespace()]),
|
||||
segment_token_ids=(151705,),
|
||||
segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}},
|
||||
)
|
||||
|
||||
assert decision is not None
|
||||
assert decision.action is DuplexOutputAction.DIRECT_RESPONSE
|
||||
assert decision.metadata["duplex_native_decision"] == "listen"
|
||||
assert decision.metadata["model_listen"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"])
|
||||
def test_minicpmo_output_decision_ignores_output_level_token_history(attr):
|
||||
output = SimpleNamespace(
|
||||
multimodal_output={"special_token_ids": {"listen_token_id": 151705}},
|
||||
outputs=[SimpleNamespace()],
|
||||
**{attr: [42, 151705]},
|
||||
)
|
||||
|
||||
assert _decide_minicpmo_output(output) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"])
|
||||
def test_minicpmo_output_decision_uses_completion_token_ids(attr):
|
||||
output = SimpleNamespace(
|
||||
multimodal_output={"special_token_ids": {"listen_token_id": 151705}},
|
||||
outputs=[SimpleNamespace(**{attr: [42, 151705]})],
|
||||
)
|
||||
|
||||
assert _decide_minicpmo_output(output) is not None
|
||||
|
||||
|
||||
def test_minicpmo_output_decision_uses_completion_stop_reason():
|
||||
output = SimpleNamespace(
|
||||
multimodal_output={"special_token_ids": {"listen_token_id": 151705}},
|
||||
outputs=[SimpleNamespace(stop_reason=151705)],
|
||||
)
|
||||
|
||||
assert _decide_minicpmo_output(output) is not None
|
||||
|
||||
|
||||
def test_duplex_runtime_cancel_fence_rejects_late_append_and_accepts_next_epoch():
|
||||
manager = DuplexSessionRuntimeManager()
|
||||
cancelled_fence = DuplexFence("sid-cancel-race")
|
||||
next_fence = DuplexFence("sid-cancel-race", epoch=1)
|
||||
session = manager.open_session(
|
||||
cancelled_fence,
|
||||
capabilities=DuplexRuntimeCapabilities(
|
||||
input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
|
||||
),
|
||||
)
|
||||
session.bind_stage_request(
|
||||
stage_id=0,
|
||||
request_id="req-cancelled",
|
||||
fence=cancelled_fence,
|
||||
)
|
||||
|
||||
stale_request_ids = session.cancel_fence(cancelled_fence, next_fence)
|
||||
|
||||
assert stale_request_ids == ["req-cancelled"]
|
||||
assert session.fence == next_fence
|
||||
assert session.stage_bindings == {}
|
||||
with pytest.raises(RuntimeError, match="fence mismatch"):
|
||||
session.append_input(
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
fence=cancelled_fence,
|
||||
)
|
||||
update = session.append_input(
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
fence=next_fence,
|
||||
)
|
||||
assert update.seq == 1
|
||||
|
||||
|
||||
def test_duplex_runtime_stale_close_preserves_live_session_and_bindings():
|
||||
manager = DuplexSessionRuntimeManager()
|
||||
current_fence = DuplexFence("sid-stale-close", epoch=1)
|
||||
session = manager.open_session(
|
||||
current_fence,
|
||||
capabilities=DuplexRuntimeCapabilities(
|
||||
input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
|
||||
),
|
||||
)
|
||||
session.bind_stage_request(0, "req-live", fence=current_fence)
|
||||
|
||||
with pytest.raises(RuntimeError, match="fence mismatch"):
|
||||
manager.close_session(DuplexFence("sid-stale-close"))
|
||||
|
||||
assert manager.get("sid-stale-close") is session
|
||||
assert session.stage_request_ids() == ["req-live"]
|
||||
|
||||
|
||||
def test_duplex_runtime_reopen_rejects_late_append_from_old_incarnation():
|
||||
manager = DuplexSessionRuntimeManager()
|
||||
old_fence = DuplexFence("sid-reopen", incarnation=0)
|
||||
old_session = manager.open_session(
|
||||
old_fence,
|
||||
capabilities=DuplexRuntimeCapabilities(
|
||||
input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
|
||||
),
|
||||
)
|
||||
manager.close_session(old_fence)
|
||||
|
||||
new_fence = DuplexFence("sid-reopen", incarnation=1)
|
||||
new_session = manager.open_session(
|
||||
new_fence,
|
||||
capabilities=old_session.capabilities,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="fence mismatch"):
|
||||
new_session.append_input(
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
fence=old_fence,
|
||||
)
|
||||
assert (
|
||||
new_session.append_input(
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
fence=new_fence,
|
||||
).seq
|
||||
== 1
|
||||
)
|
||||
|
||||
|
||||
def test_duplex_prompt_expands_incarnation_metadata():
|
||||
fence = DuplexFence("sid-incarnation", incarnation=3)
|
||||
|
||||
prompt = build_duplex_data_plane_prompt(
|
||||
request_id=duplex_resource_request_id(fence, "stage0"),
|
||||
fence=fence,
|
||||
session_config={},
|
||||
runtime_config={},
|
||||
seq=1,
|
||||
turn_seq=1,
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
payload={"is_speech": True},
|
||||
final=False,
|
||||
)
|
||||
|
||||
assert prompt["model_intermediate_buffer"]["duplex"]["incarnation"] == 3
|
||||
|
||||
|
||||
def test_duplex_runtime_tracks_turn_local_append_sequence():
|
||||
manager = DuplexSessionRuntimeManager()
|
||||
session = manager.open_session(
|
||||
DuplexFence("sid-turn-seq"),
|
||||
capabilities=DuplexRuntimeCapabilities(
|
||||
input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
|
||||
),
|
||||
)
|
||||
|
||||
first = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence)
|
||||
second = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence)
|
||||
next_turn = DuplexFence("sid-turn-seq", turn_id=1, response_seq=1)
|
||||
third = session.append_input(
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
fence=next_turn,
|
||||
)
|
||||
fourth = session.append_input(
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
fence=next_turn,
|
||||
)
|
||||
|
||||
assert [first.seq, second.seq, third.seq, fourth.seq] == [1, 2, 3, 4]
|
||||
assert [first.turn_seq, second.turn_seq, third.turn_seq, fourth.turn_seq] == [1, 2, 1, 2]
|
||||
assert [first.turn_id, second.turn_id, third.turn_id, fourth.turn_id] == [0, 0, 1, 1]
|
||||
|
||||
|
||||
def test_duplex_runtime_rejects_unsupported_append_mode():
|
||||
manager = DuplexSessionRuntimeManager()
|
||||
session = manager.open_session(
|
||||
DuplexFence("sid-2"),
|
||||
capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.TURN_COMMIT_ONLY}),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="not supported"):
|
||||
session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence)
|
||||
|
||||
|
||||
def test_duplex_data_plane_request_info_extracts_structured_stage_result():
|
||||
request_id, response_stage_id = duplex_data_plane_request_info(
|
||||
{
|
||||
"stage_results": [
|
||||
{"result": {"supported": True}},
|
||||
{
|
||||
"result": {
|
||||
"data_plane_append": True,
|
||||
"request_id": "duplex-sid-e0-stage0-s1",
|
||||
"response_stage_id": 1,
|
||||
}
|
||||
},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
assert request_id == "duplex-sid-e0-stage0-s1"
|
||||
assert response_stage_id == 1
|
||||
|
||||
|
||||
def test_duplex_data_plane_request_info_rejects_missing_request_id():
|
||||
assert duplex_data_plane_request_info(
|
||||
{
|
||||
"stage_results": [
|
||||
{
|
||||
"result": {
|
||||
"data_plane_append": True,
|
||||
"request_id": "",
|
||||
"response_stage_id": 1,
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
) == (None, None)
|
||||
|
||||
|
||||
def test_duplex_scheduler_token_budget_estimates_pcm_slots():
|
||||
assert (
|
||||
duplex_scheduler_token_budget(
|
||||
{
|
||||
"audio": "AAAAAA==",
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
}
|
||||
)
|
||||
== 16
|
||||
)
|
||||
|
||||
|
||||
def test_duplex_scheduler_token_budget_ignores_client_budget_fields():
|
||||
assert (
|
||||
duplex_scheduler_token_budget(
|
||||
{
|
||||
"audio": "AAAAAA==",
|
||||
"format": "pcm_f32le",
|
||||
"duplex_num_input_tokens": 999,
|
||||
"num_input_tokens": 999,
|
||||
}
|
||||
)
|
||||
== 16
|
||||
)
|
||||
|
||||
|
||||
def test_resource_state_rejects_fence_regression_and_requires_explicit_fence():
|
||||
current = DuplexFence("sid", epoch=2, turn_id=3, response_seq=4)
|
||||
manager = DuplexSessionRuntimeManager()
|
||||
session = manager.open_session(
|
||||
current,
|
||||
capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}),
|
||||
)
|
||||
|
||||
for stale in (
|
||||
DuplexFence("sid", epoch=1, turn_id=99, response_seq=99),
|
||||
DuplexFence("sid", epoch=2, turn_id=2, response_seq=4),
|
||||
DuplexFence("sid", epoch=2, turn_id=3, response_seq=3),
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="fence mismatch"):
|
||||
session.accept_fence(stale)
|
||||
assert session.fence == current
|
||||
|
||||
with pytest.raises(TypeError, match="fence"):
|
||||
session.bind_stage_request(0, "request")
|
||||
with pytest.raises(TypeError, match="fence"):
|
||||
session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK)
|
||||
with pytest.raises(TypeError, match="DuplexFence"):
|
||||
manager.open_session("legacy-session")
|
||||
with pytest.raises(TypeError, match="DuplexFence"):
|
||||
manager.close_session("legacy-session")
|
||||
|
||||
|
||||
def test_resource_request_id_is_derived_from_fence_and_role():
|
||||
fence = DuplexFence("sid-with-dashes", epoch=7, turn_id=11, response_seq=13)
|
||||
|
||||
assert duplex_resource_request_id(fence, "stage0") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage0"
|
||||
assert duplex_resource_request_id(fence, "stage1") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage1"
|
||||
|
||||
|
||||
def test_resource_request_id_codec_separates_session_id_from_incarnation():
|
||||
embedded_incarnation = duplex_resource_request_id(
|
||||
DuplexFence("foo-i1", incarnation=0),
|
||||
"stage0",
|
||||
)
|
||||
actual_incarnation = duplex_resource_request_id(
|
||||
DuplexFence("foo", incarnation=1),
|
||||
"stage0",
|
||||
)
|
||||
|
||||
assert embedded_incarnation != actual_incarnation
|
||||
|
||||
|
||||
def test_resource_request_id_session_membership_uses_encoded_identity():
|
||||
request_id = duplex_resource_request_id(
|
||||
DuplexFence("sid-with-dashes", incarnation=2, epoch=7),
|
||||
"stage0",
|
||||
)
|
||||
|
||||
assert duplex_resource_request_belongs_to_session(request_id, "sid-with-dashes") is True
|
||||
assert duplex_resource_request_belongs_to_session(request_id, "sid") is False
|
||||
assert duplex_resource_request_belongs_to_session("duplex-s.invalid.i.x.e.7.r.stage0", "sid") is False
|
||||
|
||||
|
||||
def test_placeholder_budget_is_planned_inside_omni_engine_boundary():
|
||||
fence = DuplexFence("sid", turn_id=1, response_seq=1)
|
||||
prompt = build_duplex_data_plane_prompt(
|
||||
request_id=duplex_resource_request_id(fence, "stage0"),
|
||||
fence=fence,
|
||||
session_config={},
|
||||
runtime_config={},
|
||||
seq=2,
|
||||
turn_seq=1,
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
payload={
|
||||
"audio": "AAAAAA==",
|
||||
"format": "pcm_f32le",
|
||||
"duplex_num_input_tokens": 999,
|
||||
"num_input_tokens": 999,
|
||||
},
|
||||
final=False,
|
||||
)
|
||||
|
||||
assert len(prompt["prompt_token_ids"]) == 16
|
||||
assert prompt["model_intermediate_buffer"]["duplex"]["fence"] == fence
|
||||
assert prompt["model_intermediate_buffer"]["duplex"]["scheduler_token_budget"] == 16
|
||||
Whitespace-only changes.
@@ -0,0 +1,99 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.openai.commit_policy import (
|
||||
CommitAction,
|
||||
CommitSnapshot,
|
||||
decide_commit_action,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("case", "snapshot", "expected"),
|
||||
[
|
||||
(
|
||||
"first idle commit",
|
||||
CommitSnapshot(
|
||||
auto_responds=True,
|
||||
speech_since_commit=True,
|
||||
active_response_id=None,
|
||||
overlap_speech_ms=0,
|
||||
native_response_in_progress=False,
|
||||
playback_active=False,
|
||||
),
|
||||
CommitAction.START_AUTO_RESPONSE,
|
||||
),
|
||||
(
|
||||
"native auto-response remains model-owned during playback",
|
||||
CommitSnapshot(
|
||||
auto_responds=True,
|
||||
speech_since_commit=True,
|
||||
active_response_id=None,
|
||||
overlap_speech_ms=800,
|
||||
native_response_in_progress=True,
|
||||
playback_active=True,
|
||||
),
|
||||
CommitAction.START_AUTO_RESPONSE,
|
||||
),
|
||||
(
|
||||
"native auto-response commit without speech",
|
||||
CommitSnapshot(
|
||||
auto_responds=True,
|
||||
speech_since_commit=False,
|
||||
active_response_id=None,
|
||||
overlap_speech_ms=0,
|
||||
native_response_in_progress=True,
|
||||
playback_active=False,
|
||||
),
|
||||
CommitAction.COMMIT_ONLY,
|
||||
),
|
||||
(
|
||||
"non-auto response overlap",
|
||||
CommitSnapshot(
|
||||
auto_responds=False,
|
||||
speech_since_commit=True,
|
||||
active_response_id=None,
|
||||
overlap_speech_ms=800,
|
||||
native_response_in_progress=True,
|
||||
playback_active=False,
|
||||
),
|
||||
CommitAction.DEFER_ACTIVE_RESPONSE,
|
||||
),
|
||||
(
|
||||
"non-auto playback overlap",
|
||||
CommitSnapshot(
|
||||
auto_responds=False,
|
||||
speech_since_commit=True,
|
||||
active_response_id=None,
|
||||
overlap_speech_ms=800,
|
||||
native_response_in_progress=False,
|
||||
playback_active=True,
|
||||
),
|
||||
CommitAction.DEFER_ACTIVE_RESPONSE,
|
||||
),
|
||||
(
|
||||
"non-auto commit without overlap",
|
||||
CommitSnapshot(
|
||||
auto_responds=False,
|
||||
speech_since_commit=True,
|
||||
active_response_id=None,
|
||||
overlap_speech_ms=0,
|
||||
native_response_in_progress=False,
|
||||
playback_active=False,
|
||||
),
|
||||
CommitAction.COMMIT_ONLY,
|
||||
),
|
||||
],
|
||||
ids=lambda value: value if isinstance(value, str) else None,
|
||||
)
|
||||
def test_decide_commit_action_covers_realtime_lifecycle_sequences(
|
||||
case: str,
|
||||
snapshot: CommitSnapshot,
|
||||
expected: CommitAction,
|
||||
) -> None:
|
||||
del case
|
||||
|
||||
assert decide_commit_action(snapshot) is expected
|
||||
@@ -0,0 +1,131 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import base64
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.minicpmo45.input import MiniCPMO45PcmAppendBuffer
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def pcm_payload(samples: int, *, speech: bool = True) -> dict[str, object]:
|
||||
audio = np.ones(samples, dtype=np.float32).tobytes()
|
||||
return {
|
||||
"type": "audio",
|
||||
"audio": base64.b64encode(audio).decode("ascii"),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16_000,
|
||||
"is_speech": speech,
|
||||
}
|
||||
|
||||
|
||||
def test_commit_does_not_add_silence_after_incremental_audio_was_drained():
|
||||
buffer = MiniCPMO45PcmAppendBuffer()
|
||||
|
||||
emitted = buffer.append(pcm_payload(16_000), chunk_period_ms=1_000)
|
||||
committed = buffer.commit(chunk_period_ms=1_000)
|
||||
|
||||
assert emitted is not None
|
||||
assert not buffer.has_pending()
|
||||
assert committed is None
|
||||
|
||||
|
||||
def test_append_emits_one_model_unit_when_multiple_units_are_buffered():
|
||||
buffer = MiniCPMO45PcmAppendBuffer()
|
||||
|
||||
emitted = buffer.append(pcm_payload(32_000), chunk_period_ms=1_000)
|
||||
|
||||
assert emitted is not None
|
||||
assert len(base64.b64decode(emitted["audio"])) == 16_000 * 4
|
||||
assert buffer.pending_byte_count == 16_000 * 4
|
||||
|
||||
|
||||
def test_commit_without_speech_does_not_synthesize_terminal_audio():
|
||||
buffer = MiniCPMO45PcmAppendBuffer()
|
||||
buffer.append(pcm_payload(8_000, speech=False), chunk_period_ms=1_000)
|
||||
|
||||
committed = buffer.commit(chunk_period_ms=1_000)
|
||||
|
||||
assert committed is None
|
||||
|
||||
|
||||
def test_commit_resets_cumulative_turn_accounting():
|
||||
buffer = MiniCPMO45PcmAppendBuffer()
|
||||
buffer.append(pcm_payload(16_000), chunk_period_ms=1_000)
|
||||
buffer.commit(chunk_period_ms=1_000)
|
||||
|
||||
empty = buffer.commit(chunk_period_ms=1_000)
|
||||
|
||||
assert empty is None
|
||||
|
||||
|
||||
def test_pcm_append_reservation_rollback_restores_emitted_audio():
|
||||
buffer = MiniCPMO45PcmAppendBuffer()
|
||||
original = pcm_payload(16_000)
|
||||
|
||||
reservation = buffer.prepare_append(
|
||||
original,
|
||||
operation_id="append-1",
|
||||
chunk_period_ms=1_000,
|
||||
)
|
||||
|
||||
assert reservation is not None
|
||||
assert not buffer.has_pending()
|
||||
reservation.rollback()
|
||||
assert buffer.has_pending()
|
||||
|
||||
retried = buffer.flush(chunk_period_ms=1_000)
|
||||
assert retried is not None
|
||||
assert base64.b64decode(retried["audio"]) == base64.b64decode(original["audio"])
|
||||
|
||||
|
||||
def test_pcm_commit_keeps_prior_append_reservation_active():
|
||||
buffer = MiniCPMO45PcmAppendBuffer()
|
||||
append_reservation = buffer.prepare_append(
|
||||
pcm_payload(16_000),
|
||||
operation_id="append-before-commit",
|
||||
chunk_period_ms=1_000,
|
||||
)
|
||||
|
||||
commit_reservation = buffer.prepare_commit(
|
||||
operation_id="commit-after-append",
|
||||
chunk_period_ms=1_000,
|
||||
)
|
||||
|
||||
assert append_reservation is not None
|
||||
assert append_reservation.active
|
||||
assert commit_reservation.payload is None
|
||||
append_reservation.commit()
|
||||
commit_reservation.commit()
|
||||
|
||||
|
||||
def test_pcm_commit_reservation_rollback_restores_residual_audio():
|
||||
buffer = MiniCPMO45PcmAppendBuffer()
|
||||
original = pcm_payload(8_000)
|
||||
assert (
|
||||
buffer.prepare_append(
|
||||
original,
|
||||
operation_id="buffer-half-chunk",
|
||||
chunk_period_ms=1_000,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
reservation = buffer.prepare_commit(
|
||||
operation_id="final-half-chunk",
|
||||
chunk_period_ms=1_000,
|
||||
)
|
||||
|
||||
assert reservation.payload is not None
|
||||
assert reservation.payload["final"] is True
|
||||
reservation.rollback()
|
||||
|
||||
retried = buffer.prepare_commit(
|
||||
operation_id="retry-final-half-chunk",
|
||||
chunk_period_ms=1_000,
|
||||
)
|
||||
assert retried.payload is not None
|
||||
assert base64.b64decode(retried.payload["audio"]) == (base64.b64decode(original["audio"]) + b"\x00" * (8_000 * 4))
|
||||
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import base64
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.openai.audio import convert_input_audio_with_rate
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_pcm16_input_conversion_is_pure_and_reports_target_rate():
|
||||
pcm16 = np.array([-32768, 0, 32767], dtype="<i2").tobytes()
|
||||
encoded = base64.b64encode(pcm16).decode("ascii")
|
||||
|
||||
converted, fmt, sample_rate = convert_input_audio_with_rate(
|
||||
encoded,
|
||||
"pcm16",
|
||||
sample_rate_hz=16_000,
|
||||
)
|
||||
|
||||
assert encoded == base64.b64encode(pcm16).decode("ascii")
|
||||
assert fmt == "pcm_f32le"
|
||||
assert sample_rate == 16_000
|
||||
samples = np.frombuffer(base64.b64decode(converted), dtype="<f4")
|
||||
assert np.allclose(samples, [-1.0, 0.0, 32767 / 32768])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sample_rate_hz", [8_000, 192_000])
|
||||
def test_pcm16_input_conversion_accepts_supported_sample_rate_boundaries(sample_rate_hz: int):
|
||||
encoded = base64.b64encode(np.zeros(8, dtype="<i2").tobytes()).decode("ascii")
|
||||
|
||||
_, fmt, converted_rate = convert_input_audio_with_rate(
|
||||
encoded,
|
||||
"pcm16",
|
||||
sample_rate_hz=sample_rate_hz,
|
||||
)
|
||||
|
||||
assert fmt == "pcm_f32le"
|
||||
assert converted_rate == 16_000
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sample_rate_hz", [7_999, 192_001, 10**1000])
|
||||
def test_pcm16_input_conversion_rejects_sample_rates_outside_supported_range(sample_rate_hz: int):
|
||||
encoded = base64.b64encode(np.zeros(8, dtype="<i2").tobytes()).decode("ascii")
|
||||
|
||||
with pytest.raises(ValueError, match="sample_rate_hz"):
|
||||
convert_input_audio_with_rate(
|
||||
encoded,
|
||||
"pcm16",
|
||||
sample_rate_hz=sample_rate_hz,
|
||||
)
|
||||
@@ -0,0 +1,139 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.openai.runtime_adapter import (
|
||||
validate_serving_runtime_adapter,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"module_name",
|
||||
[
|
||||
"vllm_omni.experimental.fullduplex.openai.session_runner",
|
||||
"vllm_omni.experimental.fullduplex.openai.runtime_bridge",
|
||||
"vllm_omni.experimental.fullduplex.openai.serving",
|
||||
],
|
||||
)
|
||||
def test_generic_openai_runtime_import_does_not_load_minicpmo45(module_name: str) -> None:
|
||||
repo_root = Path(__file__).resolve().parents[3]
|
||||
script = f"""
|
||||
import sys
|
||||
|
||||
import {module_name}
|
||||
|
||||
loaded = sorted(
|
||||
name
|
||||
for name in sys.modules
|
||||
if name == "vllm_omni.experimental.fullduplex.minicpmo45"
|
||||
or name.startswith("vllm_omni.experimental.fullduplex.minicpmo45.")
|
||||
)
|
||||
if loaded:
|
||||
raise SystemExit("generic OpenAI runtime loaded MiniCPM-o: " + ", ".join(loaded))
|
||||
"""
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
cwd=repo_root,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60,
|
||||
check=False,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_generic_handler_requires_explicit_serving_runtime_adapter() -> None:
|
||||
repo_root = Path(__file__).resolve().parents[3]
|
||||
script = """
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler
|
||||
|
||||
chat_service = SimpleNamespace(engine_client=SimpleNamespace())
|
||||
try:
|
||||
OmniDuplexSessionHandler(chat_service=chat_service)
|
||||
except ValueError as exc:
|
||||
if "serving runtime adapter" not in str(exc).lower():
|
||||
raise
|
||||
else:
|
||||
raise SystemExit("generic handler silently selected a model serving adapter")
|
||||
"""
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
cwd=repo_root,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60,
|
||||
check=False,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def _valid_runtime_adapter() -> SimpleNamespace:
|
||||
data_plane = SimpleNamespace(
|
||||
begin_request=lambda request_id: None,
|
||||
is_terminal=lambda request_id: False,
|
||||
mark_terminal=lambda request_id: None,
|
||||
close_stream=lambda request_id: None,
|
||||
close_session=lambda session_id, **kwargs: None,
|
||||
project=lambda result, **kwargs: (),
|
||||
)
|
||||
return SimpleNamespace(
|
||||
adapter_id="test-adapter",
|
||||
session_states={},
|
||||
data_plane=data_plane,
|
||||
clean_response_done_prefix="clean",
|
||||
interrupted_tts_prefix="interrupted",
|
||||
private_runtime_config_keys=frozenset({"private_key"}),
|
||||
create_session_state=lambda: object(),
|
||||
session_state=lambda session_id: object(),
|
||||
remove_session_state=lambda session_id: None,
|
||||
is_enabled=lambda config: True,
|
||||
capabilities=lambda **kwargs: object(),
|
||||
validate_client_extra_body=lambda extra_body: None,
|
||||
prepare_runtime_config=lambda config, **kwargs: {},
|
||||
runtime_config_for_update=lambda config, current: {},
|
||||
data_plane_context=lambda **kwargs: object(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("attribute", "invalid_value"),
|
||||
[
|
||||
("session_states", None),
|
||||
("clean_response_done_prefix", None),
|
||||
("interrupted_tts_prefix", None),
|
||||
("private_runtime_config_keys", {"private_key"}),
|
||||
],
|
||||
)
|
||||
def test_runtime_adapter_validator_rejects_invalid_protocol_attributes(
|
||||
attribute: str,
|
||||
invalid_value: object,
|
||||
) -> None:
|
||||
adapter = _valid_runtime_adapter()
|
||||
setattr(adapter, attribute, invalid_value)
|
||||
|
||||
with pytest.raises(TypeError, match=attribute):
|
||||
validate_serving_runtime_adapter(adapter)
|
||||
|
||||
|
||||
def test_runtime_adapter_validator_rejects_incomplete_data_plane() -> None:
|
||||
adapter = _valid_runtime_adapter()
|
||||
adapter.data_plane.close_session = None
|
||||
|
||||
with pytest.raises(TypeError, match="data_plane.*close_session"):
|
||||
validate_serving_runtime_adapter(adapter)
|
||||
@@ -0,0 +1,145 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.openai.websocket import (
|
||||
DuplexWebSocketActor,
|
||||
normalize_duplex_input_event,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeWebSocket:
|
||||
def __init__(self) -> None:
|
||||
self.sent: list[dict[str, object]] = []
|
||||
|
||||
async def send_json(self, payload: dict[str, object]) -> None:
|
||||
self.sent.append(dict(payload))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("event", "expected"),
|
||||
[
|
||||
({"type": "signal_turn", "event": "barge_in"}, {"type": "turn.signal", "event": "barge_in"}),
|
||||
({"type": "close_session"}, {"type": "session.close"}),
|
||||
({"type": "audio.playback_ack", "played_ms": 1}, {"type": "playback.ack", "played_ms": 1}),
|
||||
({"type": "input_text.append", "text": "a"}, {"type": "input.text.append", "text": "a"}),
|
||||
({"type": "push_text", "text": "b"}, {"type": "input.text.append", "text": "b"}),
|
||||
(
|
||||
{"type": "input.audio.append", "audio": "wav"},
|
||||
{"type": "input_audio_buffer.append", "audio": "wav", "format": "wav"},
|
||||
),
|
||||
(
|
||||
{"type": "push_chunk", "audio": "wav", "format": "pcm_f32le"},
|
||||
{"type": "input_audio_buffer.append", "audio": "wav", "format": "pcm_f32le"},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_input_aliases_normalize_once_at_mailbox_boundary(event, expected):
|
||||
assert normalize_duplex_input_event(event) == expected
|
||||
|
||||
|
||||
def test_actor_uses_one_fifo_mailbox_for_inbound_events():
|
||||
actor = DuplexWebSocketActor(FakeWebSocket())
|
||||
|
||||
assert isinstance(actor.mailbox, asyncio.Queue)
|
||||
assert not isinstance(actor.mailbox, asyncio.PriorityQueue)
|
||||
assert not hasattr(actor, "control_queue")
|
||||
assert not hasattr(actor, "input_queue")
|
||||
assert not hasattr(actor, "event_queue")
|
||||
assert not hasattr(actor, "lifecycle_state")
|
||||
assert not hasattr(actor, "last_response_id")
|
||||
assert not hasattr(actor, "overlap_speech_ms")
|
||||
assert not hasattr(actor, "transition")
|
||||
assert not hasattr(actor, "output_generation_in_flight")
|
||||
assert not hasattr(actor, "session")
|
||||
assert not hasattr(actor, "runtime_opened")
|
||||
assert not hasattr(actor, "runtime_closed")
|
||||
assert not hasattr(actor, "drain_input_queue")
|
||||
assert not hasattr(actor, "control_events_seen")
|
||||
assert not hasattr(actor, "input_events_seen")
|
||||
assert not hasattr(actor, "cancel_count")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_control_event_does_not_overtake_earlier_audio_input():
|
||||
actor = DuplexWebSocketActor(FakeWebSocket())
|
||||
await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"})
|
||||
await actor.enqueue_event({"type": "response.cancel"})
|
||||
|
||||
assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"}
|
||||
assert await actor.next_event() == {"type": "response.cancel"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("event_type", ["input.cancel", "session.close", "close_session"])
|
||||
async def test_terminal_control_preserves_wire_order(event_type: str):
|
||||
actor = DuplexWebSocketActor(FakeWebSocket())
|
||||
await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"})
|
||||
await actor.enqueue_event({"type": event_type})
|
||||
|
||||
assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"}
|
||||
assert await actor.next_event() == {"type": event_type}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mailbox_delivers_every_enqueued_event_exactly_once():
|
||||
actor = DuplexWebSocketActor(FakeWebSocket())
|
||||
events = [
|
||||
{"type": "input_audio_buffer.append", "audio": str(index)}
|
||||
if index % 2
|
||||
else {"type": "response.created", "response_id": f"resp-{index}"}
|
||||
for index in range(20)
|
||||
]
|
||||
|
||||
await asyncio.gather(*(actor.enqueue_event(dict(event)) for event in events))
|
||||
received = [await actor.next_event() for _ in events]
|
||||
|
||||
assert sorted(received, key=repr) == sorted(events, key=repr)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_writer_is_single_owner_of_websocket_send():
|
||||
websocket = FakeWebSocket()
|
||||
actor = DuplexWebSocketActor(websocket)
|
||||
writer = asyncio.create_task(actor.writer_loop())
|
||||
|
||||
await actor.send_json({"type": "one"})
|
||||
await actor.send_json({"type": "two"})
|
||||
await actor.close_writer()
|
||||
await writer
|
||||
|
||||
assert websocket.sent == [{"type": "one"}, {"type": "two"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_fence_payload_is_dropped_before_websocket_send():
|
||||
websocket = FakeWebSocket()
|
||||
actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 2)
|
||||
writer = asyncio.create_task(actor.writer_loop())
|
||||
|
||||
await actor.send_json({"type": "response.audio.delta", "epoch": 1})
|
||||
await actor.send_json({"type": "response.audio.delta", "epoch": 2})
|
||||
await actor.close_writer()
|
||||
await writer
|
||||
|
||||
assert websocket.sent == [{"type": "response.audio.delta", "epoch": 2}]
|
||||
assert actor.stale_output_dropped == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_writer_does_not_revoke_accepted_terminal_after_close_starts():
|
||||
websocket = FakeWebSocket()
|
||||
actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 1)
|
||||
await actor.send_json({"type": "response.done", "epoch": 1, "response_id": "resp-1"})
|
||||
actor.closing = True
|
||||
await actor.close_writer()
|
||||
|
||||
await actor.writer_loop()
|
||||
|
||||
assert websocket.sent == [{"type": "response.done", "epoch": 1, "response_id": "resp-1"}]
|
||||
assert actor.stale_output_dropped == 0
|
||||
@@ -0,0 +1,240 @@
|
||||
import base64
|
||||
import wave
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.client import (
|
||||
RealtimeDuplexClient,
|
||||
RealtimeEventCollector,
|
||||
build_realtime_url,
|
||||
read_pcm16_wav,
|
||||
write_pcm16_wav,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_realtime_client_builds_explicit_native_duplex_url():
|
||||
url = build_realtime_url(
|
||||
"ws://localhost:8099/v1/realtime?custom=1",
|
||||
"openbmb/MiniCPM-o-4_5",
|
||||
session_id="session-a",
|
||||
)
|
||||
|
||||
query = parse_qs(urlsplit(url).query)
|
||||
assert query == {
|
||||
"custom": ["1"],
|
||||
"duplex": ["1"],
|
||||
"model": ["openbmb/MiniCPM-o-4_5"],
|
||||
"minicpmo45_native_duplex": ["1"],
|
||||
"session_id": ["session-a"],
|
||||
}
|
||||
|
||||
|
||||
def test_realtime_client_builds_resume_only_url_when_autostart_disabled():
|
||||
url = build_realtime_url(
|
||||
"ws://localhost:8099/v1/realtime?duplex=1",
|
||||
"openbmb/MiniCPM-o-4_5",
|
||||
autostart=False,
|
||||
)
|
||||
|
||||
query = parse_qs(urlsplit(url).query)
|
||||
assert query["autostart"] == ["0"]
|
||||
assert query["minicpmo45_native_duplex"] == ["1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_client_configure_omits_ref_audio_by_default():
|
||||
class Client(RealtimeDuplexClient):
|
||||
def __init__(self):
|
||||
super().__init__("ws://unused")
|
||||
self.sent = []
|
||||
|
||||
async def send(self, event):
|
||||
self.sent.append(event)
|
||||
self.events.add({"type": "session.created"})
|
||||
|
||||
client = Client()
|
||||
|
||||
await client.configure("openbmb/MiniCPM-o-4_5", timeout_s=1)
|
||||
|
||||
session = client.sent[0]["session"]
|
||||
assert "ref_audio" not in session
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_client_configure_sends_explicit_ref_audio():
|
||||
class Client(RealtimeDuplexClient):
|
||||
def __init__(self):
|
||||
super().__init__("ws://unused")
|
||||
self.sent = []
|
||||
|
||||
async def send(self, event):
|
||||
self.sent.append(event)
|
||||
self.events.add({"type": "session.created"})
|
||||
|
||||
client = Client()
|
||||
|
||||
await client.configure(
|
||||
"openbmb/MiniCPM-o-4_5",
|
||||
ref_audio="data:audio/wav;base64,AAAA",
|
||||
timeout_s=1,
|
||||
)
|
||||
|
||||
session = client.sent[0]["session"]
|
||||
assert session["ref_audio"] == "data:audio/wav;base64,AAAA"
|
||||
|
||||
|
||||
def test_realtime_event_collector_partitions_audio_by_response():
|
||||
collector = RealtimeEventCollector()
|
||||
collector.add({"type": "response.created", "response": {"id": "resp-a"}})
|
||||
collector.add(
|
||||
{
|
||||
"type": "response.audio.delta",
|
||||
"response_id": "resp-a",
|
||||
"delta": base64.b64encode(b"audio-a").decode("ascii"),
|
||||
"sample_rate_hz": 16_000,
|
||||
}
|
||||
)
|
||||
|
||||
assert collector.response_ids == ["resp-a"]
|
||||
assert collector.audio_bytes("resp-a") == b"audio-a"
|
||||
assert collector.output_sample_rate_hz == 16_000
|
||||
assert collector.first_received_at("response.created") is not None
|
||||
assert collector.last_received_at("response.audio.delta") is not None
|
||||
|
||||
|
||||
def test_realtime_event_collector_reports_engine_token_and_audio_intervals():
|
||||
collector = RealtimeEventCollector()
|
||||
collector.add(
|
||||
{"type": "response.created", "response": {"id": "resp-a"}},
|
||||
received_at_s=10.0,
|
||||
)
|
||||
stage_metrics = {
|
||||
"0": {
|
||||
"num_tokens_out": 4,
|
||||
"vllm_ttft_ms": 120.0,
|
||||
"vllm_tpot_ms": 15.0,
|
||||
"vllm_itl_ms": 14.0,
|
||||
"vllm_itls_ms": [10.0, 14.0, 18.0],
|
||||
}
|
||||
}
|
||||
for received_at_s, cumulative_audio_ms in ((10.2, 80), (10.25, 160), (10.36, 240)):
|
||||
collector.add(
|
||||
{
|
||||
"type": "response.audio.delta",
|
||||
"response_id": "resp-a",
|
||||
"delta": base64.b64encode(b"audio").decode("ascii"),
|
||||
"sample_rate_hz": 16_000,
|
||||
"metadata": {
|
||||
"audio_duration_ms": cumulative_audio_ms,
|
||||
"vllm_omni": {"stage_metrics": stage_metrics},
|
||||
},
|
||||
},
|
||||
received_at_s=received_at_s,
|
||||
)
|
||||
|
||||
timing = collector.timing_summary(
|
||||
after_s=10.0,
|
||||
input_committed_at_s=9.9,
|
||||
response_id="resp-a",
|
||||
)
|
||||
|
||||
assert timing["stage0_tokens"] == {
|
||||
"source": "engine_stage_metrics",
|
||||
"output_token_count": 4,
|
||||
"ttft_ms": 120.0,
|
||||
"tpot_ms": 15.0,
|
||||
"inter_token_interval_ms": {
|
||||
"count": 3,
|
||||
"mean": 14.0,
|
||||
"p50": 14.0,
|
||||
"p95": 18.0,
|
||||
"max": 18.0,
|
||||
},
|
||||
}
|
||||
assert timing["audio_output"] == {
|
||||
"source": "client_monotonic_receive",
|
||||
"chunk_count": 3,
|
||||
"response_created_to_first_audio_ms": 200.0,
|
||||
"commit_to_first_audio_ms": 300.0,
|
||||
"inter_chunk_interval_ms": {
|
||||
"count": 2,
|
||||
"mean": 80.0,
|
||||
"p50": 50.0,
|
||||
"p95": 110.0,
|
||||
"max": 110.0,
|
||||
},
|
||||
"chunk_duration_ms": {
|
||||
"count": 3,
|
||||
"mean": 80.0,
|
||||
"p50": 80.0,
|
||||
"p95": 80.0,
|
||||
"max": 80.0,
|
||||
},
|
||||
"max_chunk_gap_ms": 110.0,
|
||||
}
|
||||
|
||||
|
||||
def test_response_timing_ignores_unowned_session_level_metrics():
|
||||
collector = RealtimeEventCollector()
|
||||
collector.add(
|
||||
{"type": "response.created", "response": {"id": "resp-a"}},
|
||||
received_at_s=10.0,
|
||||
)
|
||||
collector.add(
|
||||
{
|
||||
"type": "response.audio.delta",
|
||||
"response_id": "resp-a",
|
||||
"delta": base64.b64encode(b"audio").decode("ascii"),
|
||||
"metadata": {
|
||||
"vllm_omni": {
|
||||
"stage_metrics": {
|
||||
"0": {
|
||||
"num_tokens_out": 20,
|
||||
"vllm_ttft_ms": 157.0,
|
||||
"vllm_tpot_ms": 16.0,
|
||||
"vllm_itls_ms": [15.0, 17.0],
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
received_at_s=10.2,
|
||||
)
|
||||
collector.add(
|
||||
{
|
||||
"type": "response.listen",
|
||||
"metadata": {
|
||||
"vllm_omni": {
|
||||
"stage_metrics": {
|
||||
"0": {
|
||||
"num_tokens_out": 2,
|
||||
"vllm_ttft_ms": 106.0,
|
||||
"vllm_tpot_ms": 0.0,
|
||||
"vllm_itls_ms": [],
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
received_at_s=10.3,
|
||||
)
|
||||
|
||||
timing = collector.timing_summary(after_s=10.0, response_id="resp-a")
|
||||
|
||||
assert timing["stage0_tokens"]["output_token_count"] == 20
|
||||
assert timing["stage0_tokens"]["ttft_ms"] == 157.0
|
||||
|
||||
|
||||
def test_realtime_client_pcm16_wav_round_trip(tmp_path):
|
||||
path = tmp_path / "audio.wav"
|
||||
pcm16 = b"\x01\x00\x02\x00"
|
||||
|
||||
write_pcm16_wav(path, pcm16, sample_rate_hz=16_000)
|
||||
|
||||
with wave.open(str(path), "rb") as wav_file:
|
||||
assert wav_file.getnchannels() == 1
|
||||
assert wav_file.getframerate() == 16_000
|
||||
assert read_pcm16_wav(path) == pcm16
|
||||
@@ -5,7 +5,6 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
|
||||
MiniCPMO45OmniTTSForConditionalGeneration,
|
||||
@@ -14,10 +13,19 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _FakeMiniCPMTTS:
|
||||
def __init__(self, config, audio_tokenizer):
|
||||
self.config = config
|
||||
self.audio_tokenizer = audio_tokenizer
|
||||
self.emb_text = object()
|
||||
self.projector_semantic = object()
|
||||
|
||||
def generate(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def test_hf_model_id_path_caching(mocker, tmp_path):
|
||||
"""Verify HF model-ID path is resolved via download_weights_from_hf_specific
|
||||
once during _lazy_init_tts, cached in self._model_path, and reused in generate_speech.
|
||||
"""
|
||||
"""Verify the resolved HF path is cached and reused for default ref audio."""
|
||||
resolved_path = str(tmp_path / "resolved_hf_model")
|
||||
assets_dir = tmp_path / "resolved_hf_model" / "assets"
|
||||
assets_dir.mkdir(parents=True, exist_ok=True)
|
||||
@@ -29,11 +37,9 @@ def test_hf_model_id_path_caching(mocker, tmp_path):
|
||||
return_value=resolved_path,
|
||||
)
|
||||
|
||||
# Mock dynamic class loading
|
||||
fake_tts_class = mocker.MagicMock()
|
||||
mocker.patch(
|
||||
"transformers.dynamic_module_utils.get_class_from_dynamic_module",
|
||||
return_value=fake_tts_class,
|
||||
return_value=_FakeMiniCPMTTS,
|
||||
)
|
||||
|
||||
# Build fake vllm_config with HF model ID
|
||||
@@ -59,40 +65,10 @@ def test_hf_model_id_path_caching(mocker, tmp_path):
|
||||
|
||||
tts_model = MiniCPMO45OmniTTSForConditionalGeneration(vllm_config=vllm_config)
|
||||
|
||||
# 1. Trigger lazy init
|
||||
tts_model._lazy_init_tts()
|
||||
tts_model._lazy_init_tts()
|
||||
|
||||
# Verify download_weights_from_hf_specific was called with HF model ID and cached
|
||||
mock_download.assert_called_once_with(model_id, None, ["*"])
|
||||
assert tts_model._model_path == resolved_path
|
||||
assert tts_model._assets_loaded is True
|
||||
|
||||
# 2. Mock TTS object generate method and audio tokenizer for generate_speech
|
||||
fake_tts_obj = mocker.MagicMock()
|
||||
fake_tts_obj.emb_text.weight.device = torch.device("cpu")
|
||||
fake_tts_obj.emb_text.side_effect = lambda t: torch.zeros(t.shape[0], 768)
|
||||
fake_tts_obj.projector_semantic.side_effect = lambda h, **kw: torch.zeros(h.shape[0], 768)
|
||||
fake_tts_obj.config = tts_cfg
|
||||
fake_tts_obj.audio_bos_token_id = 151687
|
||||
|
||||
fake_output = SimpleNamespace(new_ids=torch.tensor([[[100]]]))
|
||||
fake_tts_obj.generate.return_value = fake_output
|
||||
|
||||
tts_model.tts_obj = fake_tts_obj
|
||||
fake_tokenizer = mocker.MagicMock(return_value=b"dummy_wav_bytes")
|
||||
tts_model.audio_tokenizer = fake_tokenizer
|
||||
|
||||
# Mock sf.read to return fake numpy audio
|
||||
mocker.patch(
|
||||
"soundfile.read",
|
||||
return_value=(torch.zeros(16000).numpy(), 24000),
|
||||
)
|
||||
|
||||
# Call generate_speech
|
||||
tts_token_ids = torch.tensor([1, 2, 3])
|
||||
tts_hidden_states = torch.zeros(3, 768)
|
||||
waveform = tts_model.generate_speech(tts_token_ids, tts_hidden_states)
|
||||
|
||||
# Verify download_weights_from_hf_specific was NOT called a second time
|
||||
assert mock_download.call_count == 1
|
||||
assert waveform is not None
|
||||
assert tts_model._resolve_prompt_wav_path(None, None) == (str(ref_audio_file), None)
|
||||
@@ -7,11 +7,9 @@ Covers ``vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2t
|
||||
- empty ``source_outputs`` raises
|
||||
- latent fallback to ``hidden_states`` when ``multimodal_output`` is empty
|
||||
- both inputs missing -> raises
|
||||
- additional_information payload carries the keys the talker expects
|
||||
(prompt_embeds, prompt_token_ids, llm_output_token_ids, llm_output_text)
|
||||
- dummy talker prompt ``[BOS, PAD, EOS] = [1, 0, 2]`` (single prefill step)
|
||||
- structured model_intermediate_buffer carries the thinker ids and text
|
||||
- scheduler prompt tokens follow the selected TTS region or output tokens
|
||||
- MiniCPM-o 4.5 TTS region detection on 151703 / 151704 tokens
|
||||
- MiniCPM-o 2.6 fallback detection on 151691 / 151692 when no 4.5 markers
|
||||
- No TTS markers present -> no ``tts_token_ids`` / ``tts_hidden_states`` keys
|
||||
- prompt arg is normalized to a list and ``multi_modal_data`` is gated by
|
||||
``requires_multimodal_data``
|
||||
@@ -24,6 +22,9 @@ from types import SimpleNamespace
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
|
||||
MiniCPMO45OmniForConditionalGeneration,
|
||||
)
|
||||
from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import llm2tts
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -32,6 +33,55 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
_HIDDEN_DIM = 4
|
||||
|
||||
|
||||
def test_native_duplex_talker_emits_empty_text_metadata_to_clear_previous_segment() -> None:
|
||||
class _Talker:
|
||||
_ar_last_chunk_flags = [True]
|
||||
_ar_turn_end_flags = [True]
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
del kwargs
|
||||
return None, torch.zeros(0, dtype=torch.float32)
|
||||
|
||||
model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.model_stage = "tts"
|
||||
model.config = SimpleNamespace(hidden_size=_HIDDEN_DIM)
|
||||
model.talker = _Talker()
|
||||
|
||||
output = model.forward(
|
||||
input_ids=torch.tensor([1]),
|
||||
positions=torch.tensor([0]),
|
||||
runtime_additional_information=[
|
||||
{
|
||||
"native_duplex": True,
|
||||
"duplex": {"turn_id": 1, "epoch": 0},
|
||||
"meta": {"native_duplex_segment_text": ""},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
assert output.multimodal_outputs is not None
|
||||
assert output.multimodal_outputs["meta.llm_output_text_utf8"].numel() == 0
|
||||
assert output.multimodal_outputs["meta.audio_text_total_chars"].tolist() == [0]
|
||||
|
||||
|
||||
def test_tts_embed_input_ids_uses_active_talker_embedding(monkeypatch) -> None:
|
||||
model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.model_stage = "tts"
|
||||
expected = torch.ones((2, _HIDDEN_DIM))
|
||||
|
||||
monkeypatch.setattr(
|
||||
model,
|
||||
"get_input_embeddings",
|
||||
lambda input_ids: expected,
|
||||
)
|
||||
|
||||
actual = model.embed_input_ids(torch.tensor([11, 12]))
|
||||
|
||||
assert actual is expected
|
||||
|
||||
|
||||
def _make_thinker_output(
|
||||
*,
|
||||
prompt_token_ids: list[int],
|
||||
@@ -40,6 +90,7 @@ def _make_thinker_output(
|
||||
request_id: str = "req-0",
|
||||
latent: torch.Tensor | None = None,
|
||||
hidden_states: torch.Tensor | None = None,
|
||||
multimodal_output: dict[str, object] | None = None,
|
||||
):
|
||||
"""Construct a minimal mock of a thinker engine output entry.
|
||||
|
||||
@@ -48,8 +99,11 @@ def _make_thinker_output(
|
||||
- per-output: ``multimodal_output`` (dict), ``hidden_states`` (opt),
|
||||
``text``, ``token_ids``
|
||||
"""
|
||||
mm_output = dict(multimodal_output or {})
|
||||
if latent is not None:
|
||||
mm_output["latent"] = latent
|
||||
output = SimpleNamespace(
|
||||
multimodal_output={"latent": latent} if latent is not None else {},
|
||||
multimodal_output=mm_output,
|
||||
token_ids=output_token_ids,
|
||||
text=text,
|
||||
)
|
||||
@@ -90,17 +144,16 @@ class TestBasicShape:
|
||||
)
|
||||
assert len(out) == 2
|
||||
|
||||
def test_talker_prompt_token_ids_dummy_bos_pad_eos(self) -> None:
|
||||
def test_talker_scheduler_prompt_uses_output_tokens(self) -> None:
|
||||
hidden = torch.zeros((2, _HIDDEN_DIM))
|
||||
out = llm2tts(
|
||||
[_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden)],
|
||||
prompt=None,
|
||||
)
|
||||
# The talker AR framework needs *some* tokens to do a single prefill;
|
||||
# this is the agreed [BOS, PAD, EOS] minimal payload.
|
||||
assert out[0]["prompt_token_ids"] == [1, 0, 2]
|
||||
assert out[0]["prompt_token_ids"] == [20]
|
||||
assert "stream_output" not in out[0]["model_intermediate_buffer"]
|
||||
|
||||
def test_additional_information_carries_thinker_outputs(self) -> None:
|
||||
def test_model_intermediate_buffer_carries_thinker_outputs(self) -> None:
|
||||
prompt_ids = [10, 11, 12]
|
||||
out_ids = [20, 21]
|
||||
hidden = torch.randn(len(prompt_ids) + len(out_ids), _HIDDEN_DIM)
|
||||
@@ -116,21 +169,17 @@ class TestBasicShape:
|
||||
],
|
||||
prompt=None,
|
||||
)
|
||||
ai = result[0]["additional_information"]
|
||||
assert ai["prompt_token_ids"] == prompt_ids
|
||||
assert ai["llm_output_token_ids"] == out_ids
|
||||
assert ai["llm_output_text"] == ["hello"]
|
||||
# prompt_embeds = float32 view of the prompt-portion of hidden states.
|
||||
assert ai["prompt_embeds"].dtype == torch.float32
|
||||
assert ai["prompt_embeds"].shape == (len(prompt_ids), _HIDDEN_DIM)
|
||||
assert torch.equal(ai["prompt_embeds"], hidden[: len(prompt_ids)].to(torch.float32))
|
||||
buffer = result[0]["model_intermediate_buffer"]
|
||||
assert buffer["ids"]["prompt"] == prompt_ids
|
||||
assert buffer["ids"]["output"] == out_ids
|
||||
assert buffer["llm_output_text"] == ["hello"]
|
||||
|
||||
def test_latent_in_multimodal_output_takes_precedence(self) -> None:
|
||||
# When both ``multimodal_output["latent"]`` and ``hidden_states`` are
|
||||
# present, the latent payload must win (this is the steady-state path
|
||||
# produced by the thinker stage).
|
||||
prompt_ids = [10, 11]
|
||||
out_ids = [20]
|
||||
out_ids = [151703, 30, 151704]
|
||||
latent = torch.ones((len(prompt_ids) + len(out_ids), _HIDDEN_DIM))
|
||||
hidden = torch.zeros_like(latent)
|
||||
|
||||
@@ -145,19 +194,17 @@ class TestBasicShape:
|
||||
],
|
||||
prompt=None,
|
||||
)
|
||||
ai = result[0]["additional_information"]
|
||||
buffer = result[0]["model_intermediate_buffer"]
|
||||
# latent (ones) won over hidden_states (zeros)
|
||||
assert torch.equal(ai["prompt_embeds"], latent[: len(prompt_ids)].to(torch.float32))
|
||||
assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), latent[3:4].to(torch.float32))
|
||||
|
||||
|
||||
class TestTtsRegionDetection:
|
||||
"""The bridge auto-detects which MiniCPM-o generation produced the output
|
||||
by sniffing for 4.5-specific BOS/EOS token ids (151703 / 151704); when
|
||||
those are absent it falls back to the 2.6 ids (151691 / 151692).
|
||||
"""The bridge detects MiniCPM-o 4.5 BOS/EOS ids (151703 / 151704).
|
||||
|
||||
This is a regression guard: a single off-by-one or wrong branch
|
||||
precedence here causes the talker to receive an empty / wrong-region
|
||||
slice and emit silent audio.
|
||||
causes the talker to receive an empty / wrong-region slice and emit
|
||||
silent audio.
|
||||
"""
|
||||
|
||||
def _run(self, prompt_token_ids, output_token_ids):
|
||||
@@ -173,7 +220,7 @@ class TestTtsRegionDetection:
|
||||
],
|
||||
prompt=None,
|
||||
)
|
||||
return result[0]["additional_information"], hidden
|
||||
return result[0]["model_intermediate_buffer"], hidden
|
||||
|
||||
def test_4_5_markers_detected(self) -> None:
|
||||
# prompt: [10, 11]
|
||||
@@ -181,45 +228,26 @@ class TestTtsRegionDetection:
|
||||
# full sequence: [10, 11, 151703, 30, 31, 151704, 40]
|
||||
# 0 1 2 3 4 5 6
|
||||
# 4.5 BOS at idx 2 -> slice starts at 3; EOS at idx 5 -> slice ends at 5.
|
||||
ai, hidden = self._run([10, 11], [151703, 30, 31, 151704, 40])
|
||||
assert "tts_token_ids" in ai and "tts_hidden_states" in ai
|
||||
assert ai["tts_token_ids"].tolist() == [30, 31]
|
||||
assert torch.equal(ai["tts_hidden_states"], hidden[3:5])
|
||||
|
||||
def test_4_5_takes_precedence_when_both_markers_present(self) -> None:
|
||||
# If both 2.6 and 4.5 markers appear (shouldn't really happen on a
|
||||
# real checkpoint, but exercises the loop's break), the 4.5 markers
|
||||
# win because the source loop pins on the first 4.5 id seen.
|
||||
# sequence index: 0 1 2 3 4 5 6 7 8
|
||||
ai, hidden = self._run([10], [151691, 70, 151692, 151703, 30, 151704, 40])
|
||||
assert ai["tts_token_ids"].tolist() == [30]
|
||||
assert torch.equal(ai["tts_hidden_states"], hidden[5:6])
|
||||
|
||||
def test_2_6_fallback(self) -> None:
|
||||
# No 4.5 markers anywhere -> use 2.6 BOS/EOS pair.
|
||||
# full sequence: [10, 151691, 30, 31, 151692, 40]
|
||||
# 0 1 2 3 4 5
|
||||
# BOS at idx 1 -> slice starts at 2; EOS at idx 4 -> slice ends at 4.
|
||||
ai, hidden = self._run([10], [151691, 30, 31, 151692, 40])
|
||||
assert ai["tts_token_ids"].tolist() == [30, 31]
|
||||
assert torch.equal(ai["tts_hidden_states"], hidden[2:4])
|
||||
buffer, hidden = self._run([10, 11], [151703, 30, 31, 151704, 40])
|
||||
assert buffer["ids"]["tts"] == [30, 31]
|
||||
assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[3:5])
|
||||
|
||||
def test_bos_without_eos_runs_to_end(self) -> None:
|
||||
# When BOS is found but EOS is missing (typical for an in-flight or
|
||||
# truncated decode), the slice should extend to the end of the
|
||||
# hidden-state matrix instead of being dropped.
|
||||
# sequence: [10, 11, 151703, 30, 31]
|
||||
ai, hidden = self._run([10, 11], [151703, 30, 31])
|
||||
assert ai["tts_token_ids"].tolist() == [30, 31]
|
||||
assert torch.equal(ai["tts_hidden_states"], hidden[3:5])
|
||||
buffer, hidden = self._run([10, 11], [151703, 30, 31])
|
||||
assert buffer["ids"]["tts"] == [30, 31]
|
||||
assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[3:5])
|
||||
|
||||
def test_no_tts_markers_omits_slice_keys(self) -> None:
|
||||
# If neither marker pair is present, the bridge should NOT populate
|
||||
# ``tts_token_ids`` / ``tts_hidden_states`` — the talker should fall
|
||||
# through to the dummy path.
|
||||
ai, _ = self._run([10, 11], [20, 21, 22])
|
||||
assert "tts_token_ids" not in ai
|
||||
assert "tts_hidden_states" not in ai
|
||||
buffer, _ = self._run([10, 11], [20, 21, 22])
|
||||
assert "tts" not in buffer["ids"]
|
||||
assert "tts" not in buffer.get("hidden_states", {})
|
||||
|
||||
|
||||
class TestPromptAndMultiModal:
|
||||
@@ -251,14 +279,11 @@ class TestPromptAndMultiModal:
|
||||
)
|
||||
assert out[0]["multi_modal_data"] == mm
|
||||
|
||||
def test_streaming_context_is_accepted_and_ignored(self) -> None:
|
||||
# The framework passes ``streaming_context`` for streaming-capable
|
||||
# pipelines; MiniCPM-o 4.5 is turn-taking so the bridge must accept
|
||||
# but discard it without crashing.
|
||||
def test_internal_streaming_context_is_accepted(self) -> None:
|
||||
hidden = torch.zeros((2, _HIDDEN_DIM))
|
||||
out = llm2tts(
|
||||
[_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden)],
|
||||
prompt=None,
|
||||
streaming_context=object(),
|
||||
_streaming_context=object(),
|
||||
)
|
||||
assert len(out) == 1
|
||||
@@ -44,6 +44,20 @@ class TestRegistryDeclaration:
|
||||
assert pipeline.model_type == _PIPELINE_KEY
|
||||
assert pipeline.model_arch == "MiniCPMO45OmniForConditionalGeneration"
|
||||
|
||||
def test_native_duplex_control_is_explicit_without_a_fixed_session_cap(self) -> None:
|
||||
pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
|
||||
assert pipeline.duplex_control_enabled is True
|
||||
assert pipeline.duplex_serving_adapter == (
|
||||
"vllm_omni.experimental.fullduplex.minicpmo45.serving_adapter.MiniCPMO45ServingRuntimeAdapter"
|
||||
)
|
||||
assert not hasattr(pipeline, "max_native_duplex_sessions")
|
||||
|
||||
def test_ordinary_pipeline_defaults_to_no_duplex_control(self) -> None:
|
||||
pipeline = PipelineConfig(model_type="ordinary")
|
||||
assert pipeline.duplex_control_enabled is False
|
||||
assert pipeline.duplex_serving_adapter is None
|
||||
assert not hasattr(pipeline, "max_native_duplex_sessions")
|
||||
|
||||
|
||||
class TestPipelineTopology:
|
||||
@pytest.fixture(scope="class")
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -131,6 +131,18 @@ class TestStageConfig:
|
||||
assert "max_num_seqs" not in omega_config.engine_args
|
||||
# Legacy field name for backward compatibility
|
||||
assert omega_config.engine_input_source == []
|
||||
assert omega_config.session_mode == "turn"
|
||||
|
||||
def test_to_omegaconf_duplex_session_mode(self):
|
||||
"""Test that session mode is preserved as stage metadata."""
|
||||
config = StageConfig(
|
||||
stage_id=0,
|
||||
model_stage="thinker",
|
||||
session_mode="duplex",
|
||||
)
|
||||
omega_config = config.to_omegaconf()
|
||||
|
||||
assert omega_config.session_mode == "duplex"
|
||||
|
||||
def test_to_omegaconf_with_runtime_overrides(self):
|
||||
"""Test that runtime overrides are applied to OmegaConf output."""
|
||||
@@ -1107,6 +1119,57 @@ class TestResolveScheduler:
|
||||
|
||||
|
||||
class TestDeployConfigLoading:
|
||||
def test_load_minicpmo_duplex_deploy_config(self):
|
||||
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "minicpmo_4_5_duplex.yaml"
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert deploy.async_chunk is False
|
||||
assert deploy.active_stream_window == 1
|
||||
assert [stage.session_mode for stage in stages] == ["duplex", "duplex"]
|
||||
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex"]
|
||||
assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False]
|
||||
assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
|
||||
assert deploy.stages[1].enforce_eager is False
|
||||
assert deploy.stages[1].compilation_config == {"cudagraph_mode": "PIECEWISE"}
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [151645]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "stage0_devices", "stage1_devices", "stage1_replicas"),
|
||||
[
|
||||
("minicpmo_4_5_3gpu_stage1_replicas.yaml", "0", "1,2", 2),
|
||||
("minicpmo_4_5_4gpu_stage1_replicas.yaml", "0", "1,2,3", 3),
|
||||
("minicpmo_4_5_8x4090_stage1_replicas.yaml", "0,1,2,3", "4,5,6,7", 4),
|
||||
],
|
||||
)
|
||||
def test_load_minicpmo_replica_deploy_configs(
|
||||
self,
|
||||
filename: str,
|
||||
stage0_devices: str,
|
||||
stage1_devices: str,
|
||||
stage1_replicas: int,
|
||||
):
|
||||
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / filename
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert [stage.yaml_runtime["devices"] for stage in stages] == [stage0_devices, stage1_devices]
|
||||
assert stages[0].yaml_runtime["num_replicas"] == 1
|
||||
assert stages[1].yaml_runtime["num_replicas"] == stage1_replicas
|
||||
assert all(stage.yaml_engine_args["async_scheduling"] is False for stage in stages)
|
||||
assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
|
||||
|
||||
def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path):
|
||||
deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml"
|
||||
custom_voice_dir = tmp_path / "voices"
|
||||
|
||||
@@ -21,6 +21,8 @@ from vllm_omni.data_entry_keys import (
|
||||
)
|
||||
from vllm_omni.engine import AdditionalInformationPayload
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestOmniPayloadStruct:
|
||||
"""Runtime-validated mirror of OmniPayload (msgspec.Struct)."""
|
||||
@@ -43,6 +45,16 @@ class TestOmniPayloadStruct:
|
||||
with pytest.raises(msgspec.ValidationError, match="unknown field"):
|
||||
to_struct({"meta": {"finisheed": True}})
|
||||
|
||||
def test_to_struct_rejects_minicpmo_duplex_runtime_handoff_fields(self):
|
||||
with pytest.raises(msgspec.ValidationError, match="unknown field"):
|
||||
to_struct({"hidden_states": {"tts": torch.zeros(1)}})
|
||||
with pytest.raises(msgspec.ValidationError, match="unknown field"):
|
||||
to_struct({"ids": {"tts": [1]}})
|
||||
with pytest.raises(msgspec.ValidationError, match="unknown field"):
|
||||
to_struct({"meta": {"ref_audio_sr": 16000}})
|
||||
with pytest.raises(msgspec.ValidationError, match="unknown field"):
|
||||
to_struct({"meta": {"listen_token_id": 151685}})
|
||||
|
||||
def test_to_struct_rejects_wrong_type(self):
|
||||
with pytest.raises(msgspec.ValidationError, match="Expected"):
|
||||
to_struct({"meta": {"left_context_size": "not_an_int"}})
|
||||
@@ -66,10 +78,10 @@ class TestOmniPayloadStruct:
|
||||
|
||||
def test_struct_with_all_categories(self):
|
||||
d = {
|
||||
"hidden_states": {"output": torch.zeros(1)},
|
||||
"hidden_states": {"output": torch.zeros(1), "last": torch.ones(2, 4)},
|
||||
"embed": {"prefill": torch.zeros(1), "tts_bos": torch.zeros(1)},
|
||||
"ids": {"all": [1, 2], "prompt": [1]},
|
||||
"codes": {"audio": torch.zeros(1)},
|
||||
"ids": {"all": [1, 2], "prompt": [1], "speech_token": [3, 4]},
|
||||
"codes": {"audio": torch.zeros(1), "ref": torch.ones(4)},
|
||||
"meta": {"left_context_size": 3, "num_processed_tokens": 7},
|
||||
}
|
||||
s = to_struct(d)
|
||||
@@ -79,7 +91,10 @@ class TestOmniPayloadStruct:
|
||||
assert isinstance(s.codes, CodesStruct)
|
||||
assert isinstance(s.meta, MetaStruct)
|
||||
assert s.ids.all == [1, 2]
|
||||
assert s.ids.speech_token == [3, 4]
|
||||
assert torch.equal(s.hidden_states.last, torch.ones(2, 4))
|
||||
assert s.meta.num_processed_tokens == 7
|
||||
assert torch.equal(s.codes.ref, torch.ones(4))
|
||||
|
||||
|
||||
class TestValidatePayload:
|
||||
@@ -332,7 +347,7 @@ class TestSerializeDeserializePayload:
|
||||
"hidden_states": {"output": torch.tensor([1.0, 2.0])},
|
||||
"ids": {"all": [1, 2, 3]},
|
||||
"meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 4},
|
||||
"codes": {"audio": torch.tensor([3.0])},
|
||||
"codes": {"audio": torch.tensor([3.0]), "ref": torch.tensor([0.1, 0.2])},
|
||||
}
|
||||
wire = serialize_payload(original)
|
||||
restored = deserialize_payload(wire)
|
||||
@@ -341,6 +356,7 @@ class TestSerializeDeserializePayload:
|
||||
assert restored["meta"]["finished"].item() is False
|
||||
assert restored["meta"]["ar_width"] == 4
|
||||
assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
|
||||
assert torch.equal(restored["codes"]["ref"], original["codes"]["ref"])
|
||||
|
||||
def test_hidden_states_layers_round_trip(self):
|
||||
original = {
|
||||
|
||||
@@ -20,7 +20,7 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.request import Request
|
||||
|
||||
from vllm_omni.engine import PromptEmbedsPayload
|
||||
from vllm_omni.request import OmniRequest
|
||||
from vllm_omni.request import OmniRequest, OmniStreamingUpdate
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -32,7 +32,12 @@ def test_omni_params_are_keyword_only():
|
||||
came first positionally and broke positional ``Request(...)`` construction.
|
||||
"""
|
||||
params = inspect.signature(OmniRequest.__init__).parameters
|
||||
for name in ("prompt_embeds", "external_req_id", "additional_information"):
|
||||
for name in (
|
||||
"prompt_embeds",
|
||||
"external_req_id",
|
||||
"additional_information",
|
||||
"model_intermediate_buffer",
|
||||
):
|
||||
assert params[name].kind is inspect.Parameter.KEYWORD_ONLY, name
|
||||
|
||||
|
||||
@@ -87,3 +92,21 @@ def test_keyword_omni_params_round_trip():
|
||||
# The serialized payload is retained and decoded into a tensor on the base.
|
||||
assert req.prompt_embeds_payload is payload
|
||||
assert torch.equal(req.prompt_embeds, torch.from_numpy(arr))
|
||||
|
||||
|
||||
def test_model_intermediate_buffer_round_trips_to_streaming_update():
|
||||
buffer = {"duplex": {"turn_id": 3}}
|
||||
request = OmniRequest(
|
||||
request_id="req-model-buffer",
|
||||
prompt_token_ids=[1, 2],
|
||||
sampling_params=SamplingParams(max_tokens=16),
|
||||
pooling_params=None,
|
||||
resumable=True,
|
||||
model_intermediate_buffer=buffer,
|
||||
)
|
||||
|
||||
update = OmniStreamingUpdate.from_request(request)
|
||||
|
||||
assert update is not None
|
||||
assert request.model_intermediate_buffer is buffer
|
||||
assert update.model_intermediate_buffer is buffer
|
||||
@@ -1,7 +1,10 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.utils.mm_outputs import partition_flat_payload, partition_payload_list
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_partition_thinker_latent_payload():
|
||||
payload = {
|
||||
@@ -42,3 +45,27 @@ def test_partition_payload_list_preserves_request_alignment():
|
||||
inter_list, client_list = partition_payload_list(payloads)
|
||||
assert inter_list == [payloads[0], None]
|
||||
assert client_list == [None, payloads[1]]
|
||||
|
||||
|
||||
def test_partition_duplex_audio_transcript_metadata_to_client_mm():
|
||||
payload = {
|
||||
"model_outputs": torch.zeros(1, 2400),
|
||||
"meta.duplex_epoch": torch.tensor([3], dtype=torch.int32),
|
||||
"meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32),
|
||||
"meta.llm_output_text_utf8": torch.tensor([104, 105], dtype=torch.uint8),
|
||||
"meta.audio_text_total_chars": torch.tensor([2], dtype=torch.int32),
|
||||
"meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32),
|
||||
"meta.turn_end": torch.tensor([1], dtype=torch.int32),
|
||||
"meta.native_duplex_segment_text": "hi",
|
||||
}
|
||||
|
||||
inter, client = partition_flat_payload(payload)
|
||||
|
||||
assert "meta.llm_output_text_utf8" in client
|
||||
assert "meta.audio_text_total_chars" in client
|
||||
assert "meta.duplex_epoch" in client
|
||||
assert "meta.duplex_turn_id" in client
|
||||
assert "meta.tts_is_last_chunk" in client
|
||||
assert "meta.turn_end" in client
|
||||
assert "meta.native_duplex_segment_text" not in client
|
||||
assert "meta.native_duplex_segment_text" in inter
|
||||
File diff suppressed because it is too large.
Load diff
@@ -196,8 +196,7 @@ def _make_runner_for_mimo(req_id="r_mimo"):
|
||||
|
||||
|
||||
def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
|
||||
# `_talker_mtp_forward` calls `current_omni_platform.set_forward_context`,
|
||||
# which would otherwise dispatch to the real device implementation.
|
||||
# `_talker_mtp_forward` dispatches through the active platform.
|
||||
import vllm_omni.worker.gpu_model_runner as mod # Must be the same module that defines OmniGPUModelRunner
|
||||
|
||||
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
||||
@@ -459,28 +458,33 @@ def test_update_intermediate_buffer_skips_unknown_req_id():
|
||||
assert "unknown_req" not in runner.model_intermediate_buffer
|
||||
|
||||
|
||||
def test_maybe_run_batch_preprocess_calls_model_hook():
|
||||
runner = object.__new__(OmniGPUModelRunner)
|
||||
runner.model_intermediate_buffer = {"r1": {"text": ["hello"]}}
|
||||
calls = []
|
||||
def test_streaming_input_update_merges_model_intermediate_buffer():
|
||||
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
||||
runner.model_intermediate_buffer["r1"] = {
|
||||
"duplex": {
|
||||
"session_id": "sid",
|
||||
"seq": 1,
|
||||
}
|
||||
}
|
||||
runner.requests["r1"].additional_information_cpu = runner.model_intermediate_buffer["r1"]
|
||||
new_req_data = SimpleNamespace(
|
||||
model_intermediate_buffer={
|
||||
"duplex": {
|
||||
"session_id": "sid",
|
||||
"seq": 2,
|
||||
"payload": {"type": "audio"},
|
||||
}
|
||||
},
|
||||
additional_information=None,
|
||||
)
|
||||
|
||||
class DummyModel:
|
||||
def preprocess_batch(self, *, req_ids, model_intermediate_buffer, device):
|
||||
calls.append((req_ids, model_intermediate_buffer, device))
|
||||
OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req_data, "r1")
|
||||
|
||||
runner.model = DummyModel()
|
||||
|
||||
OmniGPUModelRunner._maybe_run_batch_preprocess(runner, ["r1"], torch.device("cpu"))
|
||||
|
||||
assert calls == [(["r1"], runner.model_intermediate_buffer, torch.device("cpu"))]
|
||||
|
||||
|
||||
def test_maybe_run_batch_preprocess_skips_missing_hook():
|
||||
runner = object.__new__(OmniGPUModelRunner)
|
||||
runner.model_intermediate_buffer = {}
|
||||
runner.model = object()
|
||||
|
||||
OmniGPUModelRunner._maybe_run_batch_preprocess(runner, ["r1"], torch.device("cpu"))
|
||||
info = runner.model_intermediate_buffer["r1"]
|
||||
assert info["duplex"]["session_id"] == "sid"
|
||||
assert info["duplex"]["seq"] == 2
|
||||
assert info["duplex"]["payload"] == {"type": "audio"}
|
||||
assert runner.requests["r1"].additional_information_cpu is info
|
||||
|
||||
|
||||
def _make_full_payload_accumulation_runner(
|
||||
@@ -501,8 +505,6 @@ def _make_full_payload_accumulation_runner(
|
||||
runner._custom_process_func = object()
|
||||
runner._pending_full_payload_send = {}
|
||||
runner._stage_id = 1
|
||||
# Non-None sentinel: the gate short-circuits to False when no connector
|
||||
# is configured at all (terminal stages in pipelines with no connector).
|
||||
runner._omni_connector = object()
|
||||
return runner
|
||||
|
||||
@@ -519,7 +521,7 @@ def test_accumulate_full_payload_output_preserves_aligned_all_zero_qwen3_omni_co
|
||||
|
||||
|
||||
def test_accumulate_full_payload_output_keeps_misaligned_all_zero_qwen3_omni_codec_rows():
|
||||
# After removing the sender-side zero filter, the full-payload accumulator keeps every
|
||||
# After removing the sender-side zero filter, the accumulator keeps every
|
||||
# codec row including misaligned all-zero rows. The downstream consumer
|
||||
# (_extract_qwen3_full_payload_codec_rows) is the authoritative crop and
|
||||
# filters by output_token_ids.
|
||||
@@ -566,45 +568,25 @@ def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeh
|
||||
|
||||
|
||||
def test_full_payload_output_accumulation_hook_matrix():
|
||||
"""Producer-side gate: fires iff an explicit next-stage payload hook is loaded.
|
||||
|
||||
A derived `*_full_payload` helper from `custom_process_input_func` is not
|
||||
enough: terminal/input-only consumer stages must not enqueue orphan
|
||||
downstream payloads.
|
||||
"""
|
||||
# Thinker / talker producer stages: explicit next-stage payload hook -> gate fires.
|
||||
assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output()
|
||||
assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output()
|
||||
|
||||
# Terminal stage: even if _load_custom_func derived a builder from
|
||||
# custom_process_input_func, final output stages are not producers.
|
||||
runner = _make_full_payload_accumulation_runner(model_stage="code2wav", final_output=True)
|
||||
assert not runner._should_accumulate_full_payload_output()
|
||||
|
||||
# Input-only consumer stage without an explicit producer hook must not
|
||||
# accumulate/send just because a same-module *_full_payload helper exists.
|
||||
runner = _make_full_payload_accumulation_runner(
|
||||
assert not _make_full_payload_accumulation_runner(
|
||||
model_stage="code2wav", final_output=True
|
||||
)._should_accumulate_full_payload_output()
|
||||
assert not _make_full_payload_accumulation_runner(
|
||||
model_stage="token2audio",
|
||||
custom_process_next_stage_input_func=None,
|
||||
)
|
||||
assert not runner._should_accumulate_full_payload_output()
|
||||
|
||||
# async_chunk mode -> gate off.
|
||||
)._should_accumulate_full_payload_output()
|
||||
assert not _make_full_payload_accumulation_runner(
|
||||
model_stage="talker", async_chunk=True
|
||||
)._should_accumulate_full_payload_output()
|
||||
|
||||
# Non-qwen3 arches: gate is arch-agnostic, but if the fixture's arch
|
||||
# does not configure a connector payload builder, its runtime
|
||||
# `_custom_process_func` is None. Emulate that.
|
||||
runner = _make_full_payload_accumulation_runner(model_arch="Qwen3TTSForConditionalGeneration")
|
||||
runner._custom_process_func = None
|
||||
runner._should_accumulate_full_payload_output_cached = None
|
||||
assert not runner._should_accumulate_full_payload_output()
|
||||
runner = _make_full_payload_accumulation_runner(model_arch="Qwen2_5OmniForConditionalGeneration")
|
||||
runner._custom_process_func = None
|
||||
runner._should_accumulate_full_payload_output_cached = None
|
||||
assert not runner._should_accumulate_full_payload_output()
|
||||
for model_arch in (
|
||||
"Qwen3TTSForConditionalGeneration",
|
||||
"Qwen2_5OmniForConditionalGeneration",
|
||||
):
|
||||
runner = _make_full_payload_accumulation_runner(model_arch=model_arch)
|
||||
runner._custom_process_func = None
|
||||
assert not runner._should_accumulate_full_payload_output()
|
||||
|
||||
|
||||
def test_sync_local_stage_payloads_retains_payload_until_request_is_active():
|
||||
|
||||
@@ -270,6 +270,15 @@ class PipelineConfig:
|
||||
# multi-component repos (e.g. GLM-Image). ``None`` = not a diffusers model.
|
||||
diffusers_class_name: str | None = None
|
||||
endpoint_restrictions: tuple[EndpointRestriction, ...] = ()
|
||||
# Optional model-owned duplex planner loaded by the stable engine runtime.
|
||||
duplex_runtime_extension: str | None = None
|
||||
# Optional model-owned Serving adapter loaded only when the Realtime duplex
|
||||
# endpoint is enabled. Generic OpenAI modules must not select a model.
|
||||
duplex_serving_adapter: str | None = None
|
||||
# Explicitly enable the stable duplex control mechanism. This is separate
|
||||
# from the optional model extension because turn-commit-only deployments
|
||||
# do not require a model planner.
|
||||
duplex_control_enabled: bool = False
|
||||
# Bundled deploy defaults for this concrete pipeline topology. The file is
|
||||
# loaded from vllm_omni/deploy; None uses DeployConfig defaults.
|
||||
default_deploy_config_name: str | None = None
|
||||
@@ -406,6 +415,38 @@ class StageDeployConfig:
|
||||
engine_extras: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexSessionRuntimeConfig:
|
||||
"""Server-owned lifecycle and per-session buffering limits."""
|
||||
|
||||
idle_ttl_s: float | None = 300.0
|
||||
disconnect_grace_s: float = 30.0
|
||||
reaper_interval_s: float = 5.0
|
||||
resume_replay_ttl_s: float = 60.0
|
||||
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_sessions: int = 1
|
||||
completed_append_cache_size: int = 256
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
positive = {
|
||||
"disconnect_grace_s": self.disconnect_grace_s,
|
||||
"reaper_interval_s": self.reaper_interval_s,
|
||||
"resume_replay_ttl_s": self.resume_replay_ttl_s,
|
||||
"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_sessions": self.max_sessions,
|
||||
"completed_append_cache_size": self.completed_append_cache_size,
|
||||
}
|
||||
if self.idle_ttl_s is not None and self.idle_ttl_s <= 0:
|
||||
raise ValueError("duplex_session.idle_ttl_s must be positive or null")
|
||||
for name, value in positive.items():
|
||||
if value <= 0:
|
||||
raise ValueError(f"duplex_session.{name} must be positive")
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeployConfig:
|
||||
"""Loaded from deploy/<model>.yaml — the only config file users edit.
|
||||
@@ -419,8 +460,10 @@ class DeployConfig:
|
||||
"""
|
||||
|
||||
async_chunk: bool = True
|
||||
session_mode: str = "turn"
|
||||
# Stage-1 active stream slots; 0 preserves legacy all-stream cycling.
|
||||
active_stream_window: int = 0
|
||||
duplex_session: DuplexSessionRuntimeConfig = field(default_factory=DuplexSessionRuntimeConfig)
|
||||
connectors: dict[str, Any] | None = None
|
||||
edges: list[dict[str, Any]] | None = None
|
||||
stages: list[StageDeployConfig] = field(default_factory=list)
|
||||
@@ -610,7 +653,9 @@ def load_deploy_config(path: str | Path) -> DeployConfig:
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"async_chunk": raw_dict.get("async_chunk", True),
|
||||
"session_mode": raw_dict.get("session_mode", "turn"),
|
||||
"active_stream_window": int(raw_dict.get("active_stream_window", 0) or 0),
|
||||
"duplex_session": DuplexSessionRuntimeConfig(**(raw_dict.get("duplex_session") or {})),
|
||||
"connectors": raw_dict.get("connectors", None),
|
||||
"edges": raw_dict.get("edges", None),
|
||||
"stages": stages,
|
||||
@@ -896,6 +941,7 @@ def merge_pipeline_deploy(
|
||||
StageConfig(
|
||||
stage_id=ps.stage_id,
|
||||
model_stage=ps.model_stage,
|
||||
session_mode=deploy.session_mode,
|
||||
stage_type=stage_type,
|
||||
input_sources=list(ps.input_sources),
|
||||
custom_process_input_func=input_proc,
|
||||
@@ -922,6 +968,7 @@ class StageConfig:
|
||||
|
||||
stage_id: int
|
||||
model_stage: str
|
||||
session_mode: str = "turn"
|
||||
stage_type: StageType = StageType.LLM
|
||||
input_sources: list[int] = field(default_factory=list)
|
||||
custom_process_input_func: str | None = None
|
||||
@@ -984,6 +1031,7 @@ class StageConfig:
|
||||
config_dict: dict[str, Any] = {
|
||||
"stage_id": self.stage_id,
|
||||
"stage_type": StageType(self.stage_type).value,
|
||||
"session_mode": self.session_mode,
|
||||
"engine_args": create_config(engine_args),
|
||||
"runtime": create_config(runtime),
|
||||
"engine_input_source": self.input_sources, # Legacy field name
|
||||
|
||||
@@ -2,9 +2,28 @@
|
||||
Scheduling components for vLLM-Omni.
|
||||
"""
|
||||
|
||||
from .omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
|
||||
from .omni_generation_scheduler import OmniGenerationScheduler
|
||||
from .output import OmniNewRequestData
|
||||
from __future__ import annotations
|
||||
|
||||
from importlib import import_module
|
||||
from typing import Any
|
||||
|
||||
_LAZY_EXPORTS = {
|
||||
"OmniARAsyncScheduler": (".omni_ar_scheduler", "OmniARAsyncScheduler"),
|
||||
"OmniARScheduler": (".omni_ar_scheduler", "OmniARScheduler"),
|
||||
"OmniGenerationScheduler": (".omni_generation_scheduler", "OmniGenerationScheduler"),
|
||||
"OmniNewRequestData": (".output", "OmniNewRequestData"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
target = _LAZY_EXPORTS.get(name)
|
||||
if target is None:
|
||||
raise AttributeError(name)
|
||||
module_name, attribute_name = target
|
||||
value = getattr(import_module(module_name, __name__), attribute_name)
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
|
||||
__all__ = [
|
||||
"OmniARAsyncScheduler",
|
||||
|
||||
@@ -259,6 +259,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
|
||||
prompt_is_token_ids=nr.prompt_is_token_ids,
|
||||
additional_information=(getattr(request, "additional_information", None) if request else None),
|
||||
model_intermediate_buffer=(
|
||||
getattr(request, "model_intermediate_buffer", None) if request else None
|
||||
),
|
||||
)
|
||||
new_list.append(omni_nr)
|
||||
|
||||
@@ -718,6 +721,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
session.record_event(EngineCoreEventType.QUEUED)
|
||||
return
|
||||
super()._update_request_as_session(session, update)
|
||||
if hasattr(update, "model_intermediate_buffer"):
|
||||
session.model_intermediate_buffer = update.model_intermediate_buffer
|
||||
|
||||
def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None:
|
||||
# TODO(wzliu)! for offline mode, we should not end process until all data is transferred
|
||||
|
||||
@@ -319,6 +319,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
|
||||
prompt_is_token_ids=nr.prompt_is_token_ids,
|
||||
additional_information=(getattr(request, "additional_information", None) if request else None),
|
||||
model_intermediate_buffer=(
|
||||
getattr(request, "model_intermediate_buffer", None) if request else None
|
||||
),
|
||||
)
|
||||
new_list.append(omni_nr)
|
||||
|
||||
|
||||
@@ -40,16 +40,16 @@ except ValueError:
|
||||
class OmniSchedulerMixin:
|
||||
"""Shared scheduler helpers for omni-specific request handling."""
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Shared scheduler/output helpers (lift the AR / generation duplicates)
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _free_input_coordinator_request(self, request_id: str) -> None:
|
||||
"""Prune full-payload coordinator state for a completed request."""
|
||||
input_coordinator = getattr(self, "input_coordinator", None)
|
||||
if input_coordinator is not None:
|
||||
input_coordinator.free_finished_request(request_id)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Shared scheduler/output helpers (lift the AR / generation duplicates)
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _consume_pending_connector_output(self, model_mode: str) -> None:
|
||||
"""Drain ``self._latest_omni_connector_output`` into the coordinator.
|
||||
|
||||
|
||||
@@ -20,10 +20,13 @@ class OmniNewRequestData(NewRequestData):
|
||||
external_req_id: Optional external request ID for tracking
|
||||
additional_information: Optional serialized additional information
|
||||
dictionary containing tensors or lists
|
||||
model_intermediate_buffer: Optional runner-owned payload for
|
||||
GPUModelRunner.model_intermediate_buffer
|
||||
"""
|
||||
|
||||
external_req_id: str | None = None
|
||||
additional_information: AdditionalInformationPayload | None = None
|
||||
model_intermediate_buffer: dict[str, object] | None = None
|
||||
|
||||
@classmethod
|
||||
def from_request(
|
||||
@@ -56,6 +59,7 @@ class OmniNewRequestData(NewRequestData):
|
||||
prompt_is_token_ids=getattr(request, "prompt_is_token_ids", None),
|
||||
prefill_token_ids=prefill_token_ids,
|
||||
additional_information=getattr(request, "additional_information", None),
|
||||
model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
# Experimental validation profile: one thinker and two Stage1 replicas.
|
||||
# This is not a recommended production entrypoint.
|
||||
#
|
||||
# Stage 0: GPU 0
|
||||
# Stage 1: replicas on GPU 1 and 2
|
||||
base_config: minicpmo_4_5.yaml
|
||||
pipeline: minicpmo_4_5
|
||||
session_mode: duplex
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
async_scheduling: false
|
||||
|
||||
- stage_id: 1
|
||||
devices: "1,2"
|
||||
num_replicas: 2
|
||||
async_scheduling: false
|
||||
@@ -0,0 +1,17 @@
|
||||
# Experimental validation profile: one thinker and three Stage1 replicas.
|
||||
# This is not a recommended production entrypoint.
|
||||
#
|
||||
# Stage 0: GPU 0
|
||||
# Stage 1: replicas on GPU 1, 2, and 3
|
||||
base_config: minicpmo_4_5.yaml
|
||||
pipeline: minicpmo_4_5
|
||||
session_mode: duplex
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
async_scheduling: false
|
||||
|
||||
- stage_id: 1
|
||||
devices: "1,2,3"
|
||||
num_replicas: 3
|
||||
async_scheduling: false
|
||||
@@ -0,0 +1,20 @@
|
||||
# Experimental validation profile: 4-way thinker TP and four Stage1 replicas.
|
||||
# This is not a recommended production entrypoint.
|
||||
#
|
||||
# Stage 0: GPU 0, 1, 2, and 3
|
||||
# Stage 1: replicas on GPU 4, 5, 6, and 7
|
||||
base_config: minicpmo_4_5_8x4090.yaml
|
||||
pipeline: minicpmo_4_5
|
||||
session_mode: duplex
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
async_scheduling: false
|
||||
gpu_memory_utilization: 0.40
|
||||
max_num_batched_tokens: 16384
|
||||
|
||||
- stage_id: 1
|
||||
devices: "4,5,6,7"
|
||||
num_replicas: 4
|
||||
async_scheduling: false
|
||||
gpu_memory_utilization: 0.8
|
||||
@@ -0,0 +1,36 @@
|
||||
# MiniCPM-o 4.5 native duplex deployment.
|
||||
#
|
||||
# The registered minicpmo_4_5 pipeline owns the fixed Stage0 -> Stage1
|
||||
# topology and llm2tts bridge. This overlay only changes session lifetime and
|
||||
# the Stage1 token budget needed to drain streaming audio segments.
|
||||
base_config: minicpmo_4_5.yaml
|
||||
pipeline: minicpmo_4_5
|
||||
session_mode: duplex
|
||||
active_stream_window: 1
|
||||
duplex_session:
|
||||
idle_ttl_s: 300
|
||||
disconnect_grace_s: 30
|
||||
reaper_interval_s: 5
|
||||
resume_replay_ttl_s: 60
|
||||
resume_replay_max_bytes_per_session: 8388608
|
||||
max_pending_input_bytes_per_session: 16777216
|
||||
max_pending_turns_per_session: 4
|
||||
# Matches the configured Stage0/Stage1 sequence capacity below.
|
||||
max_sessions: 2
|
||||
completed_append_cache_size: 256
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 2
|
||||
async_scheduling: false
|
||||
|
||||
- stage_id: 1
|
||||
max_num_seqs: 2
|
||||
async_scheduling: false
|
||||
# The streaming talker advances Python-owned per-request generators on
|
||||
# every step, so only fixed-topology model regions are graph-capturable.
|
||||
compilation_config:
|
||||
cudagraph_mode: PIECEWISE
|
||||
default_sampling_params:
|
||||
max_tokens: 4096
|
||||
stop_token_ids: [151645]
|
||||
@@ -771,8 +771,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
self.waiting_for_chunk_running_requests = deque(
|
||||
request for request in self.waiting_for_chunk_running_requests if request.request_id not in request_ids
|
||||
)
|
||||
self._held_non_active = deque(
|
||||
request for request in self._held_non_active if request.request_id not in request_ids
|
||||
)
|
||||
|
||||
for req_id in request_ids:
|
||||
self._active_streams.pop(req_id, None)
|
||||
self.requests_with_ready_chunks.discard(req_id)
|
||||
self.finished_requests.discard(req_id)
|
||||
self._finished_load_reqs.discard(req_id)
|
||||
|
||||
@@ -75,6 +75,10 @@ class OmniEngineCoreRequest(EngineCoreRequest):
|
||||
|
||||
# Optional additional information dictionary (serialized)
|
||||
additional_information: AdditionalInformationPayload | None = None
|
||||
# Runner-owned runtime payload. This is materialized directly into
|
||||
# GPUModelRunner.model_intermediate_buffer instead of using the deprecated
|
||||
# additional_information request transport.
|
||||
model_intermediate_buffer: dict[str, Any] | None = None
|
||||
|
||||
@classmethod
|
||||
def from_request(
|
||||
@@ -83,6 +87,7 @@ class OmniEngineCoreRequest(EngineCoreRequest):
|
||||
*,
|
||||
prompt_embeds: torch.Tensor | None = None,
|
||||
additional_information: AdditionalInformationPayload | None = None,
|
||||
model_intermediate_buffer: dict[str, Any] | None = None,
|
||||
) -> "OmniEngineCoreRequest":
|
||||
"""Clone an EngineCoreRequest into an OmniEngineCoreRequest with optional payload overrides."""
|
||||
|
||||
@@ -90,6 +95,8 @@ class OmniEngineCoreRequest(EngineCoreRequest):
|
||||
prompt_embeds = request.prompt_embeds
|
||||
if additional_information is None:
|
||||
additional_information = getattr(request, "additional_information", None)
|
||||
if model_intermediate_buffer is None:
|
||||
model_intermediate_buffer = getattr(request, "model_intermediate_buffer", None)
|
||||
|
||||
return cls(
|
||||
request_id=request.request_id,
|
||||
@@ -113,6 +120,7 @@ class OmniEngineCoreRequest(EngineCoreRequest):
|
||||
reasoning_parser_kwargs=request.reasoning_parser_kwargs,
|
||||
abort_immediately=request.abort_immediately,
|
||||
additional_information=additional_information,
|
||||
model_intermediate_buffer=model_intermediate_buffer,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
"""Stateless request and shutdown helpers for :mod:`async_omni_engine`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
import janus
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
|
||||
from vllm_omni.engine import OmniEngineCoreRequest
|
||||
from vllm_omni.engine.messages import EngineQueueMessage, ShutdownRequestMessage
|
||||
from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient
|
||||
from vllm_omni.engine.serialization import (
|
||||
deserialize_additional_information,
|
||||
serialize_additional_information,
|
||||
)
|
||||
from vllm_omni.engine.stage_runtime import StageRuntime
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
SHUTDOWN_ENQUEUE_TIMEOUT_S = 1.0
|
||||
SHUTDOWN_JOIN_TIMEOUT_S = 30.0
|
||||
_WEAK_SHUTDOWN_JOIN_TIMEOUT_S = 1.0
|
||||
|
||||
|
||||
def inject_global_id(target: Any, request_id: str) -> None:
|
||||
"""Inject ``global_request_id`` into a prompt's additional information."""
|
||||
if isinstance(target, dict):
|
||||
if "additional_information" not in target:
|
||||
target["additional_information"] = {}
|
||||
if target["additional_information"] is None:
|
||||
target["additional_information"] = {}
|
||||
if isinstance(target["additional_information"], dict):
|
||||
target["additional_information"]["global_request_id"] = [str(request_id)]
|
||||
|
||||
|
||||
def upgrade_to_omni_request(
|
||||
request: EngineCoreRequest,
|
||||
raw_prompt: Any,
|
||||
) -> EngineCoreRequest:
|
||||
"""Restore omni-only fields omitted by the upstream input processor."""
|
||||
prompt_embeds = request.prompt_embeds
|
||||
additional_information = None
|
||||
wire_payload = None
|
||||
model_intermediate_buffer = None
|
||||
|
||||
if isinstance(raw_prompt, dict):
|
||||
if prompt_embeds is None:
|
||||
raw_prompt_embeds = raw_prompt.get("prompt_embeds")
|
||||
if isinstance(raw_prompt_embeds, torch.Tensor):
|
||||
prompt_embeds = raw_prompt_embeds
|
||||
raw_info = raw_prompt.get("additional_information")
|
||||
raw_buffer = raw_prompt.get("model_intermediate_buffer")
|
||||
if isinstance(raw_info, dict):
|
||||
wire_payload = dict(raw_info)
|
||||
if isinstance(raw_buffer, dict):
|
||||
model_intermediate_buffer = raw_buffer
|
||||
additional_information = serialize_additional_information(
|
||||
wire_payload,
|
||||
log_prefix="AsyncOmniEngine",
|
||||
)
|
||||
|
||||
if prompt_embeds is None and additional_information is None and model_intermediate_buffer is None:
|
||||
return request
|
||||
|
||||
return OmniEngineCoreRequest.from_request(
|
||||
request,
|
||||
prompt_embeds=prompt_embeds,
|
||||
additional_information=additional_information,
|
||||
model_intermediate_buffer=model_intermediate_buffer,
|
||||
)
|
||||
|
||||
|
||||
def apply_omni_final_stage_metadata(
|
||||
request: EngineCoreRequest,
|
||||
final_stage_id: int,
|
||||
) -> EngineCoreRequest:
|
||||
"""Tag a request so the AR scheduler can skip DiT KV for stage zero."""
|
||||
merged: dict[str, Any] = {}
|
||||
if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None:
|
||||
merged = deserialize_additional_information(request.additional_information)
|
||||
merged["omni_final_stage_id"] = final_stage_id
|
||||
payload = serialize_additional_information(merged)
|
||||
return OmniEngineCoreRequest.from_request(
|
||||
request,
|
||||
additional_information=payload,
|
||||
)
|
||||
|
||||
|
||||
def weak_shutdown_async_omni_engine(
|
||||
orchestrator_thread: threading.Thread | None,
|
||||
request_queue: janus.Queue[EngineQueueMessage] | None,
|
||||
output_queue: janus.Queue[EngineQueueMessage] | None,
|
||||
rpc_output_queue: janus.Queue[EngineQueueMessage] | None,
|
||||
rpc_client: CorrelatedRpcClient | None,
|
||||
) -> None:
|
||||
"""Best-effort orchestrator cleanup for garbage-collection finalization."""
|
||||
request_queue_closed = False
|
||||
shutdown_enqueued = enqueue_orchestrator_shutdown(
|
||||
request_queue,
|
||||
timeout=SHUTDOWN_ENQUEUE_TIMEOUT_S,
|
||||
)
|
||||
if request_queue is not None and not shutdown_enqueued:
|
||||
try:
|
||||
request_queue.close()
|
||||
request_queue_closed = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if rpc_client is not None:
|
||||
rpc_client.close()
|
||||
|
||||
try:
|
||||
if orchestrator_thread is not None and orchestrator_thread.is_alive():
|
||||
orchestrator_thread.join(timeout=_WEAK_SHUTDOWN_JOIN_TIMEOUT_S)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for q in (request_queue, output_queue, rpc_output_queue):
|
||||
try:
|
||||
if q is not None and not (q is request_queue and request_queue_closed):
|
||||
q.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def enqueue_orchestrator_shutdown(
|
||||
request_queue: janus.Queue[EngineQueueMessage] | None,
|
||||
*,
|
||||
timeout: float,
|
||||
) -> bool:
|
||||
"""Deliver shutdown without waiting forever behind a full request queue."""
|
||||
if request_queue is None:
|
||||
return False
|
||||
try:
|
||||
request_queue.sync_q.put(ShutdownRequestMessage(), timeout=timeout)
|
||||
except Exception:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def shutdown_runtime_after_orchestrator(
|
||||
orchestrator_thread: threading.Thread,
|
||||
runtime: StageRuntime,
|
||||
) -> None:
|
||||
"""Release a stage runtime once its orchestrator can no longer use it."""
|
||||
try:
|
||||
orchestrator_thread.join()
|
||||
except Exception:
|
||||
logger.exception("[AsyncOmniEngine] Failed while awaiting deferred Orchestrator shutdown")
|
||||
if orchestrator_thread.is_alive():
|
||||
logger.error("[AsyncOmniEngine] Orchestrator is still alive; StageRuntime cleanup remains deferred")
|
||||
return
|
||||
try:
|
||||
runtime.shutdown()
|
||||
except Exception:
|
||||
logger.exception("[AsyncOmniEngine] Failed to shutdown deferred StageRuntime")
|
||||
@@ -18,7 +18,7 @@ import uuid
|
||||
import weakref
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import asdict
|
||||
from typing import Any, Literal, cast
|
||||
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
|
||||
import janus
|
||||
import torch
|
||||
@@ -31,10 +31,23 @@ from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.input_processor import InputProcessor
|
||||
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
from vllm_omni.config.stage_config import strip_parent_engine_args
|
||||
from vllm_omni.config.stage_config import (
|
||||
DuplexSessionRuntimeConfig,
|
||||
load_deploy_config,
|
||||
strip_parent_engine_args,
|
||||
)
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
|
||||
from vllm_omni.diffusion.diffusion_engine import supports_audio_output
|
||||
from vllm_omni.engine import OmniEngineCoreRequest
|
||||
from vllm_omni.engine.async_engine_utils import (
|
||||
SHUTDOWN_ENQUEUE_TIMEOUT_S,
|
||||
SHUTDOWN_JOIN_TIMEOUT_S,
|
||||
apply_omni_final_stage_metadata,
|
||||
enqueue_orchestrator_shutdown,
|
||||
inject_global_id,
|
||||
shutdown_runtime_after_orchestrator,
|
||||
upgrade_to_omni_request,
|
||||
weak_shutdown_async_omni_engine,
|
||||
)
|
||||
from vllm_omni.engine.messages import (
|
||||
AbortRequestMessage,
|
||||
AddCompanionRequestMessage,
|
||||
@@ -42,14 +55,10 @@ from vllm_omni.engine.messages import (
|
||||
CollectiveRPCResultMessage,
|
||||
EngineQueueMessage,
|
||||
ErrorMessage,
|
||||
ShutdownRequestMessage,
|
||||
StageSubmissionMessage,
|
||||
)
|
||||
from vllm_omni.engine.orchestrator import Orchestrator
|
||||
from vllm_omni.engine.serialization import (
|
||||
deserialize_additional_information,
|
||||
serialize_additional_information,
|
||||
)
|
||||
from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient
|
||||
from vllm_omni.engine.stage_client import StageClient
|
||||
from vllm_omni.engine.stage_init_utils import build_stage0_input_processor
|
||||
from vllm_omni.engine.stage_pool import StagePool
|
||||
@@ -67,10 +76,13 @@ from vllm_omni.metrics.prometheus import OmniRequestCounter
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_control_client import DuplexControlClient
|
||||
from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
|
||||
_STARTUP_POLL_INTERVAL_S = 1.0
|
||||
_REQUEST_QUEUE_MAXSIZE = 256
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Parent-EngineArgs field-routing contracts (consumed by
|
||||
# AsyncOmniEngine._strip_parent_engine_args when ``stage_configs_path`` is set).
|
||||
@@ -106,88 +118,6 @@ _PARENT_ARGS_STRIP: frozenset[str] = frozenset({"stage_configs_path"})
|
||||
_PARENT_ARGS_NO_WARN: frozenset[str] = frozenset({"model"})
|
||||
|
||||
|
||||
def _inject_global_id(target: Any, request_id: str) -> None:
|
||||
"""Inject global_request_id into a prompt dict's additional_information."""
|
||||
if isinstance(target, dict):
|
||||
if "additional_information" not in target:
|
||||
target["additional_information"] = {}
|
||||
if target["additional_information"] is None:
|
||||
target["additional_information"] = {}
|
||||
if isinstance(target["additional_information"], dict):
|
||||
target["additional_information"]["global_request_id"] = [str(request_id)]
|
||||
|
||||
|
||||
def _upgrade_to_omni_request(
|
||||
request: EngineCoreRequest,
|
||||
raw_prompt: Any,
|
||||
) -> EngineCoreRequest:
|
||||
"""Restore omni-only fields omitted by upstream InputProcessor."""
|
||||
prompt_embeds = request.prompt_embeds
|
||||
additional_information = None
|
||||
|
||||
if isinstance(raw_prompt, dict):
|
||||
if prompt_embeds is None:
|
||||
raw_prompt_embeds = raw_prompt.get("prompt_embeds")
|
||||
if isinstance(raw_prompt_embeds, torch.Tensor):
|
||||
prompt_embeds = raw_prompt_embeds
|
||||
additional_information = serialize_additional_information(
|
||||
raw_prompt.get("additional_information"),
|
||||
log_prefix="AsyncOmniEngine",
|
||||
)
|
||||
|
||||
if prompt_embeds is None and additional_information is None:
|
||||
return request
|
||||
|
||||
return OmniEngineCoreRequest.from_request(
|
||||
request,
|
||||
prompt_embeds=prompt_embeds,
|
||||
additional_information=additional_information,
|
||||
)
|
||||
|
||||
|
||||
def _apply_omni_final_stage_metadata(
|
||||
request: EngineCoreRequest,
|
||||
final_stage_id: int,
|
||||
) -> EngineCoreRequest:
|
||||
"""Tag EngineCoreRequest so OmniARScheduler can skip DiT KV when final_stage_id is 0."""
|
||||
merged: dict[str, Any] = {}
|
||||
if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None:
|
||||
merged = deserialize_additional_information(request.additional_information)
|
||||
merged["omni_final_stage_id"] = final_stage_id
|
||||
payload = serialize_additional_information(merged)
|
||||
return OmniEngineCoreRequest.from_request(
|
||||
request,
|
||||
additional_information=payload,
|
||||
)
|
||||
|
||||
|
||||
def _weak_shutdown_async_omni_engine(
|
||||
orchestrator_thread: threading.Thread | None,
|
||||
request_queue: janus.Queue[EngineQueueMessage] | None,
|
||||
output_queue: janus.Queue[EngineQueueMessage] | None,
|
||||
rpc_output_queue: janus.Queue[EngineQueueMessage] | None,
|
||||
) -> None:
|
||||
"""Best-effort orchestrator cleanup for GC finalization."""
|
||||
try:
|
||||
if request_queue is not None:
|
||||
request_queue.sync_q.put_nowait(ShutdownRequestMessage())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
if orchestrator_thread is not None and orchestrator_thread.is_alive():
|
||||
orchestrator_thread.join()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for q in (request_queue, output_queue, rpc_output_queue):
|
||||
try:
|
||||
if q is not None:
|
||||
q.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class AsyncOmniEngine:
|
||||
"""Thin proxy that launches an Orchestrator in a background thread.
|
||||
|
||||
@@ -282,11 +212,22 @@ class AsyncOmniEngine:
|
||||
# restriction beforehand. TODO (Alex) make this cleaner and refactor
|
||||
# stage config resolution to remove kwargs hacks.
|
||||
deploy_config_path = kwargs.get("deploy_config")
|
||||
self.endpoint_restrictions = StageConfigFactory.get_pipeline_endpoint_restrictions(
|
||||
pipeline_config = StageConfigFactory.get_pipeline_config(
|
||||
model=model,
|
||||
trust_remote_code=trust_remote_code,
|
||||
deploy_config_path=deploy_config_path,
|
||||
)
|
||||
self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else ()
|
||||
self._duplex_runtime_extension_path = (
|
||||
pipeline_config.duplex_runtime_extension if pipeline_config is not None else None
|
||||
)
|
||||
self.duplex_serving_adapter_path = (
|
||||
pipeline_config.duplex_serving_adapter if pipeline_config is not None else None
|
||||
)
|
||||
self._duplex_control_enabled = bool(pipeline_config and pipeline_config.duplex_control_enabled)
|
||||
self.duplex_session_config = DuplexSessionRuntimeConfig()
|
||||
if deploy_config_path is not None:
|
||||
self.duplex_session_config = load_deploy_config(deploy_config_path).duplex_session
|
||||
|
||||
kwargs["trust_remote_code"] = trust_remote_code
|
||||
self.config_path, self.stage_configs = self._resolve_stage_configs(
|
||||
@@ -315,7 +256,8 @@ class AsyncOmniEngine:
|
||||
self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue()
|
||||
self._shutdown_called = False
|
||||
self._weak_finalizer: weakref.finalize | None = None
|
||||
self._rpc_lock = threading.Lock()
|
||||
self._correlated_rpc_client: CorrelatedRpcClient | None = None
|
||||
self._duplex_control_client: DuplexControlClient | None = None
|
||||
self._running_counter = OmniRequestCounter()
|
||||
|
||||
logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages")
|
||||
@@ -334,15 +276,20 @@ class AsyncOmniEngine:
|
||||
)
|
||||
self.orchestrator_thread.start()
|
||||
self._wait_for_orchestrator_init(startup_future, startup_timeout)
|
||||
self._correlated_rpc_client = CorrelatedRpcClient(
|
||||
self.request_queue.sync_q,
|
||||
self.rpc_output_queue.sync_q,
|
||||
)
|
||||
|
||||
# Stage runtime fields are assigned directly on self by the bootstrap thread.
|
||||
self._weak_finalizer = weakref.finalize(
|
||||
self,
|
||||
_weak_shutdown_async_omni_engine,
|
||||
weak_shutdown_async_omni_engine,
|
||||
self.orchestrator_thread,
|
||||
self.request_queue,
|
||||
self.output_queue,
|
||||
self.rpc_output_queue,
|
||||
self._correlated_rpc_client,
|
||||
)
|
||||
|
||||
logger.info(f"[AsyncOmniEngine] Orchestrator ready with {self.num_stages} stages")
|
||||
@@ -435,6 +382,21 @@ class AsyncOmniEngine:
|
||||
pd_config = self._detect_pd_config()
|
||||
|
||||
membership_controller = self._runtime.create_membership_controller()
|
||||
duplex_runtime_extension = None
|
||||
if self._duplex_control_enabled:
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
|
||||
load_duplex_runtime_extension,
|
||||
validate_duplex_runtime_extension,
|
||||
)
|
||||
|
||||
duplex_runtime_extension = load_duplex_runtime_extension(
|
||||
getattr(self, "_duplex_runtime_extension_path", None)
|
||||
)
|
||||
if duplex_runtime_extension is not None:
|
||||
validate_duplex_runtime_extension(
|
||||
duplex_runtime_extension,
|
||||
sampling_defaults=tuple(pool.stage_client.default_sampling_params for pool in self.stage_pools),
|
||||
)
|
||||
|
||||
orchestrator = Orchestrator(
|
||||
request_async_queue=self.request_queue.async_q,
|
||||
@@ -448,6 +410,9 @@ class AsyncOmniEngine:
|
||||
transfer_emitter=self._transfer_emitter,
|
||||
log_stats=self._log_stats,
|
||||
enable_orch_monitor=self._enable_orch_monitor,
|
||||
duplex_runtime_extension=duplex_runtime_extension,
|
||||
enable_duplex_control=self._duplex_control_enabled,
|
||||
duplex_session_config=self.duplex_session_config,
|
||||
)
|
||||
if not startup_future.done():
|
||||
startup_future.set_result(asyncio.get_running_loop())
|
||||
@@ -718,10 +683,10 @@ class AsyncOmniEngine:
|
||||
if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest):
|
||||
# Inject global_request_id into the raw prompt.
|
||||
if isinstance(prompt, dict):
|
||||
_inject_global_id(prompt, request_id)
|
||||
inject_global_id(prompt, request_id)
|
||||
elif isinstance(prompt, list):
|
||||
for item in prompt:
|
||||
_inject_global_id(item, request_id)
|
||||
inject_global_id(item, request_id)
|
||||
|
||||
preselected_stage0_replica = self._scope_stage0_multimodal_cache_to_replica(
|
||||
request_id,
|
||||
@@ -751,7 +716,7 @@ class AsyncOmniEngine:
|
||||
_preprocess_ms = (time.perf_counter() - _t_preprocess) * 1000.0
|
||||
# TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have
|
||||
# additional_information field in the prompt.
|
||||
request = _upgrade_to_omni_request(request, prompt)
|
||||
request = upgrade_to_omni_request(request, prompt)
|
||||
|
||||
if reasoning_ended is not None:
|
||||
request.reasoning_ended = reasoning_ended
|
||||
@@ -763,7 +728,7 @@ class AsyncOmniEngine:
|
||||
# to match the key used in Orchestrator.request_states so that
|
||||
# output routing (output.request_id lookup) can find the req_state.
|
||||
request.external_req_id = request_id
|
||||
request = _apply_omni_final_stage_metadata(request, final_stage_id)
|
||||
request = apply_omni_final_stage_metadata(request, final_stage_id)
|
||||
|
||||
# Registration with stage 0's output processor is deferred to the
|
||||
# orchestrator thread (see Orchestrator._handle_add_request), which
|
||||
@@ -811,7 +776,7 @@ class AsyncOmniEngine:
|
||||
companion_params, companion_spl = ep.apply_overrides(stage0_params, sampling_params_list)
|
||||
|
||||
if isinstance(companion_prompt, dict):
|
||||
_inject_global_id(companion_prompt, cid)
|
||||
inject_global_id(companion_prompt, cid)
|
||||
|
||||
request = self.input_processor.process_inputs(
|
||||
request_id=cid,
|
||||
@@ -1434,6 +1399,268 @@ class AsyncOmniEngine:
|
||||
resumable=resumable,
|
||||
)
|
||||
|
||||
def open_duplex_session(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
session_mode: str = "duplex",
|
||||
capabilities: dict[str, object] | None = None,
|
||||
session_config: dict[str, object] | None = None,
|
||||
runtime_config: dict[str, object] | None = None,
|
||||
fence: DuplexFence,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Open an engine-level duplex session."""
|
||||
return self._get_duplex_control_client().open(
|
||||
session_id,
|
||||
session_mode=session_mode,
|
||||
capabilities=capabilities,
|
||||
session_config=session_config,
|
||||
runtime_config=runtime_config,
|
||||
fence=fence,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def open_duplex_session_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
session_mode: str = "duplex",
|
||||
capabilities: dict[str, object] | None = None,
|
||||
session_config: dict[str, object] | None = None,
|
||||
runtime_config: dict[str, object] | None = None,
|
||||
fence: DuplexFence,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Async wrapper for opening an engine-level duplex session."""
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.open_duplex_session(
|
||||
session_id,
|
||||
session_mode=session_mode,
|
||||
capabilities=capabilities,
|
||||
session_config=session_config,
|
||||
runtime_config=runtime_config,
|
||||
fence=fence,
|
||||
timeout=timeout,
|
||||
),
|
||||
)
|
||||
|
||||
def append_duplex_input(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
mode: str,
|
||||
payload: object,
|
||||
operation_id: str | None = None,
|
||||
final: bool = False,
|
||||
expected_epoch: int | None = None,
|
||||
fence: DuplexFence,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Append input to an engine-level duplex session."""
|
||||
return self._get_duplex_control_client().append(
|
||||
session_id,
|
||||
mode=mode,
|
||||
payload=payload,
|
||||
operation_id=operation_id,
|
||||
final=final,
|
||||
expected_epoch=expected_epoch,
|
||||
fence=fence,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def append_duplex_input_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
mode: str,
|
||||
payload: object,
|
||||
operation_id: str | None = None,
|
||||
final: bool = False,
|
||||
expected_epoch: int | None = None,
|
||||
fence: DuplexFence,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Async wrapper for appending duplex input."""
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.append_duplex_input(
|
||||
session_id,
|
||||
mode=mode,
|
||||
payload=payload,
|
||||
operation_id=operation_id,
|
||||
final=final,
|
||||
expected_epoch=expected_epoch,
|
||||
fence=fence,
|
||||
timeout=timeout,
|
||||
),
|
||||
)
|
||||
|
||||
def signal_duplex_turn(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
event: str,
|
||||
fence: DuplexFence,
|
||||
next_fence: DuplexFence | None = None,
|
||||
session_config: dict[str, object] | None = None,
|
||||
runtime_config: dict[str, object] | None = None,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Signal an engine-level duplex turn."""
|
||||
return self._get_duplex_control_client().signal(
|
||||
session_id,
|
||||
event=event,
|
||||
fence=fence,
|
||||
next_fence=next_fence,
|
||||
session_config=session_config,
|
||||
runtime_config=runtime_config,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def signal_duplex_turn_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
event: str,
|
||||
fence: DuplexFence,
|
||||
next_fence: DuplexFence | None = None,
|
||||
session_config: dict[str, object] | None = None,
|
||||
runtime_config: dict[str, object] | None = None,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Async wrapper for signaling a duplex turn."""
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.signal_duplex_turn(
|
||||
session_id,
|
||||
event=event,
|
||||
fence=fence,
|
||||
next_fence=next_fence,
|
||||
session_config=session_config,
|
||||
runtime_config=runtime_config,
|
||||
timeout=timeout,
|
||||
),
|
||||
)
|
||||
|
||||
def close_duplex_session(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
reason: str = "client_close",
|
||||
fence: DuplexFence,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Close an engine-level duplex session."""
|
||||
return self._get_duplex_control_client().close(
|
||||
session_id,
|
||||
reason=reason,
|
||||
fence=fence,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def touch_duplex_session(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
activity: DuplexLeaseActivity,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
return self._get_duplex_control_client().touch(
|
||||
session_id,
|
||||
fence=fence,
|
||||
activity=activity,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def touch_duplex_session_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
activity: DuplexLeaseActivity,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.touch_duplex_session(
|
||||
session_id,
|
||||
fence=fence,
|
||||
activity=activity,
|
||||
timeout=timeout,
|
||||
),
|
||||
)
|
||||
|
||||
def resume_duplex_session(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
expected_lease_generation: int,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
return self._get_duplex_control_client().resume(
|
||||
session_id,
|
||||
fence=fence,
|
||||
expected_lease_generation=expected_lease_generation,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def resume_duplex_session_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
expected_lease_generation: int,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.resume_duplex_session(
|
||||
session_id,
|
||||
fence=fence,
|
||||
expected_lease_generation=expected_lease_generation,
|
||||
timeout=timeout,
|
||||
),
|
||||
)
|
||||
|
||||
def _get_duplex_control_client(self) -> DuplexControlClient:
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_control_client import DuplexControlClient
|
||||
|
||||
client = getattr(self, "_duplex_control_client", None)
|
||||
if client is None:
|
||||
transport = getattr(self, "_correlated_rpc_client", None)
|
||||
if transport is None:
|
||||
raise RuntimeError("correlated RPC client is not initialized")
|
||||
client = DuplexControlClient(
|
||||
transport,
|
||||
control_id_factory=lambda: uuid.uuid4().hex,
|
||||
)
|
||||
self._duplex_control_client = client
|
||||
return client
|
||||
|
||||
async def close_duplex_session_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
reason: str = "client_close",
|
||||
fence: DuplexFence,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Async wrapper for closing an engine-level duplex session."""
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.close_duplex_session(session_id, reason=reason, fence=fence, timeout=timeout),
|
||||
)
|
||||
|
||||
def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None:
|
||||
"""Read one output message from the Orchestrator output queue."""
|
||||
try:
|
||||
@@ -1489,36 +1716,19 @@ class AsyncOmniEngine:
|
||||
stage_ids=stage_ids,
|
||||
)
|
||||
|
||||
with self._rpc_lock:
|
||||
self.request_queue.sync_q.put(msg)
|
||||
deadline = None if timeout is None else time.monotonic() + timeout
|
||||
|
||||
while True:
|
||||
remaining = None if deadline is None else max(0.0, deadline - time.monotonic())
|
||||
try:
|
||||
result_msg = self.rpc_output_queue.sync_q.get(timeout=remaining)
|
||||
except queue.Empty as exc:
|
||||
raise TimeoutError(f"collective_rpc timed out after {timeout} seconds") from exc
|
||||
|
||||
if isinstance(result_msg, ErrorMessage):
|
||||
raise RuntimeError(result_msg.error)
|
||||
|
||||
if not isinstance(result_msg, CollectiveRPCResultMessage):
|
||||
logger.warning(
|
||||
"[AsyncOmniEngine] Dropping unexpected rpc queue message type=%s",
|
||||
getattr(result_msg, "type", type(result_msg).__name__),
|
||||
)
|
||||
continue
|
||||
|
||||
if result_msg.rpc_id != rpc_id:
|
||||
logger.warning(
|
||||
"[AsyncOmniEngine] Dropping mismatched rpc result rpc_id=%s expected=%s",
|
||||
result_msg.rpc_id,
|
||||
rpc_id,
|
||||
)
|
||||
continue
|
||||
|
||||
return list(result_msg.results)
|
||||
transport = self._correlated_rpc_client
|
||||
if transport is None:
|
||||
raise RuntimeError("correlated RPC client is not initialized")
|
||||
result_msg = transport.execute(
|
||||
("collective", rpc_id),
|
||||
msg,
|
||||
timeout=timeout,
|
||||
timeout_message=f"collective_rpc timed out after {timeout} seconds",
|
||||
block_on_submit=True,
|
||||
)
|
||||
if not isinstance(result_msg, CollectiveRPCResultMessage):
|
||||
raise RuntimeError(f"unexpected collective RPC result type: {type(result_msg).__name__}")
|
||||
return list(result_msg.results)
|
||||
|
||||
async def collective_rpc_async(
|
||||
self,
|
||||
@@ -1555,22 +1765,61 @@ class AsyncOmniEngine:
|
||||
finalizer.detach()
|
||||
|
||||
logger.info("[AsyncOmniEngine] Shutting down Orchestrator")
|
||||
if self.request_queue is not None:
|
||||
self.request_queue.sync_q.put_nowait(ShutdownRequestMessage())
|
||||
if self.is_alive():
|
||||
self.orchestrator_thread.join()
|
||||
request_queue_closed = False
|
||||
shutdown_enqueued = enqueue_orchestrator_shutdown(
|
||||
self.request_queue,
|
||||
timeout=SHUTDOWN_ENQUEUE_TIMEOUT_S,
|
||||
)
|
||||
if self.request_queue is not None and not shutdown_enqueued:
|
||||
logger.error(
|
||||
"[AsyncOmniEngine] Failed to enqueue orchestrator shutdown; "
|
||||
"closing the request queue to wake the request handler"
|
||||
)
|
||||
try:
|
||||
self.request_queue.close()
|
||||
request_queue_closed = True
|
||||
except Exception:
|
||||
logger.exception("[AsyncOmniEngine] Failed to close the request queue")
|
||||
|
||||
if self._correlated_rpc_client is not None:
|
||||
try:
|
||||
self._correlated_rpc_client.close()
|
||||
except Exception:
|
||||
logger.exception("[AsyncOmniEngine] Failed to close correlated RPC client")
|
||||
|
||||
orchestrator_stopped = False
|
||||
try:
|
||||
if self.is_alive():
|
||||
self.orchestrator_thread.join(timeout=SHUTDOWN_JOIN_TIMEOUT_S)
|
||||
orchestrator_stopped = not self.is_alive()
|
||||
if not orchestrator_stopped:
|
||||
logger.error(
|
||||
"[AsyncOmniEngine] Orchestrator did not stop within %.1f seconds; continuing cleanup",
|
||||
SHUTDOWN_JOIN_TIMEOUT_S,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("[AsyncOmniEngine] Failed to join Orchestrator thread")
|
||||
|
||||
for q in (self.request_queue, self.output_queue, self.rpc_output_queue):
|
||||
try:
|
||||
q.close()
|
||||
if not (q is self.request_queue and request_queue_closed):
|
||||
q.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if hasattr(self, "_runtime") and self._runtime is not None:
|
||||
if hasattr(self, "_runtime") and self._runtime is not None and orchestrator_stopped:
|
||||
try:
|
||||
self._runtime.shutdown()
|
||||
except Exception:
|
||||
logger.exception("[AsyncOmniEngine] Failed to shutdown StageRuntime")
|
||||
elif hasattr(self, "_runtime") and self._runtime is not None:
|
||||
logger.warning("[AsyncOmniEngine] Deferring StageRuntime shutdown until the Orchestrator exits")
|
||||
threading.Thread(
|
||||
target=shutdown_runtime_after_orchestrator,
|
||||
args=(self.orchestrator_thread, self._runtime),
|
||||
daemon=True,
|
||||
name="omni-stage-runtime-shutdown",
|
||||
).start()
|
||||
|
||||
# ── Release CuMem allocator memory pool ──────────────────────────────
|
||||
# When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory
|
||||
|
||||
@@ -157,7 +157,7 @@ class MembershipController:
|
||||
input_addr = StagePool._client_input_addr(client)
|
||||
if input_addr is None:
|
||||
raise RuntimeError(f"remote replica factory for stage {stage_id} produced a client without input address")
|
||||
pool.add_client(input_addr, client)
|
||||
pool.add_client(input_addr, client, replica_id=replica_id)
|
||||
logger.info(
|
||||
"[MembershipController] attached remote replica stage=%d replica=%d addr=%s",
|
||||
stage_id,
|
||||
|
||||
@@ -106,3 +106,7 @@ class CollectiveRPCResultMessage(EngineQueueMessage, kw_only=True):
|
||||
method: str
|
||||
stage_ids: list[int]
|
||||
results: list[object]
|
||||
|
||||
@property
|
||||
def rpc_correlation_key(self) -> tuple[str, str]:
|
||||
return ("collective", self.rpc_id)
|
||||
@@ -13,8 +13,9 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time as _time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import janus
|
||||
import torch
|
||||
@@ -22,11 +23,12 @@ from vllm.config import ModelConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
from vllm.pooling_params import PoolingParams
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from vllm.v1.engine import EngineCoreOutputs
|
||||
from vllm.v1.engine.exceptions import EngineDeadError
|
||||
from vllm.v1.metrics.stats import IterationStats
|
||||
|
||||
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
|
||||
from vllm_omni.engine import OmniEngineCoreRequest
|
||||
from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
|
||||
from vllm_omni.engine.membership_controller import MembershipController
|
||||
@@ -53,6 +55,23 @@ from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import (
|
||||
DuplexControlPlanePort,
|
||||
DuplexOutputContext,
|
||||
DuplexOutputDecision,
|
||||
DuplexRequestIdentity,
|
||||
DuplexRuntimeExtension,
|
||||
DuplexStageRequestContext,
|
||||
DuplexStageSubmission,
|
||||
DuplexStageSubmissionResult,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_session import (
|
||||
DuplexSessionRuntimeManager,
|
||||
DuplexSessionRuntimeState,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
|
||||
|
||||
def _build_terminal_empty_output(
|
||||
request_id: str,
|
||||
@@ -141,8 +160,13 @@ def build_engine_core_request_from_tokens(
|
||||
pooling_params = params.clone()
|
||||
|
||||
prompt_embeds: torch.Tensor | None = prompt.get("prompt_embeds")
|
||||
raw_additional_information = prompt.get("additional_information")
|
||||
model_intermediate_buffer = prompt.get("model_intermediate_buffer")
|
||||
wire_payload: dict[str, Any] | None = None
|
||||
if isinstance(raw_additional_information, dict):
|
||||
wire_payload = dict(raw_additional_information)
|
||||
additional_info_payload = serialize_additional_information(
|
||||
prompt.get("additional_information"),
|
||||
wire_payload,
|
||||
log_prefix=f"build_engine_core_request_from_tokens req={request_id}",
|
||||
)
|
||||
|
||||
@@ -159,6 +183,7 @@ def build_engine_core_request_from_tokens(
|
||||
prompt_embeds=prompt_embeds,
|
||||
resumable=resumable,
|
||||
additional_information=additional_info_payload,
|
||||
model_intermediate_buffer=model_intermediate_buffer if isinstance(model_intermediate_buffer, dict) else None,
|
||||
)
|
||||
|
||||
|
||||
@@ -186,6 +211,10 @@ class OrchestratorRequestState:
|
||||
|
||||
# Per-request pipeline timing accumulator (milliseconds)
|
||||
pipeline_timings: dict[str, float] = field(default_factory=dict)
|
||||
duplex_identity: DuplexRequestIdentity | None = None
|
||||
duplex_stage_fences: dict[int, DuplexFence] = field(default_factory=dict)
|
||||
duplex_config_generation: int = -1
|
||||
running_counter_registered: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -194,10 +223,124 @@ class StreamingInputState:
|
||||
enabled: bool = False
|
||||
# Flag of segment of streaming input finished
|
||||
segment_finished: bool = False
|
||||
# Tokens from the current raw segment boundary. The vLLM output processor
|
||||
# does not guarantee that EngineCoreOutput.new_token_ids survives on the
|
||||
# processed RequestOutput used by the routing layer.
|
||||
segment_token_ids: list[int] = field(default_factory=list)
|
||||
segment_output_metadata: dict[str, Any] = field(default_factory=dict)
|
||||
# Streaming update prompt length
|
||||
new_prompt_len_snapshot: int | None = None
|
||||
# Model/bridge-specific runtime states (e.g., thinker->talker)
|
||||
bridge_states: dict[str, Any] = field(default_factory=dict)
|
||||
# Synchronous stage-transition capability installed by the orchestrator
|
||||
# while the downstream input processor consumes upstream token output.
|
||||
source_token_decoder: Callable[..., str] | None = None
|
||||
|
||||
|
||||
class _OrchestratorDuplexStagePort:
|
||||
"""Adapts generic stage pools to the model-neutral duplex control plane."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
stage_pools: list[StagePool],
|
||||
request_states: dict[str, OrchestratorRequestState],
|
||||
running_counter: OmniRequestCounter | None,
|
||||
cleanup_request_ids: Callable[..., Any],
|
||||
) -> None:
|
||||
self._stage_pools = stage_pools
|
||||
self._request_states = request_states
|
||||
self._running_counter = running_counter
|
||||
self._cleanup_request_ids = cleanup_request_ids
|
||||
|
||||
@property
|
||||
def stage_count(self) -> int:
|
||||
return len(self._stage_pools)
|
||||
|
||||
def sampling_defaults(self) -> tuple[object, ...]:
|
||||
return tuple(pool.stage_client.default_sampling_params for pool in self._stage_pools)
|
||||
|
||||
@staticmethod
|
||||
def _sync_bridge_state(
|
||||
request_state: OrchestratorRequestState,
|
||||
context: DuplexStageRequestContext,
|
||||
) -> None:
|
||||
duplex_state = request_state.streaming.bridge_states.setdefault("duplex", {})
|
||||
if not isinstance(duplex_state, dict):
|
||||
duplex_state = {}
|
||||
request_state.streaming.bridge_states["duplex"] = duplex_state
|
||||
previous_epoch = duplex_state.get("epoch")
|
||||
if not isinstance(duplex_state.get("model_turn_id"), int) or previous_epoch != context.fence.epoch:
|
||||
duplex_state["model_turn_id"] = context.fence.turn_id
|
||||
duplex_state.update(
|
||||
{
|
||||
"session_id": context.session_id,
|
||||
"fence": context.fence,
|
||||
"incarnation": context.fence.incarnation,
|
||||
"epoch": context.fence.epoch,
|
||||
"turn_id": context.fence.turn_id,
|
||||
"response_seq": context.fence.response_seq,
|
||||
"session_config": dict(context.session_config),
|
||||
"runtime_config": dict(context.runtime_config),
|
||||
}
|
||||
)
|
||||
|
||||
def ensure_request(self, context: DuplexStageRequestContext) -> None:
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity
|
||||
|
||||
request_state = self._request_states.get(context.request_id)
|
||||
if request_state is None:
|
||||
request_state = OrchestratorRequestState(
|
||||
request_id=context.request_id,
|
||||
prompt=None,
|
||||
sampling_params_list=list(context.sampling_params),
|
||||
final_stage_id=context.final_stage_id,
|
||||
duplex_config_generation=context.config_generation,
|
||||
)
|
||||
request_state.streaming.enabled = True
|
||||
self._request_states[context.request_id] = request_state
|
||||
elif request_state.duplex_config_generation != context.config_generation:
|
||||
request_state.sampling_params_list = list(context.sampling_params)
|
||||
request_state.duplex_config_generation = context.config_generation
|
||||
request_state.duplex_identity = DuplexRequestIdentity(
|
||||
session_id=context.session_id,
|
||||
fence=context.fence,
|
||||
)
|
||||
self._sync_bridge_state(request_state, context)
|
||||
|
||||
async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult:
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import DuplexStageSubmissionResult
|
||||
|
||||
context = submission.context
|
||||
request_state = self._request_states.get(context.request_id)
|
||||
if request_state is None:
|
||||
raise RuntimeError(f"duplex request was not preregistered: {context.request_id}")
|
||||
request = build_engine_core_request_from_tokens(
|
||||
request_id=context.request_id,
|
||||
prompt=dict(submission.prompt),
|
||||
params=context.stage_sampling_params,
|
||||
model_config=self._stage_pools[context.stage_id].stage_vllm_config.model_config,
|
||||
resumable=True,
|
||||
)
|
||||
request.external_req_id = request.request_id
|
||||
pool = self._stage_pools[context.stage_id]
|
||||
if submission.already_submitted:
|
||||
replica_id = await pool.submit_update(context.request_id, request_state, request)
|
||||
else:
|
||||
replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None)
|
||||
request_state.duplex_stage_fences[context.stage_id] = context.fence
|
||||
request_state.stage_submit_ts[context.stage_id] = _time.time()
|
||||
if not request_state.running_counter_registered and self._running_counter is not None:
|
||||
self._running_counter.increment()
|
||||
request_state.running_counter_registered = True
|
||||
return DuplexStageSubmissionResult(
|
||||
request_id=context.request_id,
|
||||
stage_id=context.stage_id,
|
||||
replica_id=replica_id,
|
||||
)
|
||||
|
||||
async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None:
|
||||
await self._cleanup_request_ids(request_ids, abort=abort)
|
||||
|
||||
|
||||
class Orchestrator:
|
||||
@@ -208,6 +351,7 @@ class Orchestrator:
|
||||
_running_counter: OmniRequestCounter | None = None
|
||||
_transfer_emitter: Any = None
|
||||
_stat_logger: OmniPrometheusStatLogger | None = None
|
||||
duplex_control_plane: DuplexControlPlanePort | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -223,6 +367,9 @@ class Orchestrator:
|
||||
transfer_emitter: Any = None,
|
||||
log_stats: bool = False,
|
||||
enable_orch_monitor: bool = False,
|
||||
duplex_runtime_extension: DuplexRuntimeExtension | None = None,
|
||||
enable_duplex_control: bool = False,
|
||||
duplex_session_config: DuplexSessionRuntimeConfig | None = None,
|
||||
) -> None:
|
||||
self.request_async_queue = request_async_queue
|
||||
self.output_async_queue = output_async_queue
|
||||
@@ -236,7 +383,7 @@ class Orchestrator:
|
||||
replica_sampler=self._sample_replica_metrics,
|
||||
)
|
||||
for stage_id, pool in enumerate(self.stage_pools):
|
||||
for replica_id in pool.live_replica_ids():
|
||||
for replica_id in pool.available_replica_ids():
|
||||
self._orch_monitor.register_replica(stage_id, replica_id)
|
||||
|
||||
# PD disaggregation state
|
||||
@@ -254,6 +401,33 @@ class Orchestrator:
|
||||
self._cfg_tracker = CfgCompanionTracker()
|
||||
self._stage_input_processors: dict[int, Any] = {}
|
||||
|
||||
self.duplex_control_plane: DuplexControlPlanePort | None = None
|
||||
self._duplex_reaper_interval_s = 1.0
|
||||
if enable_duplex_control:
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_control_plane import DuplexControlPlane
|
||||
from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseConfig
|
||||
|
||||
runtime_session_config = duplex_session_config or DuplexSessionRuntimeConfig()
|
||||
self._duplex_reaper_interval_s = runtime_session_config.reaper_interval_s
|
||||
|
||||
self.duplex_control_plane = DuplexControlPlane(
|
||||
extension=duplex_runtime_extension,
|
||||
stage_port=_OrchestratorDuplexStagePort(
|
||||
stage_pools=self.stage_pools,
|
||||
request_states=self.request_states,
|
||||
running_counter=self._running_counter,
|
||||
cleanup_request_ids=self._cleanup_request_ids,
|
||||
),
|
||||
result_sink=self.rpc_async_queue,
|
||||
lifecycle_sink=self.output_async_queue,
|
||||
lease_config=DuplexLeaseConfig(
|
||||
idle_ttl_s=runtime_session_config.idle_ttl_s,
|
||||
disconnect_grace_s=runtime_session_config.disconnect_grace_s,
|
||||
),
|
||||
max_sessions=runtime_session_config.max_sessions,
|
||||
completed_append_limit=runtime_session_config.completed_append_cache_size,
|
||||
)
|
||||
|
||||
self._shutdown_event = asyncio.Event()
|
||||
self._stages_shutdown = False
|
||||
self._fatal_error: str | None = None
|
||||
@@ -327,6 +501,17 @@ class Orchestrator:
|
||||
logger.exception("[Orchestrator] OmniPrometheusStatLogger init failed; metrics wrap disabled")
|
||||
self._stat_logger = None
|
||||
|
||||
@property
|
||||
def duplex_sessions(self) -> DuplexSessionRuntimeManager:
|
||||
if self.duplex_control_plane is None:
|
||||
raise RuntimeError("duplex control plane is disabled")
|
||||
return self.duplex_control_plane.sessions
|
||||
|
||||
def _require_duplex_control_plane(self) -> DuplexControlPlanePort:
|
||||
if self.duplex_control_plane is None:
|
||||
raise RuntimeError("duplex control plane is disabled")
|
||||
return self.duplex_control_plane
|
||||
|
||||
async def run(self) -> None:
|
||||
"""Main entry point for the Orchestrator event loop."""
|
||||
logger.info("[Orchestrator] Starting event loop")
|
||||
@@ -347,6 +532,8 @@ class Orchestrator:
|
||||
membership_watcher = self._membership.start()
|
||||
|
||||
tasks = [request_task, output_task]
|
||||
if self.duplex_control_plane is not None:
|
||||
tasks.append(asyncio.create_task(self._duplex_reaper_loop(), name="orchestrator-duplex-reaper"))
|
||||
if membership_watcher is not None:
|
||||
tasks.append(membership_watcher)
|
||||
|
||||
@@ -361,6 +548,15 @@ class Orchestrator:
|
||||
# During teardown this is expected; the finally block handles
|
||||
# proper cleanup. Do not re-raise.
|
||||
logger.info("[Orchestrator] Engine dead during shutdown: %s", e)
|
||||
if self._fatal_error is None:
|
||||
self._fatal_error = str(e) or "Stage engine died"
|
||||
await self.rpc_async_queue.put(
|
||||
ErrorMessage(
|
||||
error=self._fatal_error or str(e),
|
||||
fatal=True,
|
||||
stage_id=self._fatal_error_stage_id,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("[Orchestrator] Fatal error in orchestrator tasks")
|
||||
raise
|
||||
@@ -373,6 +569,8 @@ class Orchestrator:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
except Exception:
|
||||
pass
|
||||
if self.duplex_control_plane is not None:
|
||||
await self.duplex_control_plane.shutdown()
|
||||
|
||||
if self._fatal_error is not None:
|
||||
await self._drain_pending_requests_on_fatal()
|
||||
@@ -405,6 +603,8 @@ class Orchestrator:
|
||||
await self._handle_streaming_update(msg)
|
||||
elif msg_type == "add_companion_request":
|
||||
await self._handle_add_companion(msg)
|
||||
elif self.duplex_control_plane is not None and self.duplex_control_plane.accepts(msg):
|
||||
self.duplex_control_plane.dispatch(msg)
|
||||
elif msg_type == "abort":
|
||||
await self._handle_abort(msg)
|
||||
elif msg_type == "collective_rpc":
|
||||
@@ -432,6 +632,21 @@ class Orchestrator:
|
||||
else:
|
||||
logger.warning("[Orchestrator] Unknown message type: %s", msg_type)
|
||||
|
||||
async def _duplex_reaper_loop(self) -> None:
|
||||
while not self._shutdown_event.is_set():
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._shutdown_event.wait(),
|
||||
timeout=self._duplex_reaper_interval_s,
|
||||
)
|
||||
except TimeoutError:
|
||||
plane = self.duplex_control_plane
|
||||
if plane is not None:
|
||||
try:
|
||||
await plane.reap_expired()
|
||||
except Exception:
|
||||
logger.exception("[Orchestrator] Duplex expiry cleanup failed; retrying on next tick")
|
||||
|
||||
async def _handle_add_request(self, msg: StageSubmissionMessage) -> None:
|
||||
"""Handle an add_request message from the main thread."""
|
||||
stage_id = 0
|
||||
@@ -466,8 +681,7 @@ class Orchestrator:
|
||||
mm_features=getattr(prompt, "mm_features", None),
|
||||
)
|
||||
self.request_states[request_id] = req_state
|
||||
if self._running_counter is not None:
|
||||
self._running_counter.increment()
|
||||
self._register_running_request(req_state)
|
||||
req_state.streaming.enabled = bool(getattr(prompt, "resumable", False))
|
||||
req_state.stage_submit_ts[stage_id] = _time.time()
|
||||
enqueue_ts = msg.enqueue_ts
|
||||
@@ -590,6 +804,7 @@ class Orchestrator:
|
||||
return
|
||||
for pool in self.stage_pools:
|
||||
await pool.abort_requests(request_ids)
|
||||
pool.release_bindings(request_ids)
|
||||
|
||||
def _release_request_bindings(self, request_ids: list[str]) -> None:
|
||||
"""Release all stage-local route bindings for the given request ids."""
|
||||
@@ -662,7 +877,7 @@ class Orchestrator:
|
||||
idle = True
|
||||
for stage_id in range(self.num_stages):
|
||||
pool = self.stage_pools[stage_id]
|
||||
for replica_id in pool.live_replica_ids():
|
||||
for replica_id in pool.available_replica_ids():
|
||||
if self._shutdown_event.is_set():
|
||||
return
|
||||
|
||||
@@ -686,6 +901,17 @@ class Orchestrator:
|
||||
if req_state is None or not req_state.streaming.enabled:
|
||||
continue
|
||||
req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False))
|
||||
req_state.streaming.segment_token_ids = (
|
||||
self._coerce_int_list(getattr(eco, "new_token_ids", None))
|
||||
if req_state.streaming.segment_finished
|
||||
else []
|
||||
)
|
||||
raw_mm = self._completion_multimodal_output(eco, None)
|
||||
req_state.streaming.segment_output_metadata = (
|
||||
dict(raw_mm)
|
||||
if req_state.streaming.segment_finished and isinstance(raw_mm, dict)
|
||||
else {}
|
||||
)
|
||||
req_state.streaming.new_prompt_len_snapshot = getattr(
|
||||
eco,
|
||||
"new_prompt_len_snapshot",
|
||||
@@ -715,32 +941,62 @@ class Orchestrator:
|
||||
raise
|
||||
except EngineDeadError as e:
|
||||
logger.error(
|
||||
"[Orchestrator] Stage-%s is dead: %s",
|
||||
"[Orchestrator] Stage-%s replica-%s is dead: %s",
|
||||
stage_id,
|
||||
replica_id,
|
||||
e,
|
||||
)
|
||||
# TODO: Fault handling is intentionally fail-stop at
|
||||
# the orchestrator level today. If one replica in a
|
||||
# logical stage dies, we promote it to `_fatal_error`,
|
||||
# notify requests already admitted to that stage, and
|
||||
# re-raise so `run()` shuts down all stages. This is
|
||||
# conservative but means a single unhealthy replica in
|
||||
# a multi-replica deployment can take down otherwise
|
||||
# healthy replicas in other stages. Revisit this when
|
||||
# adding per-replica fault isolation / eviction.
|
||||
self._fatal_error = str(e)
|
||||
self._fatal_error_stage_id = stage_id
|
||||
for req_id, req_state in list(self.request_states.items()):
|
||||
if stage_id in req_state.stage_submit_ts:
|
||||
affected_request_ids = pool.mark_replica_unavailable(replica_id)
|
||||
closed_sessions = (
|
||||
self.duplex_control_plane.close_sessions_for_request_ids(
|
||||
affected_request_ids,
|
||||
abort=False,
|
||||
)
|
||||
if self.duplex_control_plane is not None
|
||||
else {}
|
||||
)
|
||||
for session_id, stale_request_ids in closed_sessions.items():
|
||||
affected_request_ids.extend(stale_request_ids)
|
||||
logger.warning(
|
||||
"[Orchestrator] closed duplex session %s after stage-%s replica-%s died; "
|
||||
"stale_request_ids=%s",
|
||||
session_id,
|
||||
stage_id,
|
||||
replica_id,
|
||||
stale_request_ids,
|
||||
)
|
||||
affected_request_ids = list(dict.fromkeys(affected_request_ids))
|
||||
if pool.available_replica_ids():
|
||||
for req_id in affected_request_ids:
|
||||
await self.output_async_queue.put(
|
||||
ErrorMessage(
|
||||
error=str(e),
|
||||
fatal=True,
|
||||
fatal=False,
|
||||
request_id=req_id,
|
||||
stage_id=stage_id,
|
||||
)
|
||||
)
|
||||
self.request_states.pop(req_id, None)
|
||||
await self._cleanup_request_ids(
|
||||
affected_request_ids,
|
||||
close_duplex_sessions=True,
|
||||
)
|
||||
continue
|
||||
|
||||
self._fatal_error = str(e)
|
||||
self._fatal_error_stage_id = stage_id
|
||||
for req_id in affected_request_ids:
|
||||
await self.output_async_queue.put(
|
||||
ErrorMessage(
|
||||
error=str(e),
|
||||
fatal=True,
|
||||
request_id=req_id,
|
||||
stage_id=stage_id,
|
||||
)
|
||||
)
|
||||
await self._cleanup_request_ids(
|
||||
affected_request_ids,
|
||||
close_duplex_sessions=True,
|
||||
)
|
||||
self._shutdown_event.set()
|
||||
raise
|
||||
except Exception:
|
||||
@@ -781,7 +1037,8 @@ class Orchestrator:
|
||||
continue
|
||||
|
||||
stage_metrics = None
|
||||
if output.finished:
|
||||
segment_finished = req_state.streaming.enabled and req_state.streaming.segment_finished
|
||||
if output.finished or segment_finished:
|
||||
stage_metrics = pool.build_stage_metrics(
|
||||
[output],
|
||||
submit_ts=req_state.stage_submit_ts.get(stage_id, _time.time()),
|
||||
@@ -811,22 +1068,56 @@ class Orchestrator:
|
||||
await self._cleanup_request_ids(
|
||||
[parent_id, *self._cfg_tracker.cleanup_parent(parent_id)],
|
||||
abort=True,
|
||||
close_duplex_sessions=True,
|
||||
)
|
||||
|
||||
# ---- Shared helpers ----
|
||||
|
||||
async def _cleanup_request_ids(self, request_ids: list[str], *, abort: bool = False) -> None:
|
||||
async def _cleanup_request_ids(
|
||||
self,
|
||||
request_ids: list[str],
|
||||
*,
|
||||
abort: bool = False,
|
||||
close_duplex_sessions: bool = False,
|
||||
) -> None:
|
||||
"""Release pool bindings and logical request state for the given ids."""
|
||||
if not request_ids:
|
||||
return
|
||||
|
||||
if abort:
|
||||
await self._abort_request_ids(request_ids)
|
||||
self._release_request_bindings(request_ids)
|
||||
for request_id in request_ids:
|
||||
self._pd_kv_params.pop(request_id, None)
|
||||
if self.request_states.pop(request_id, None) is not None and self._running_counter is not None:
|
||||
self._running_counter.decrement()
|
||||
cleanup_ids = list(dict.fromkeys(request_ids))
|
||||
closing_session_ids: list[str] = []
|
||||
if close_duplex_sessions and self.duplex_control_plane is not None:
|
||||
closed_sessions = self.duplex_control_plane.close_sessions_for_request_ids(
|
||||
cleanup_ids,
|
||||
abort=abort,
|
||||
cleanup_in_progress=True,
|
||||
)
|
||||
closing_session_ids.extend(closed_sessions)
|
||||
for session_id, stale_request_ids in closed_sessions.items():
|
||||
logger.info(
|
||||
"[Orchestrator] closed duplex session %s while cleaning failed request ids %s",
|
||||
session_id,
|
||||
stale_request_ids,
|
||||
)
|
||||
cleanup_ids.extend(stale_request_ids)
|
||||
cleanup_ids = list(dict.fromkeys(cleanup_ids))
|
||||
|
||||
try:
|
||||
if abort:
|
||||
await self._abort_request_ids(cleanup_ids)
|
||||
self._release_request_bindings(cleanup_ids)
|
||||
for request_id in cleanup_ids:
|
||||
self._pd_kv_params.pop(request_id, None)
|
||||
req_state = self.request_states.pop(request_id, None)
|
||||
if req_state is not None and req_state.running_counter_registered and self._running_counter is not None:
|
||||
self._running_counter.decrement()
|
||||
req_state.running_counter_registered = False
|
||||
except BaseException:
|
||||
if closing_session_ids and self.duplex_control_plane is not None:
|
||||
self.duplex_control_plane.defer_request_cleanups(closing_session_ids)
|
||||
raise
|
||||
if closing_session_ids and self.duplex_control_plane is not None:
|
||||
self.duplex_control_plane.finalize_closed_sessions(closing_session_ids)
|
||||
|
||||
async def _apply_raw_terminal_stage_finish(
|
||||
self,
|
||||
@@ -873,6 +1164,34 @@ class Orchestrator:
|
||||
params.cfg_kv_request_ids = companion_request_ids
|
||||
return params
|
||||
|
||||
def _duplex_session_for_req_state(self, req_state: OrchestratorRequestState) -> DuplexSessionRuntimeState | None:
|
||||
if self.duplex_control_plane is None:
|
||||
return None
|
||||
return self.duplex_control_plane.session_for_identity(req_state.duplex_identity)
|
||||
|
||||
def _record_duplex_stage_submission(
|
||||
self,
|
||||
stage_id: int,
|
||||
request_id: str,
|
||||
replica_id: int,
|
||||
req_state: OrchestratorRequestState,
|
||||
) -> None:
|
||||
del replica_id
|
||||
identity = req_state.duplex_identity
|
||||
session = self._duplex_session_for_req_state(req_state)
|
||||
if identity is None or session is None:
|
||||
return
|
||||
req_state.duplex_stage_fences[stage_id] = identity.fence
|
||||
session.bind_stage_request(stage_id, request_id, fence=identity.fence)
|
||||
req_state.stage_submit_ts[stage_id] = _time.time()
|
||||
self._register_running_request(req_state)
|
||||
|
||||
def _register_running_request(self, req_state: OrchestratorRequestState) -> None:
|
||||
if req_state.running_counter_registered or self._running_counter is None:
|
||||
return
|
||||
self._running_counter.increment()
|
||||
req_state.running_counter_registered = True
|
||||
|
||||
async def _route_output(
|
||||
self,
|
||||
stage_id: int,
|
||||
@@ -885,7 +1204,6 @@ class Orchestrator:
|
||||
req_id = output.request_id
|
||||
finished = output.finished
|
||||
submit_ts = req_state.stage_submit_ts.get(stage_id)
|
||||
|
||||
# CFG companion: stash output so parent can bundle [parent, *companions]
|
||||
# into source_outputs for the bridge (e.g. thinker2imagegen).
|
||||
if finished and self._cfg_tracker.is_companion(req_id):
|
||||
@@ -899,7 +1217,17 @@ class Orchestrator:
|
||||
req_state.finished_final_output_stage_ids.add(stage_id)
|
||||
final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
|
||||
request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids)
|
||||
if self.stage_pools[stage_id].final_output:
|
||||
# Duplex stage-0 segment boundaries are not client-visible outputs:
|
||||
# direct decisions are emitted by the model runtime extension below,
|
||||
# while spoken content flows through the next stage. Forwarding
|
||||
# the raw stage-0 output as well injects one cumulative-text,
|
||||
# no-audio message per unit that every downstream consumer must
|
||||
# filter out again (the official implementation returns exactly one
|
||||
# result per audio chunk).
|
||||
is_duplex_stage0_segment = (
|
||||
stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished
|
||||
)
|
||||
if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment:
|
||||
await self.output_async_queue.put(
|
||||
OutputMessage(
|
||||
request_id=req_id,
|
||||
@@ -907,7 +1235,10 @@ class Orchestrator:
|
||||
replica_id=replica_id,
|
||||
engine_outputs=output,
|
||||
metrics=stage_metrics,
|
||||
finished=request_finished,
|
||||
finished=(
|
||||
request_finished
|
||||
or (self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished)
|
||||
),
|
||||
stage_submit_ts=submit_ts,
|
||||
)
|
||||
)
|
||||
@@ -928,6 +1259,18 @@ class Orchestrator:
|
||||
self._pd_kv_params[req_id] = kv_params if isinstance(kv_params, dict) else dict(kv_params)
|
||||
req_state.pd_prefill_multimodal_output = getattr(output, "multimodal_output", None)
|
||||
|
||||
duplex_output_decision = self._duplex_output_decision(stage_id, output, req_state)
|
||||
if duplex_output_decision is not None:
|
||||
await self._emit_duplex_direct_output(
|
||||
stage_id,
|
||||
req_id,
|
||||
output,
|
||||
duplex_output_decision,
|
||||
stage_metrics,
|
||||
submit_ts,
|
||||
)
|
||||
return
|
||||
|
||||
if (
|
||||
(finished or (req_state.streaming.enabled and req_state.streaming.segment_finished))
|
||||
and stage_id < req_state.final_stage_id
|
||||
@@ -941,6 +1284,12 @@ class Orchestrator:
|
||||
):
|
||||
self._cfg_tracker.defer_parent(req_id, output, stage_id)
|
||||
else:
|
||||
stage_params = req_state.sampling_params_list[stage_id]
|
||||
final_only_finished = (
|
||||
req_state.streaming.enabled
|
||||
and finished
|
||||
and getattr(stage_params, "output_kind", None) == RequestOutputKind.FINAL_ONLY
|
||||
)
|
||||
await self._forward_to_next_stage(
|
||||
req_id,
|
||||
stage_id,
|
||||
@@ -948,9 +1297,14 @@ class Orchestrator:
|
||||
req_state,
|
||||
src_replica_id=replica_id,
|
||||
is_streaming_session=req_state.streaming.enabled,
|
||||
is_final_update=False,
|
||||
is_final_update=final_only_finished,
|
||||
)
|
||||
if req_state.streaming.enabled and finished:
|
||||
if (
|
||||
req_state.streaming.enabled
|
||||
and finished
|
||||
and not final_only_finished
|
||||
and not self._is_duplex_session_request(req_state)
|
||||
):
|
||||
# For streaming sessions, send the terminal (resumable=False) update only on a finish
|
||||
await self._forward_to_next_stage(
|
||||
req_id,
|
||||
@@ -962,7 +1316,7 @@ class Orchestrator:
|
||||
is_final_update=True,
|
||||
)
|
||||
|
||||
if request_finished:
|
||||
if request_finished and not self._is_duplex_session_request(req_state):
|
||||
await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)])
|
||||
|
||||
def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool:
|
||||
@@ -1039,6 +1393,137 @@ class Orchestrator:
|
||||
request.external_req_id = req_id
|
||||
return request
|
||||
|
||||
@staticmethod
|
||||
def _duplex_output_context(
|
||||
req_state: OrchestratorRequestState,
|
||||
*,
|
||||
stage_id: int | None = None,
|
||||
) -> DuplexOutputContext | None:
|
||||
identity = req_state.duplex_identity
|
||||
if identity is None:
|
||||
return None
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import (
|
||||
DuplexOutputContext,
|
||||
DuplexRequestIdentity,
|
||||
)
|
||||
|
||||
fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence
|
||||
return DuplexOutputContext(
|
||||
identity=DuplexRequestIdentity(
|
||||
session_id=identity.session_id,
|
||||
fence=fence,
|
||||
),
|
||||
final_stage_id=req_state.final_stage_id,
|
||||
segment_finished=req_state.streaming.enabled and req_state.streaming.segment_finished,
|
||||
segment_token_ids=tuple(req_state.streaming.segment_token_ids),
|
||||
segment_output_metadata=req_state.streaming.segment_output_metadata,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_duplex_session_request(req_state: OrchestratorRequestState) -> bool:
|
||||
return req_state.duplex_identity is not None
|
||||
|
||||
@staticmethod
|
||||
def _duplex_fence_for_req_state(
|
||||
req_state: OrchestratorRequestState,
|
||||
*,
|
||||
stage_id: int | None = None,
|
||||
) -> DuplexFence | None:
|
||||
context = Orchestrator._duplex_output_context(req_state, stage_id=stage_id)
|
||||
return context.identity.fence if context is not None else None
|
||||
|
||||
def _duplex_output_decision(
|
||||
self,
|
||||
stage_id: int,
|
||||
output: Any,
|
||||
req_state: OrchestratorRequestState,
|
||||
) -> DuplexOutputDecision | None:
|
||||
if self.duplex_control_plane is None:
|
||||
return None
|
||||
context = self._duplex_output_context(req_state, stage_id=stage_id)
|
||||
decision = self.duplex_control_plane.decide_output(
|
||||
stage_id,
|
||||
output,
|
||||
context,
|
||||
)
|
||||
return decision
|
||||
|
||||
async def _emit_duplex_direct_output(
|
||||
self,
|
||||
stage_id: int,
|
||||
req_id: str,
|
||||
output: Any,
|
||||
decision: DuplexOutputDecision,
|
||||
stage_metrics: Any,
|
||||
submit_ts: float | None,
|
||||
) -> None:
|
||||
action = getattr(decision.action, "value", decision.action)
|
||||
if action != "direct_response":
|
||||
raise ValueError(f"Unsupported duplex output action: {action}")
|
||||
from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision
|
||||
|
||||
engine_output = attach_duplex_output_decision(
|
||||
OmniRequestOutput(
|
||||
request_id=req_id,
|
||||
finished=True,
|
||||
stage_id=stage_id,
|
||||
final_output_type=decision.final_output_type,
|
||||
request_output=output,
|
||||
),
|
||||
decision,
|
||||
)
|
||||
await self.output_async_queue.put(
|
||||
OutputMessage(
|
||||
request_id=req_id,
|
||||
stage_id=stage_id,
|
||||
engine_outputs=engine_output,
|
||||
metrics=stage_metrics,
|
||||
finished=True,
|
||||
stage_submit_ts=submit_ts,
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _completion_multimodal_output(output: Any, completion: Any) -> dict[str, Any]:
|
||||
mm_output = getattr(output, "multimodal_output", None)
|
||||
if isinstance(mm_output, dict):
|
||||
return mm_output
|
||||
mm_output = getattr(completion, "multimodal_output", None) if completion is not None else None
|
||||
return mm_output if isinstance(mm_output, dict) else {}
|
||||
|
||||
@classmethod
|
||||
def _coerce_int_list(cls, value: Any) -> list[int]:
|
||||
if value is None:
|
||||
return []
|
||||
if hasattr(value, "detach"):
|
||||
try:
|
||||
value = value.detach().cpu().reshape(-1).tolist()
|
||||
except Exception:
|
||||
return []
|
||||
if not isinstance(value, (list, tuple)):
|
||||
return []
|
||||
out: list[int] = []
|
||||
for item in value:
|
||||
token_id = cls._coerce_int(item)
|
||||
if token_id is not None:
|
||||
out.append(token_id)
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def _coerce_int(value: Any) -> int | None:
|
||||
if hasattr(value, "detach"):
|
||||
try:
|
||||
value = value.detach().cpu().reshape(-1)
|
||||
if value.numel() == 0:
|
||||
return None
|
||||
value = value[0].item()
|
||||
except Exception:
|
||||
return None
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
async def _handle_cfg_companion_ready(self, req_id: str) -> None:
|
||||
"""Mark a CFG companion as done; if all companions are done, flush deferred parent."""
|
||||
parent_id = self._cfg_tracker.on_companion_completed(req_id)
|
||||
@@ -1293,9 +1778,9 @@ class Orchestrator:
|
||||
diffusion_prompt = req_state.prompt
|
||||
|
||||
if already_submitted:
|
||||
await next_pool.submit_update(req_id, req_state, diffusion_prompt)
|
||||
replica_id = await next_pool.submit_update(req_id, req_state, diffusion_prompt)
|
||||
else:
|
||||
await next_pool.submit_initial(
|
||||
replica_id = await next_pool.submit_initial(
|
||||
req_id,
|
||||
req_state,
|
||||
diffusion_prompt,
|
||||
@@ -1307,6 +1792,12 @@ class Orchestrator:
|
||||
},
|
||||
params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
|
||||
)
|
||||
self._record_duplex_stage_submission(
|
||||
next_logical,
|
||||
req_id,
|
||||
replica_id,
|
||||
req_state,
|
||||
)
|
||||
req_state.stage_submit_ts[next_logical] = _time.time()
|
||||
_tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0
|
||||
self._emit_tx_edge(
|
||||
@@ -1351,9 +1842,15 @@ class Orchestrator:
|
||||
)
|
||||
request.external_req_id = request.request_id
|
||||
if already_submitted:
|
||||
await next_pool.submit_update(req_id, req_state, request)
|
||||
replica_id = await next_pool.submit_update(req_id, req_state, request)
|
||||
else:
|
||||
await next_pool.submit_initial(req_id, req_state, request, prompt_text=None)
|
||||
replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None)
|
||||
self._record_duplex_stage_submission(
|
||||
next_logical,
|
||||
req_id,
|
||||
replica_id,
|
||||
req_state,
|
||||
)
|
||||
|
||||
req_state.stage_submit_ts[next_logical] = _time.time()
|
||||
_tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0
|
||||
@@ -1373,6 +1870,13 @@ class Orchestrator:
|
||||
{},
|
||||
)[req_id] = req_state.pd_prefill_multimodal_output
|
||||
|
||||
previous_decoder = req_state.streaming.source_token_decoder
|
||||
source_processor = self.stage_pools[src_stage_id].output_processor
|
||||
tokenizer = getattr(source_processor, "tokenizer", None)
|
||||
decode = getattr(tokenizer, "decode", None)
|
||||
if callable(decode):
|
||||
req_state.streaming.source_token_decoder = decode
|
||||
|
||||
try:
|
||||
next_inputs = next_client.process_engine_inputs(
|
||||
source_outputs,
|
||||
@@ -1386,6 +1890,8 @@ class Orchestrator:
|
||||
next_logical,
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
req_state.streaming.source_token_decoder = previous_decoder
|
||||
|
||||
if not next_inputs:
|
||||
if not getattr(output, "finished", False):
|
||||
@@ -1446,9 +1952,15 @@ class Orchestrator:
|
||||
)
|
||||
|
||||
if already_submitted:
|
||||
await next_pool.submit_update(req_id, req_state, request)
|
||||
replica_id = await next_pool.submit_update(req_id, req_state, request)
|
||||
else:
|
||||
await next_pool.submit_initial(req_id, req_state, request, prompt_text=None)
|
||||
replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None)
|
||||
self._record_duplex_stage_submission(
|
||||
next_logical,
|
||||
req_id,
|
||||
replica_id,
|
||||
req_state,
|
||||
)
|
||||
|
||||
req_state.stage_submit_ts[next_logical] = _time.time()
|
||||
_tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import queue
|
||||
import threading
|
||||
from typing import TypeAlias
|
||||
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.engine.messages import (
|
||||
EngineQueueMessage,
|
||||
ErrorMessage,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
RpcCorrelationKey: TypeAlias = tuple[str, str]
|
||||
RpcWaiter: TypeAlias = queue.Queue[EngineQueueMessage]
|
||||
|
||||
|
||||
class RpcResultRouter:
|
||||
"""Single consumer that dispatches shared RPC results by correlation ID."""
|
||||
|
||||
def __init__(self, source_queue) -> None:
|
||||
self._source_queue = source_queue
|
||||
self._pending: dict[RpcCorrelationKey, RpcWaiter] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._stopped = threading.Event()
|
||||
self._terminal_error: ErrorMessage | None = None
|
||||
self._thread = threading.Thread(
|
||||
target=self._run,
|
||||
daemon=True,
|
||||
name="omni-rpc-result-router",
|
||||
)
|
||||
self._thread.start()
|
||||
|
||||
def register(self, key: RpcCorrelationKey) -> RpcWaiter:
|
||||
waiter: RpcWaiter = queue.Queue(maxsize=1)
|
||||
with self._lock:
|
||||
if self._stopped.is_set():
|
||||
raise RuntimeError("RPC result router is closed")
|
||||
if self._terminal_error is not None:
|
||||
raise RuntimeError(self._terminal_error.error)
|
||||
if key in self._pending:
|
||||
raise RuntimeError(f"duplicate pending RPC correlation key: {key!r}")
|
||||
self._pending[key] = waiter
|
||||
return waiter
|
||||
|
||||
def unregister(self, key: RpcCorrelationKey, waiter: RpcWaiter) -> None:
|
||||
with self._lock:
|
||||
if self._pending.get(key) is waiter:
|
||||
self._pending.pop(key, None)
|
||||
|
||||
def close(self) -> None:
|
||||
if self._stopped.is_set():
|
||||
return
|
||||
self._stopped.set()
|
||||
self._broadcast_error(ErrorMessage(error="RPC result router closed", fatal=True))
|
||||
self._thread.join(timeout=1.0)
|
||||
|
||||
def _run(self) -> None:
|
||||
while not self._stopped.is_set():
|
||||
try:
|
||||
message = self._source_queue.get(timeout=0.1)
|
||||
except queue.Empty:
|
||||
continue
|
||||
except Exception as exc:
|
||||
if not self._stopped.is_set():
|
||||
logger.exception("RPC result router source queue failed")
|
||||
self._broadcast_error(ErrorMessage(error=str(exc), fatal=True))
|
||||
return
|
||||
|
||||
if isinstance(message, ErrorMessage):
|
||||
if message.fatal:
|
||||
self._broadcast_error(message)
|
||||
else:
|
||||
logger.warning(
|
||||
"Dropping uncorrelated non-fatal RPC error request_id=%s stage_id=%s: %s",
|
||||
message.request_id,
|
||||
message.stage_id,
|
||||
message.error,
|
||||
)
|
||||
continue
|
||||
key = self._correlation_key(message)
|
||||
if key is None:
|
||||
logger.warning(
|
||||
"Dropping unexpected RPC result message type=%s",
|
||||
getattr(message, "type", type(message).__name__),
|
||||
)
|
||||
continue
|
||||
with self._lock:
|
||||
waiter = self._pending.pop(key, None)
|
||||
if waiter is None:
|
||||
logger.warning("Dropping late or unknown RPC result correlation_key=%s", key)
|
||||
continue
|
||||
waiter.put_nowait(message)
|
||||
|
||||
def _broadcast_error(self, message: ErrorMessage) -> None:
|
||||
with self._lock:
|
||||
if message.fatal:
|
||||
self._terminal_error = message
|
||||
waiters = list(self._pending.values())
|
||||
self._pending.clear()
|
||||
for waiter in waiters:
|
||||
try:
|
||||
waiter.put_nowait(message)
|
||||
except queue.Full:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _correlation_key(message: EngineQueueMessage) -> RpcCorrelationKey | None:
|
||||
key = getattr(message, "rpc_correlation_key", None)
|
||||
if not isinstance(key, tuple) or len(key) != 2 or not all(isinstance(part, str) and part for part in key):
|
||||
return None
|
||||
return key
|
||||
|
||||
|
||||
class CorrelatedRpcClient:
|
||||
"""Own request submission and correlated result waiting as one lifecycle."""
|
||||
|
||||
def __init__(self, request_queue, result_queue) -> None:
|
||||
self._request_queue = request_queue
|
||||
self._router = RpcResultRouter(result_queue)
|
||||
|
||||
def execute(
|
||||
self,
|
||||
key: RpcCorrelationKey,
|
||||
message: EngineQueueMessage,
|
||||
*,
|
||||
timeout: float | None,
|
||||
timeout_message: str,
|
||||
block_on_submit: bool = False,
|
||||
) -> EngineQueueMessage:
|
||||
waiter = self._router.register(key)
|
||||
try:
|
||||
if block_on_submit:
|
||||
self._request_queue.put(message)
|
||||
else:
|
||||
self._request_queue.put_nowait(message)
|
||||
try:
|
||||
result = waiter.get(timeout=timeout)
|
||||
except queue.Empty as exc:
|
||||
raise TimeoutError(timeout_message) from exc
|
||||
if isinstance(result, ErrorMessage):
|
||||
raise RuntimeError(result.error)
|
||||
return result
|
||||
finally:
|
||||
self._router.unregister(key, waiter)
|
||||
|
||||
def close(self) -> None:
|
||||
self._router.close()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CorrelatedRpcClient",
|
||||
"RpcCorrelationKey",
|
||||
"RpcResultRouter",
|
||||
"RpcWaiter",
|
||||
]
|
||||
+105
-30
@@ -97,6 +97,7 @@ class StagePool:
|
||||
self._stage_vllm_config = stage_vllm_config
|
||||
self._next_replica_id = 0
|
||||
self._request_bindings: dict[str, int] = {}
|
||||
self._unavailable_replicas: set[int] = set()
|
||||
self._replica_metrics: list[_ReplicaMetrics] = [_ReplicaMetrics() for _ in self.clients]
|
||||
self._output_timestamps_by_request: dict[str, list[float]] = {}
|
||||
self._non_empty_first_output_timestamps_by_request: dict[str, float] = {}
|
||||
@@ -206,27 +207,41 @@ class StagePool:
|
||||
return addr
|
||||
return None
|
||||
|
||||
def add_client(self, input_addr: str, client: Any) -> int:
|
||||
def add_client(self, input_addr: str, client: Any, *, replica_id: int | None = None) -> int:
|
||||
"""Register a head-side client for ``input_addr``.
|
||||
|
||||
Returns the assigned ``replica_id`` (index into :attr:`clients`).
|
||||
If the address is already known, replaces the existing client and
|
||||
returns its existing id (this should not happen in practice — the
|
||||
master server assigns unique slots — but the contract is idempotent
|
||||
to keep the dispatch layer robust).
|
||||
A coordinator-provided ``replica_id`` restores that exact slot after
|
||||
unregister/re-register. Without one, a known address keeps its slot
|
||||
and a new address takes the first vacant slot.
|
||||
"""
|
||||
if not input_addr:
|
||||
raise ValueError("input_addr must be a non-empty string")
|
||||
if replica_id is not None and replica_id < 0:
|
||||
raise ValueError("replica_id must be non-negative")
|
||||
|
||||
existing = self._addr_to_replica_id.get(input_addr)
|
||||
if existing is not None:
|
||||
self.clients[existing] = client
|
||||
return existing
|
||||
if replica_id is not None and existing is not None and existing != replica_id:
|
||||
raise ValueError(
|
||||
f"input address {input_addr!r} is already assigned to replica {existing}, not {replica_id}"
|
||||
)
|
||||
if replica_id is None:
|
||||
replica_id = existing
|
||||
if replica_id is None:
|
||||
replica_id = next(
|
||||
(index for index, existing_client in enumerate(self.clients) if existing_client is None),
|
||||
len(self.clients),
|
||||
)
|
||||
|
||||
replica_id = len(self.clients)
|
||||
self.clients.append(client)
|
||||
while len(self.clients) <= replica_id:
|
||||
self.clients.append(None)
|
||||
self._replica_metrics.append(_ReplicaMetrics())
|
||||
for known_addr, known_replica_id in list(self._addr_to_replica_id.items()):
|
||||
if known_replica_id == replica_id and known_addr != input_addr:
|
||||
self._addr_to_replica_id.pop(known_addr, None)
|
||||
self.clients[replica_id] = client
|
||||
self._addr_to_replica_id[input_addr] = replica_id
|
||||
self._replica_metrics.append(_ReplicaMetrics())
|
||||
self._unavailable_replicas.discard(replica_id)
|
||||
return replica_id
|
||||
|
||||
def remove_client(self, input_addr: str) -> Any | None:
|
||||
@@ -234,7 +249,7 @@ class StagePool:
|
||||
|
||||
Slot is marked ``None`` to preserve indices for outstanding bindings.
|
||||
"""
|
||||
replica_id = self._addr_to_replica_id.pop(input_addr, None)
|
||||
replica_id = self._addr_to_replica_id.get(input_addr)
|
||||
if replica_id is None:
|
||||
return None
|
||||
client = self.clients[replica_id]
|
||||
@@ -475,6 +490,39 @@ class StagePool:
|
||||
for request_id in request_ids:
|
||||
self.release_binding(request_id)
|
||||
|
||||
def release_replica_bindings(self, replica_id: int) -> list[str]:
|
||||
"""Drop all route/session bindings owned by one physical replica."""
|
||||
released_request_ids = [
|
||||
request_id
|
||||
for request_id, bound_replica_id in list(self._request_bindings.items())
|
||||
if bound_replica_id == replica_id
|
||||
]
|
||||
released_request_ids.extend(
|
||||
request_id
|
||||
for request_id, input_addr in list(self._affinity.items())
|
||||
if self._addr_to_replica_id.get(input_addr) == replica_id
|
||||
)
|
||||
released_request_ids = list(dict.fromkeys(released_request_ids))
|
||||
for request_id in released_request_ids:
|
||||
self.release_binding(request_id)
|
||||
return released_request_ids
|
||||
|
||||
def mark_replica_unavailable(self, replica_id: int) -> list[str]:
|
||||
"""Evict a failed replica from admission and release its bindings."""
|
||||
if 0 <= replica_id < self.num_replicas:
|
||||
self._unavailable_replicas.add(replica_id)
|
||||
return self.release_replica_bindings(replica_id)
|
||||
|
||||
def is_replica_available(self, replica_id: int) -> bool:
|
||||
return (
|
||||
0 <= replica_id < self.num_replicas
|
||||
and self.clients[replica_id] is not None
|
||||
and replica_id not in self._unavailable_replicas
|
||||
)
|
||||
|
||||
def available_replica_ids(self) -> list[int]:
|
||||
return [replica_id for replica_id in range(self.num_replicas) if self.is_replica_available(replica_id)]
|
||||
|
||||
def select_replica_id(
|
||||
self,
|
||||
request_id: str,
|
||||
@@ -483,23 +531,27 @@ class StagePool:
|
||||
) -> int:
|
||||
"""Pick a replica id for *request_id* and cache the choice (legacy path)."""
|
||||
cached = self.get_bound_replica_id(request_id)
|
||||
if cached is not None and self.clients[cached] is not None:
|
||||
if cached is not None and self.clients[cached] is not None and self.is_replica_available(cached):
|
||||
return cached
|
||||
if cached is not None:
|
||||
# Cached replica is dead/unavailable: drop the stale binding.
|
||||
self.release_binding(request_id)
|
||||
|
||||
chosen: int | None = None
|
||||
if affinity_request_id is not None:
|
||||
parent = self.get_bound_replica_id(affinity_request_id)
|
||||
if parent is not None and self.clients[parent] is not None:
|
||||
if parent is not None and self.clients[parent] is not None and self.is_replica_available(parent):
|
||||
chosen = parent
|
||||
|
||||
if chosen is None:
|
||||
live = self.live_replica_ids()
|
||||
# Prefer replicas that are both live (client up) and available.
|
||||
live = [r for r in self.live_replica_ids() if self.is_replica_available(r)]
|
||||
if not live:
|
||||
raise RuntimeError(f"stage {self.stage_id} has no live replicas")
|
||||
if len(live) == 1:
|
||||
chosen = live[0]
|
||||
else:
|
||||
# Round-robin over live replicas only.
|
||||
# Round-robin over live, available replicas only.
|
||||
start = self._next_replica_id % len(live)
|
||||
chosen = live[start]
|
||||
self._next_replica_id = (self._next_replica_id + 1) % len(live)
|
||||
@@ -539,18 +591,23 @@ class StagePool:
|
||||
non_empty_first_output_ts = (
|
||||
self._non_empty_first_output_timestamps_by_request.pop(request_id, None) if request_id else None
|
||||
)
|
||||
num_tokens_out = count_tokens_from_outputs(request_outputs)
|
||||
native_text_metrics = {}
|
||||
if request_id:
|
||||
pop_native_text_metrics = getattr(self.output_processor, "pop_native_text_metrics", None)
|
||||
if callable(pop_native_text_metrics):
|
||||
native_text_metrics = pop_native_text_metrics(request_id)
|
||||
native_generation_tokens = native_text_metrics.get("num_generation_tokens")
|
||||
num_tokens_out = (
|
||||
max(int(native_generation_tokens), 0)
|
||||
if isinstance(native_generation_tokens, int) and not isinstance(native_generation_tokens, bool)
|
||||
else count_tokens_from_outputs(request_outputs)
|
||||
)
|
||||
output_unit_type = self._infer_output_unit_type(request_outputs, token_count=num_tokens_out)
|
||||
output_unit_count = self._count_output_units(
|
||||
request_outputs,
|
||||
unit_type=output_unit_type,
|
||||
fallback_token_count=num_tokens_out,
|
||||
)
|
||||
native_text_metrics = {}
|
||||
if request_id:
|
||||
pop_native_text_metrics = getattr(self.output_processor, "pop_native_text_metrics", None)
|
||||
if callable(pop_native_text_metrics):
|
||||
native_text_metrics = pop_native_text_metrics(request_id)
|
||||
current_audio_frames, current_audio_sample_rate, _ = self._collect_audio_metrics(request_outputs)
|
||||
accumulated_audio_frames = self._audio_frames_by_request.pop(request_id, 0) if request_id else 0
|
||||
accumulated_audio_sample_rate = self._audio_sample_rate_by_request.pop(request_id, 0) if request_id else 0
|
||||
@@ -1029,14 +1086,28 @@ class StagePool:
|
||||
# Refresh the shared output-processor state before yielding to the
|
||||
# stage client so streaming segments are merged against the latest
|
||||
# prompt/token metadata.
|
||||
self.output_processor.add_request(
|
||||
request=request,
|
||||
prompt=prompt_text,
|
||||
parent_req=None,
|
||||
request_index=0,
|
||||
queue=None,
|
||||
)
|
||||
await self._llm_client(replica_id).add_request_async(request)
|
||||
try:
|
||||
self.output_processor.add_request(
|
||||
request=request,
|
||||
prompt=prompt_text,
|
||||
parent_req=None,
|
||||
request_index=0,
|
||||
queue=None,
|
||||
)
|
||||
await self._llm_client(replica_id).add_request_async(request)
|
||||
except Exception:
|
||||
rollback = getattr(self.output_processor, "remove_request", None)
|
||||
if callable(rollback):
|
||||
try:
|
||||
rollback(request_id)
|
||||
except Exception as rollback_error:
|
||||
logger.warning(
|
||||
"[StagePool] Failed to rollback output processor update for req=%s stage-%s: %s",
|
||||
request_id,
|
||||
self.stage_id,
|
||||
rollback_error,
|
||||
)
|
||||
raise
|
||||
return replica_id
|
||||
|
||||
async def _pick_or_select(
|
||||
@@ -1097,6 +1168,8 @@ class StagePool:
|
||||
timeout_s: float = 0.001,
|
||||
) -> EngineCoreOutputs | None:
|
||||
"""Poll raw EngineCore outputs from one LLM replica once."""
|
||||
if not self.is_replica_available(replica_id):
|
||||
return None
|
||||
raw_client = self.clients[replica_id]
|
||||
if raw_client is None:
|
||||
return None
|
||||
@@ -1120,6 +1193,8 @@ class StagePool:
|
||||
|
||||
def poll_diffusion_output(self, replica_id: int) -> Any | None:
|
||||
"""Drain one ready diffusion output from the given replica if present."""
|
||||
if not self.is_replica_available(replica_id):
|
||||
return None
|
||||
raw_client = self.clients[replica_id]
|
||||
if raw_client is None:
|
||||
return None
|
||||
|
||||
@@ -396,7 +396,6 @@ class StageRuntime:
|
||||
replica_metadata.replica_id = replica_id
|
||||
if launch_mode == "remote" and replica_metadata.stage_type != "diffusion":
|
||||
replica_metadata.runtime_cfg = None
|
||||
|
||||
replicas.append(
|
||||
ReplicaInitPlan(
|
||||
replica_id=replica_id,
|
||||
|
||||
@@ -44,6 +44,12 @@ if TYPE_CHECKING:
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.v1.engine import PauseMode
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import (
|
||||
DuplexFence,
|
||||
DuplexSessionLifecycleMessage,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient
|
||||
from vllm_omni.inputs.data import OmniPromptType
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -143,12 +149,16 @@ class AsyncOmni(EngineClient, OmniBase):
|
||||
self._paused: bool = False
|
||||
self._sleeping_tags: set[str] = set()
|
||||
self._level2_sleeping: bool = False
|
||||
self._duplex_request_client: DuplexRequestClient | None = None
|
||||
self.duplex_lifecycle_events: asyncio.Queue[DuplexSessionLifecycleMessage] = asyncio.Queue()
|
||||
self.final_output_task: asyncio.Task | None = None
|
||||
self.event_resolver = AsyncEventResolver(orchestrator=self)
|
||||
self.config_path = self.engine.config_path
|
||||
self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None)
|
||||
self.input_processor = self.engine.input_processor
|
||||
self.endpoint_restrictions = self.engine.endpoint_restrictions
|
||||
self.duplex_session_config = self.engine.duplex_session_config
|
||||
self.duplex_serving_adapter_path = self.engine.duplex_serving_adapter_path
|
||||
|
||||
stage_index = self._get_comprehension_stage_index()
|
||||
if stage_index is None:
|
||||
@@ -254,6 +264,191 @@ class AsyncOmni(EngineClient, OmniBase):
|
||||
prefix = "" if not external_request_id else f"{external_request_id}-"
|
||||
return f"{prefix}{uuid:.8}"
|
||||
|
||||
async def open_duplex_session_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
session_mode: str = "duplex",
|
||||
capabilities: dict[str, object] | None = None,
|
||||
session_config: dict[str, object] | None = None,
|
||||
runtime_config: dict[str, object] | None = None,
|
||||
fence: DuplexFence,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Open an engine-level duplex session when the backend supports it."""
|
||||
return await self._get_duplex_request_client().open(
|
||||
session_id,
|
||||
session_mode=session_mode,
|
||||
capabilities=capabilities,
|
||||
session_config=session_config,
|
||||
runtime_config=runtime_config,
|
||||
fence=fence,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def append_duplex_input_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
mode: str,
|
||||
payload: object,
|
||||
operation_id: str | None = None,
|
||||
final: bool = False,
|
||||
expected_epoch: int | None = None,
|
||||
fence: DuplexFence,
|
||||
timeout: float | None = 10.0,
|
||||
collect_outputs: bool = True,
|
||||
) -> dict[str, object]:
|
||||
"""Append input to an engine-level duplex session."""
|
||||
return await self._get_duplex_request_client().append(
|
||||
session_id,
|
||||
mode=mode,
|
||||
payload=payload,
|
||||
operation_id=operation_id,
|
||||
final=final,
|
||||
expected_epoch=expected_epoch,
|
||||
fence=fence,
|
||||
timeout=timeout,
|
||||
collect_outputs=collect_outputs,
|
||||
)
|
||||
|
||||
async def collect_duplex_data_plane_outputs_async(
|
||||
self,
|
||||
request_id: str,
|
||||
*,
|
||||
response_stage_id: int | None = None,
|
||||
timeout: float | None = 10.0,
|
||||
) -> list[OmniRequestOutput]:
|
||||
"""Collect the next duplex data-plane output batch for a live request."""
|
||||
return await self._get_duplex_request_client().collect_registered_outputs(
|
||||
request_id,
|
||||
response_stage_id=response_stage_id,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def signal_duplex_turn_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
event: str,
|
||||
fence: DuplexFence,
|
||||
next_fence: DuplexFence | None = None,
|
||||
session_config: dict[str, object] | None = None,
|
||||
runtime_config: dict[str, object] | None = None,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Send a turn/control signal to an engine-level duplex session."""
|
||||
return await self._get_duplex_request_client().signal(
|
||||
session_id,
|
||||
event=event,
|
||||
fence=fence,
|
||||
next_fence=next_fence,
|
||||
session_config=session_config,
|
||||
runtime_config=runtime_config,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def close_duplex_session_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
reason: str = "client_close",
|
||||
fence: DuplexFence,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
"""Close an engine-level duplex session."""
|
||||
return await self._get_duplex_request_client().close(
|
||||
session_id,
|
||||
reason=reason,
|
||||
fence=fence,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def touch_duplex_session_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
activity: DuplexLeaseActivity,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
return await self._get_duplex_request_client().touch(
|
||||
session_id,
|
||||
fence=fence,
|
||||
activity=activity,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def resume_duplex_session_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
expected_lease_generation: int,
|
||||
timeout: float | None = 10.0,
|
||||
) -> dict[str, object]:
|
||||
return await self._get_duplex_request_client().resume(
|
||||
session_id,
|
||||
fence=fence,
|
||||
expected_lease_generation=expected_lease_generation,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def _get_duplex_request_client(self) -> DuplexRequestClient:
|
||||
from vllm_omni.experimental.fullduplex.request_client import (
|
||||
DuplexRequestClient,
|
||||
DuplexRequestOutputPort,
|
||||
)
|
||||
|
||||
client = getattr(self, "_duplex_request_client", None)
|
||||
if client is None:
|
||||
engine = getattr(self, "engine", None)
|
||||
client = DuplexRequestClient(
|
||||
engine,
|
||||
DuplexRequestOutputPort(
|
||||
request_states=getattr(self, "request_states", {}),
|
||||
num_stages=getattr(engine, "num_stages", 1),
|
||||
log_stats=getattr(self, "log_stats", False),
|
||||
start_output_handler=self._final_output_handler,
|
||||
process_single_result=self._process_single_result,
|
||||
),
|
||||
)
|
||||
self._duplex_request_client = client
|
||||
return client
|
||||
|
||||
@staticmethod
|
||||
def _duplex_data_plane_request_info(result: dict[str, object]) -> tuple[str | None, int | None]:
|
||||
from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient
|
||||
|
||||
return DuplexRequestClient.request_info(result)
|
||||
|
||||
async def _collect_duplex_data_plane_outputs(
|
||||
self,
|
||||
request_id: str,
|
||||
req_state: ClientRequestState,
|
||||
*,
|
||||
response_stage_id: int | None,
|
||||
timeout: float | None,
|
||||
) -> list[OmniRequestOutput]:
|
||||
return await self._get_duplex_request_client().collect_outputs(
|
||||
request_id,
|
||||
req_state,
|
||||
response_stage_id=response_stage_id,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _is_direct_duplex_data_plane_response(cls, output: object) -> bool:
|
||||
from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient
|
||||
|
||||
return DuplexRequestClient.is_direct_response(output)
|
||||
|
||||
@classmethod
|
||||
def _duplex_multimodal_output(cls, output: object) -> dict[str, object]:
|
||||
from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient
|
||||
|
||||
return DuplexRequestClient.multimodal_output(output)
|
||||
|
||||
# ==================== Generate Method ====================
|
||||
|
||||
async def generate(
|
||||
@@ -693,6 +888,10 @@ class AsyncOmni(EngineClient, OmniBase):
|
||||
await self.event_resolver.resolve(msg)
|
||||
continue
|
||||
|
||||
if getattr(msg, "type", None) == "duplex_session_lifecycle":
|
||||
await self.duplex_lifecycle_events.put(msg)
|
||||
continue
|
||||
|
||||
if isinstance(msg, ErrorMessage) and not msg.fatal:
|
||||
req_state = self.request_states.get(msg.request_id)
|
||||
if req_state is not None:
|
||||
|
||||
@@ -4,7 +4,7 @@ from vllm_omni.metrics import OrchestratorAggregator
|
||||
|
||||
|
||||
class ClientRequestState:
|
||||
"""Tracks the state of an individual request in the orchestrator."""
|
||||
"""Tracks one entrypoint request and its output queue."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
||||
@@ -473,6 +473,7 @@ class OmniBase(PDDisaggregationMixin):
|
||||
req_start_ts: dict[str, float],
|
||||
wall_start_ts: float,
|
||||
final_stage_id_for_e2e: int,
|
||||
stage_event_cursor: int = 0,
|
||||
) -> OmniRequestOutput | None:
|
||||
req_id = result.request_id
|
||||
engine_outputs = result.engine_outputs
|
||||
@@ -481,7 +482,7 @@ class OmniBase(PDDisaggregationMixin):
|
||||
|
||||
# Merge AR stage timing from OrchestratorAggregator.stage_events
|
||||
if self._enable_ar_profiler:
|
||||
ar_events = metrics.stage_events.get(str(req_id), [])
|
||||
ar_events = metrics.stage_events.get(str(req_id), [])[stage_event_cursor:]
|
||||
for evt in ar_events:
|
||||
if evt.stage_id != stage_id:
|
||||
stage_durations[f"ar_stage_{evt.stage_id}"] = evt.stage_gen_time_ms / 1000.0
|
||||
@@ -494,7 +495,7 @@ class OmniBase(PDDisaggregationMixin):
|
||||
stage_durations[key] = value
|
||||
|
||||
# Merge per-stage gen times into stage_durations
|
||||
for evt in metrics.stage_events.get(str(req_id), []):
|
||||
for evt in metrics.stage_events.get(str(req_id), [])[stage_event_cursor:]:
|
||||
key = f"stage_{evt.stage_id}_gen_ms"
|
||||
if key not in stage_durations:
|
||||
stage_durations[key] = evt.stage_gen_time_ms
|
||||
@@ -546,7 +547,7 @@ class OmniBase(PDDisaggregationMixin):
|
||||
# Token counters — aggregate across all stages for this request.
|
||||
_prompt_tok = 0
|
||||
_gen_tok = 0
|
||||
for evt in metrics.stage_events.get(rid_key, []):
|
||||
for evt in metrics.stage_events.get(rid_key, [])[stage_event_cursor:]:
|
||||
if evt.stage_id == 0:
|
||||
_prompt_tok += int(evt.num_tokens_in)
|
||||
_gen_tok += int(evt.num_tokens_out)
|
||||
@@ -581,7 +582,7 @@ class OmniBase(PDDisaggregationMixin):
|
||||
response_metrics: dict[str, Any] = {}
|
||||
stage_metrics: dict[str, dict[str, Any]] = {}
|
||||
rid_key = str(req_id)
|
||||
for evt in metrics.stage_events.get(rid_key, []):
|
||||
for evt in metrics.stage_events.get(rid_key, [])[stage_event_cursor:]:
|
||||
if evt.stage_id is None:
|
||||
continue
|
||||
sid = int(evt.stage_id)
|
||||
@@ -604,6 +605,7 @@ class OmniBase(PDDisaggregationMixin):
|
||||
response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"]
|
||||
return OmniRequestOutput(
|
||||
request_id=req_id or "",
|
||||
finished=finished,
|
||||
stage_id=stage_id,
|
||||
final_output_type=output_type,
|
||||
request_output=engine_outputs,
|
||||
@@ -616,7 +618,6 @@ class OmniBase(PDDisaggregationMixin):
|
||||
metrics=response_metrics,
|
||||
stage_durations=stage_durations,
|
||||
peak_memory_mb=peak_memory_mb,
|
||||
finished=finished,
|
||||
)
|
||||
|
||||
def shutdown(self, timeout: float | None = None) -> None:
|
||||
|
||||
@@ -94,6 +94,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
|
||||
from vllm_omni.config.endpoint_policy import shutdown_unsupported_routes
|
||||
from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.openai.duplex_capability import should_enable_duplex_endpoint
|
||||
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
|
||||
from vllm_omni.entrypoints.openai.image_api_utils import (
|
||||
SUPPORTED_LAYERED_RESOLUTIONS,
|
||||
@@ -793,6 +794,7 @@ async def omni_init_app_state(
|
||||
model_name=model_name,
|
||||
stage_configs=diffusion_stage_configs,
|
||||
)
|
||||
state.openai_serving_duplex = None
|
||||
state.openai_streaming_speech = None
|
||||
state.openai_streaming_video = None
|
||||
state.openai_serving_realtime_robot = ServingRealtimeRobotOpenPI.create_policy_server(
|
||||
@@ -1112,6 +1114,18 @@ async def omni_init_app_state(
|
||||
if state.openai_serving_chat is not None
|
||||
else None
|
||||
)
|
||||
state.openai_serving_duplex = None
|
||||
if state.openai_serving_chat is not None and should_enable_duplex_endpoint(
|
||||
state.stage_configs,
|
||||
config_path=getattr(args, "stage_configs_path", None) or getattr(args, "deploy_config", None),
|
||||
):
|
||||
from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler
|
||||
|
||||
state.openai_serving_duplex = OmniDuplexSessionHandler(
|
||||
chat_service=state.openai_serving_chat,
|
||||
duplex_session_config=getattr(engine_client, "duplex_session_config", None),
|
||||
serving_runtime_adapter_path=getattr(engine_client, "duplex_serving_adapter_path", None),
|
||||
)
|
||||
state.openai_serving_realtime = OpenAIServingRealtime(
|
||||
engine_client=engine_client,
|
||||
models=state.openai_serving_models,
|
||||
@@ -1612,6 +1626,31 @@ async def streaming_video_output(websocket: WebSocket):
|
||||
@router.websocket("/v1/realtime")
|
||||
async def realtime_websocket(websocket: WebSocket):
|
||||
"""WebSocket endpoint for OpenAI-style realtime interactions."""
|
||||
duplex_handler = getattr(websocket.app.state, "openai_serving_duplex", None)
|
||||
duplex_query = websocket.query_params.get("duplex")
|
||||
use_duplex_realtime = (
|
||||
duplex_handler is not None and isinstance(duplex_query, str) and duplex_query.lower() in {"1", "true", "on"}
|
||||
)
|
||||
if use_duplex_realtime and duplex_handler is not None:
|
||||
await duplex_handler.handle_realtime_session(websocket)
|
||||
return
|
||||
|
||||
engine_client = getattr(websocket.app.state, "engine_client", None)
|
||||
if engine_client is not None and getattr(engine_client, "async_chunk", False):
|
||||
await websocket.accept()
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "error",
|
||||
"error": (
|
||||
"The /v1/realtime API is not supported when async_chunk is enabled on the server. "
|
||||
"Use a stage configuration with async_chunk disabled and restart the server before using "
|
||||
"this endpoint."
|
||||
),
|
||||
"code": "unsupported",
|
||||
}
|
||||
)
|
||||
await websocket.close()
|
||||
return
|
||||
serving = getattr(websocket.app.state, "openai_serving_realtime", None)
|
||||
if serving is None:
|
||||
await websocket.accept()
|
||||
@@ -1639,6 +1678,18 @@ async def realtime_robot_openpi(websocket: WebSocket):
|
||||
await connection.handle_connection()
|
||||
|
||||
|
||||
@router.websocket("/v1/duplex")
|
||||
async def duplex_websocket(websocket: WebSocket):
|
||||
"""WebSocket endpoint for vLLM-Omni duplex session control."""
|
||||
handler = getattr(websocket.app.state, "openai_serving_duplex", None)
|
||||
if handler is None:
|
||||
await websocket.accept()
|
||||
await websocket.send_json({"type": "error", "error": "Duplex API is not available", "code": "unsupported"})
|
||||
await websocket.close()
|
||||
return
|
||||
await handler.handle_session(websocket)
|
||||
|
||||
|
||||
# Health and Model endpoints for diffusion mode
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def should_enable_duplex_endpoint(
|
||||
stage_configs: list | None,
|
||||
*,
|
||||
config_path: str | None = None,
|
||||
) -> bool:
|
||||
"""Enable duplex routes only for deployments that explicitly opt in."""
|
||||
if stage_configs:
|
||||
for stage in stage_configs:
|
||||
session_mode = (
|
||||
stage.get("session_mode") if isinstance(stage, dict) else getattr(stage, "session_mode", None)
|
||||
)
|
||||
if session_mode == "duplex":
|
||||
return True
|
||||
if config_path:
|
||||
try:
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
raw_config = OmegaConf.load(config_path)
|
||||
session_mode = raw_config.get("session_mode") if hasattr(raw_config, "get") else None
|
||||
if session_mode == "duplex":
|
||||
return True
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to inspect duplex session_mode from %s: %s", config_path, exc)
|
||||
return False
|
||||
|
||||
|
||||
__all__ = ["should_enable_duplex_endpoint"]
|
||||
@@ -79,7 +79,7 @@ from vllm.outputs import RequestOutput
|
||||
from vllm.reasoning import ReasoningParser
|
||||
from vllm.renderers import BaseRenderer, merge_kwargs
|
||||
from vllm.renderers.inputs import TokPrompt
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tokenizers import TokenizerLike as AnyTokenizer
|
||||
from vllm.tokenizers.mistral import (
|
||||
@@ -187,6 +187,63 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
tp_cls = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None
|
||||
return request.tools and tp_cls is not None and self.enable_auto_tools and request.tool_choice in ["auto", None]
|
||||
|
||||
@staticmethod
|
||||
def _stage_get(obj: Any, key: str, default: Any = None) -> Any:
|
||||
if obj is None:
|
||||
return default
|
||||
if hasattr(obj, "get"):
|
||||
return obj.get(key, default)
|
||||
return getattr(obj, key, default)
|
||||
|
||||
def _has_minicpmo45_stage(self) -> bool:
|
||||
for stage in getattr(self.engine_client, "stage_configs", []) or []:
|
||||
engine_args = self._stage_get(stage, "engine_args")
|
||||
model_arch = self._stage_get(engine_args, "model_arch")
|
||||
if model_arch == "MiniCPMO45OmniForConditionalGeneration":
|
||||
return True
|
||||
return False
|
||||
|
||||
def _fix_minicpmo45_audio_stream_output_kinds(
|
||||
self,
|
||||
sampling_params_list: list[Any],
|
||||
output_modalities: list[str] | tuple[str, ...] | None,
|
||||
) -> list[Any]:
|
||||
"""Keep MiniCPM-o 4.5's thinker final-only during audio streaming.
|
||||
|
||||
Chat ``stream=true`` normally coerces every AR stage to DELTA output.
|
||||
That is correct for token-level text streaming, but MiniCPM-o 4.5's
|
||||
Talker consumes the completed TTS span plus aligned thinker hidden
|
||||
states. If Stage0 is DELTA, llm2tts receives one generated token at a
|
||||
time and repeatedly starts tiny, independent TTS streams. The API then
|
||||
never represents one coherent audio response.
|
||||
|
||||
For this model, stream the audio stage, not the thinker-to-talker
|
||||
boundary.
|
||||
"""
|
||||
if not output_modalities or "audio" not in output_modalities:
|
||||
return sampling_params_list
|
||||
if not self._has_minicpmo45_stage():
|
||||
return sampling_params_list
|
||||
|
||||
stage_configs = getattr(self.engine_client, "stage_configs", []) or []
|
||||
for idx, stage in enumerate(stage_configs):
|
||||
if idx >= len(sampling_params_list):
|
||||
break
|
||||
sp = sampling_params_list[idx]
|
||||
if not hasattr(sp, "output_kind"):
|
||||
continue
|
||||
|
||||
engine_args = self._stage_get(stage, "engine_args")
|
||||
if self._stage_get(engine_args, "model_arch") != "MiniCPMO45OmniForConditionalGeneration":
|
||||
continue
|
||||
|
||||
model_stage = self._stage_get(engine_args, "model_stage")
|
||||
if model_stage == "llm":
|
||||
sp.output_kind = RequestOutputKind.FINAL_ONLY
|
||||
elif model_stage == "tts":
|
||||
sp.output_kind = RequestOutputKind.DELTA
|
||||
return sampling_params_list
|
||||
|
||||
@classmethod
|
||||
def for_diffusion(
|
||||
cls,
|
||||
@@ -605,6 +662,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
# to delta to ensure emitted outputs are correctly drained. Otherwise
|
||||
# convert cumulative to Final Only to ensure the output is correct.
|
||||
sampling_params_list = coerce_param_message_types(sampling_params_list, request.stream)
|
||||
if request.stream:
|
||||
sampling_params_list = self._fix_minicpmo45_audio_stream_output_kinds(
|
||||
sampling_params_list,
|
||||
output_modalities,
|
||||
)
|
||||
|
||||
# Apply user-specified overrides to diffusion stage(s) for image generation
|
||||
for idx, sp in enumerate(sampling_params_list):
|
||||
|
||||
@@ -10,12 +10,14 @@ via :class:`VideoStreamPipelineHooks`.
|
||||
Protocol:
|
||||
Client -> Server:
|
||||
{"type": "session.config", ...} # Session config (sent once)
|
||||
{"type": "video.frame", "data": "..."} # base64 JPEG/PNG frame
|
||||
{"type": "video.frame", "data": "...", "frame_id": "...", "pts_ms": 0}
|
||||
{"type": "audio.chunk", "data": "..."} # base64 PCM16 16kHz mono
|
||||
{"type": "video.query", "text": "..."} # Submit query about buffered frames
|
||||
{"type": "video.done"} # End of session
|
||||
|
||||
Server -> Client:
|
||||
{"type": "video.frame.ack", ...} # when frame_id is provided
|
||||
{"type": "video.frames.consumed", ...} # after first engine output
|
||||
{"type": "response.start"}
|
||||
{"type": "response.text.delta", "delta": "..."}
|
||||
{"type": "response.text.done", "text": "..."}
|
||||
@@ -215,6 +217,7 @@ class OmniStreamingVideoHandler:
|
||||
return
|
||||
|
||||
frame_buffer: list[str] = [] # base64-encoded JPEG frames
|
||||
frame_metadata: list[dict[str, Any]] = []
|
||||
# Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index.
|
||||
frame_pil_cache: dict[str, tuple[Any, str] | object] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME
|
||||
frame_filter = (
|
||||
@@ -259,6 +262,8 @@ class OmniStreamingVideoHandler:
|
||||
if msg_type.startswith("_internal."):
|
||||
await self._send_error(websocket, f"Unknown type: {msg_type}")
|
||||
continue
|
||||
if msg_type == "video.frame":
|
||||
msg["_receiver_received_ts_ms"] = _time.monotonic() * 1000
|
||||
|
||||
await msg_queue.put(msg)
|
||||
if msg.get("type") == "video.done":
|
||||
@@ -308,6 +313,7 @@ class OmniStreamingVideoHandler:
|
||||
active_request_id = request_id
|
||||
interrupt_event.clear()
|
||||
query_frames = list(frame_buffer)
|
||||
query_frame_metadata = list(frame_metadata)
|
||||
query_audio_buffer = bytearray(audio_buffer)
|
||||
audio_buffer.clear()
|
||||
query_prewarmed_frames = dict(frame_pil_cache)
|
||||
@@ -315,6 +321,9 @@ class OmniStreamingVideoHandler:
|
||||
async def _run_query() -> None:
|
||||
nonlocal active_request_id, prev_request_id
|
||||
try:
|
||||
process_kwargs: dict[str, Any] = {}
|
||||
if any(metadata.get("frame_id") for metadata in query_frame_metadata):
|
||||
process_kwargs["frame_metadata"] = query_frame_metadata
|
||||
await self._process_query(
|
||||
websocket,
|
||||
config,
|
||||
@@ -325,6 +334,7 @@ class OmniStreamingVideoHandler:
|
||||
request_id,
|
||||
interrupt_event,
|
||||
query_prewarmed_frames,
|
||||
**process_kwargs,
|
||||
)
|
||||
finally:
|
||||
if active_request_id == request_id:
|
||||
@@ -349,7 +359,11 @@ class OmniStreamingVideoHandler:
|
||||
frame_data = msg.get("b64", "")
|
||||
removed = frame_data in frame_buffer
|
||||
if removed:
|
||||
frame_buffer[:] = [f for f in frame_buffer if f != frame_data]
|
||||
retained_indices = [
|
||||
index for index, frame in enumerate(frame_buffer) if frame != frame_data
|
||||
]
|
||||
frame_buffer[:] = [frame_buffer[index] for index in retained_indices]
|
||||
frame_metadata[:] = [frame_metadata[index] for index in retained_indices]
|
||||
if frame_pil_cache.get(frame_data) is _BAD_FRAME:
|
||||
frame_pil_cache.pop(frame_data, None)
|
||||
if removed:
|
||||
@@ -370,16 +384,43 @@ class OmniStreamingVideoHandler:
|
||||
if frame_filter is not None:
|
||||
try:
|
||||
if not frame_filter.should_retain(raw_bytes):
|
||||
await self._send_frame_ack(
|
||||
websocket,
|
||||
msg,
|
||||
accepted=False,
|
||||
buffered_frames=len(frame_buffer),
|
||||
reason="filtered",
|
||||
)
|
||||
continue
|
||||
except Exception:
|
||||
await self._send_error(websocket, "Invalid image data")
|
||||
continue
|
||||
max_buf = config.max_frames
|
||||
dropped_frame_id: str | None = None
|
||||
if len(frame_buffer) >= max_buf:
|
||||
dropped = frame_buffer.pop(0)
|
||||
dropped_metadata = frame_metadata.pop(0)
|
||||
dropped_frame_id = dropped_metadata.get("frame_id")
|
||||
frame_pil_cache.pop(dropped, None)
|
||||
frame_buffer.append(frame_data)
|
||||
frame_metadata.append(
|
||||
{
|
||||
"frame_id": msg.get("frame_id"),
|
||||
"pts_ms": msg.get("pts_ms"),
|
||||
"source_pts_ms": msg.get("source_pts_ms"),
|
||||
"quality_profile": msg.get("quality_profile"),
|
||||
"capture_ts_ms": msg.get("capture_ts_ms"),
|
||||
"receiver_received_ts_ms": msg.get("_receiver_received_ts_ms"),
|
||||
}
|
||||
)
|
||||
self.on_frame_buffered(raw_bytes, frame_data, message_history, config)
|
||||
await self._send_frame_ack(
|
||||
websocket,
|
||||
msg,
|
||||
accepted=True,
|
||||
buffered_frames=len(frame_buffer),
|
||||
dropped_frame_id=dropped_frame_id,
|
||||
)
|
||||
# Prewarm: decode PIL off the event loop so query-time chat_template
|
||||
# can skip base64+Image.open. uuid=md5 lets mm_cache dedupe identical frames.
|
||||
if frame_data not in frame_pil_cache:
|
||||
@@ -537,6 +578,7 @@ class OmniStreamingVideoHandler:
|
||||
request_id: str,
|
||||
interrupt_event: asyncio.Event,
|
||||
prewarmed_frames: dict[str, tuple[Any, str]],
|
||||
frame_metadata: list[dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
"""Build prompt, run inference, stream text + audio response."""
|
||||
|
||||
@@ -544,6 +586,9 @@ class OmniStreamingVideoHandler:
|
||||
await self._send_error(websocket, "Streaming video requires an engine client")
|
||||
return
|
||||
|
||||
engine_kwargs: dict[str, Any] = {}
|
||||
if frame_metadata:
|
||||
engine_kwargs["frame_metadata"] = frame_metadata
|
||||
await self._process_query_engine(
|
||||
websocket,
|
||||
config,
|
||||
@@ -554,6 +599,7 @@ class OmniStreamingVideoHandler:
|
||||
request_id,
|
||||
interrupt_event,
|
||||
prewarmed_frames,
|
||||
**engine_kwargs,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@@ -571,6 +617,7 @@ class OmniStreamingVideoHandler:
|
||||
request_id: str,
|
||||
interrupt_event: asyncio.Event,
|
||||
prewarmed_frames: dict[str, tuple[Any, str]],
|
||||
frame_metadata: list[dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
"""Direct engine_client.generate() path for async_chunk audio."""
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
@@ -613,6 +660,9 @@ class OmniStreamingVideoHandler:
|
||||
except Exception as e:
|
||||
await self._send_error(websocket, f"Preprocess failed: {e}")
|
||||
return
|
||||
decoded_ready_ts_ms = _time.monotonic() * 1000
|
||||
selected_metadata = self._sample_frame_metadata(frame_metadata or [], config.num_frames)
|
||||
model_selected_ts_ms = _time.monotonic() * 1000
|
||||
|
||||
await websocket.send_json({"type": "response.start"})
|
||||
text_parts: list[str] = []
|
||||
@@ -623,6 +673,7 @@ class OmniStreamingVideoHandler:
|
||||
audio_chunks_drained = 0
|
||||
previous_text = ""
|
||||
interrupted = False
|
||||
frames_consumed_sent = False
|
||||
t_start = _time.monotonic()
|
||||
t_first_text = None
|
||||
t_first_audio = None
|
||||
@@ -652,6 +703,33 @@ class OmniStreamingVideoHandler:
|
||||
if not isinstance(output, OmniRequestOutput):
|
||||
continue
|
||||
|
||||
if not frames_consumed_sent and frame_metadata:
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "video.frames.consumed",
|
||||
"request_id": request_id,
|
||||
"model_selected_ts_ms": model_selected_ts_ms,
|
||||
"frame_ids": [
|
||||
metadata["frame_id"]
|
||||
for metadata in selected_metadata
|
||||
if isinstance(metadata.get("frame_id"), str)
|
||||
],
|
||||
"frames": [
|
||||
{
|
||||
"frame_id": metadata.get("frame_id"),
|
||||
"pts_ms": metadata.get("pts_ms"),
|
||||
"source_pts_ms": metadata.get("source_pts_ms"),
|
||||
"quality_profile": metadata.get("quality_profile"),
|
||||
"receiver_received_ts_ms": metadata.get("receiver_received_ts_ms"),
|
||||
"decoded_ready_ts_ms": decoded_ready_ts_ms,
|
||||
}
|
||||
for metadata in selected_metadata
|
||||
],
|
||||
"latest_pts_ms": selected_metadata[-1].get("pts_ms") if selected_metadata else None,
|
||||
}
|
||||
)
|
||||
frames_consumed_sent = True
|
||||
|
||||
out_type = getattr(output, "final_output_type", "text")
|
||||
|
||||
if out_type == "audio":
|
||||
@@ -971,3 +1049,42 @@ class OmniStreamingVideoHandler:
|
||||
await websocket.send_json({"type": "error", "message": message})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _sample_frame_metadata(
|
||||
frame_metadata: list[dict[str, Any]],
|
||||
num_frames: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
if len(frame_metadata) <= num_frames:
|
||||
return list(frame_metadata)
|
||||
stride = max(1, len(frame_metadata) // num_frames)
|
||||
indices = [index * stride for index in range(num_frames - 1)] + [len(frame_metadata) - 1]
|
||||
return [frame_metadata[index] for index in indices]
|
||||
|
||||
@staticmethod
|
||||
async def _send_frame_ack(
|
||||
websocket: WebSocket,
|
||||
message: Mapping[str, Any],
|
||||
*,
|
||||
accepted: bool,
|
||||
buffered_frames: int,
|
||||
reason: str | None = None,
|
||||
dropped_frame_id: str | None = None,
|
||||
) -> None:
|
||||
frame_id = message.get("frame_id")
|
||||
if not isinstance(frame_id, str) or not frame_id:
|
||||
return
|
||||
ack: dict[str, Any] = {
|
||||
"type": "video.frame.ack",
|
||||
"frame_id": frame_id,
|
||||
"pts_ms": message.get("pts_ms"),
|
||||
"capture_ts_ms": message.get("capture_ts_ms"),
|
||||
"accepted": accepted,
|
||||
"buffered_frames": buffered_frames,
|
||||
"server_receive_ts_ms": message.get("_receiver_received_ts_ms", _time.monotonic() * 1000),
|
||||
}
|
||||
if reason is not None:
|
||||
ack["reason"] = reason
|
||||
if dropped_frame_id is not None:
|
||||
ack["dropped_frame_id"] = dropped_frame_id
|
||||
await websocket.send_json(ack)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1,80 +1,28 @@
|
||||
# Full-duplex interaction framework
|
||||
# Experimental Full-Duplex Runtime
|
||||
|
||||
A model-agnostic framework for real-time, full-duplex (streaming-in / streaming-out)
|
||||
model interaction, plus the JoyVL implementation built on it.
|
||||
This package contains two experimental integrations:
|
||||
|
||||
To **run the JoyVL model**, see the recipe:
|
||||
- the existing JoyVL framework and example integration;
|
||||
- the MiniCPM-o 4.5 native audio path used by `/v1/duplex` and
|
||||
`/v1/realtime?duplex=1`.
|
||||
|
||||
For the MiniCPM active runtime path, lifecycle invariants, capability boundary,
|
||||
Realtime response contract, and validation scope, see [`DESIGN.md`](DESIGN.md).
|
||||
|
||||
To run JoyVL, see
|
||||
[`recipes/JD/JoyAI-VL-Interaction.md`](../../../recipes/JD/JoyAI-VL-Interaction.md).
|
||||
This README covers the framework itself and how to add a new model.
|
||||
|
||||
## Layout
|
||||
## Package boundaries
|
||||
|
||||
```
|
||||
vllm_omni/experimental/fullduplex/
|
||||
core/ generic full-duplex framework (model-agnostic): DuplexRuntime (event
|
||||
loop + epoch barge-in), DuplexSession, DuplexAdapter (ABC), protocol
|
||||
joyvl/ JoyVL implementation (model-specific):
|
||||
adapter.py JoyVLDuplexAdapter (implements core.DuplexAdapter)
|
||||
decision/ policy + output_parser + prompts (speak/silence/delegate)
|
||||
memory/ InteractionBrain — 3-tier summary memory (async)
|
||||
serving/ OpenAI-compatible HTTP orchestrator
|
||||
bridges/ model backend + delegation
|
||||
```text
|
||||
core/ existing JoyVL framework and experimental compatibility exports
|
||||
engine/ AsyncOmni/orchestrator scheduler data-plane adapter
|
||||
openai/ WebSocket transport, Realtime projection, and audio codecs
|
||||
minicpmo45/ MiniCPM input framing, policy, compatibility, and Stage0 state
|
||||
joyvl/ JoyVL model-specific integration
|
||||
```
|
||||
|
||||
`core/` is the only part shared across models; data planes differ by model and are
|
||||
intentionally not shared.
|
||||
|
||||
## Scope
|
||||
|
||||
The runnable serving path is `joyvl/serving/` driving `decision/` + `memory/` directly.
|
||||
`joyvl/adapter.py` + `core/` are a **demonstration** of how a model plugs into the
|
||||
generic full-duplex framework (exercised by `tests/fullduplex/`), not the serving path
|
||||
the HTTP orchestrator currently uses — a fused-audio model (e.g. MiniCPM-o) is the case
|
||||
`core/` is built for.
|
||||
|
||||
## Adding a full-duplex model
|
||||
|
||||
The seam is `core.DuplexAdapter`. `core/` owns the session lifecycle, epoch-based
|
||||
barge-in, playback cursor, and the event protocol — you implement only model policy.
|
||||
|
||||
1. Create a sibling package `vllm_omni/experimental/fullduplex/<model>/` next to `joyvl/`; keep
|
||||
model-specific code there and do not touch `core/`.
|
||||
2. Implement one `DuplexAdapter` (three required methods; the rest have defaults):
|
||||
|
||||
```python
|
||||
from collections.abc import AsyncIterator
|
||||
from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter, DuplexCapability, OutputChunk
|
||||
from vllm_omni.experimental.fullduplex.core.session import DuplexSession
|
||||
|
||||
class MyModelAdapter(DuplexAdapter):
|
||||
def capabilities(self) -> DuplexCapability:
|
||||
return DuplexCapability(
|
||||
input_modalities=frozenset({"audio", "text"}),
|
||||
output_modalities=frozenset({"audio", "text"}),
|
||||
proactive=True, # speak without being asked?
|
||||
)
|
||||
|
||||
async def on_input(self, session: DuplexSession, modality: str, data) -> None:
|
||||
... # buffer/route an incoming chunk
|
||||
|
||||
async def respond(self, session: DuplexSession) -> AsyncIterator[OutputChunk]:
|
||||
async for piece in self._model_stream(session):
|
||||
yield OutputChunk(modality="audio", data=piece) # runtime drops stale
|
||||
# chunks after a barge-in
|
||||
|
||||
# optional: should_respond / on_barge_in / on_playback_ack
|
||||
```
|
||||
|
||||
3. Run it through the shared runtime — no new control-plane code:
|
||||
|
||||
```python
|
||||
from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime
|
||||
from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig
|
||||
|
||||
rt = DuplexRuntime(DuplexSession("sid", DuplexSessionConfig()), MyModelAdapter())
|
||||
await rt.run(input_events, emit)
|
||||
```
|
||||
|
||||
`joyvl/adapter.py` is the worked demonstration (currently exercised by tests, not the
|
||||
HTTP serving path). Promote a helper from a model package up into `core/` only once a
|
||||
second model actually needs it.
|
||||
MiniCPM does not run through the removed experimental `core.DuplexRuntime`
|
||||
facade. Its active path uses the `openai` session controller, the experimental
|
||||
engine contracts, the standard scheduler/model runners, and an injected
|
||||
MiniCPM-specific runtime extension from `minicpmo45/runtime.py`.
|
||||
@@ -0,0 +1,417 @@
|
||||
"""Reusable Realtime WebSocket, PCM, and event helpers for MiniCPM-o demos."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
import wave
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
try:
|
||||
import websockets
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
except ImportError as exc: # pragma: no cover - example dependency
|
||||
raise SystemExit("Install websockets first: pip install websockets") from exc
|
||||
|
||||
PCM16_SAMPLE_RATE = 16_000
|
||||
PCM16_BYTES_PER_SAMPLE = 2
|
||||
|
||||
|
||||
def _rounded_ms(value: float) -> float:
|
||||
return round(float(value), 3)
|
||||
|
||||
|
||||
def _interval_summary(values: list[float]) -> dict[str, float | int]:
|
||||
clean = sorted(_rounded_ms(value) for value in values if math.isfinite(value) and value >= 0)
|
||||
if not clean:
|
||||
return {"count": 0, "mean": 0.0, "p50": 0.0, "p95": 0.0, "max": 0.0}
|
||||
|
||||
def nearest_rank(percentile: float) -> float:
|
||||
index = max(0, math.ceil(percentile * len(clean)) - 1)
|
||||
return clean[min(index, len(clean) - 1)]
|
||||
|
||||
return {
|
||||
"count": len(clean),
|
||||
"mean": _rounded_ms(sum(clean) / len(clean)),
|
||||
"p50": nearest_rank(0.50),
|
||||
"p95": nearest_rank(0.95),
|
||||
"max": clean[-1],
|
||||
}
|
||||
|
||||
|
||||
def _event_stage_metrics(event: dict[str, object]) -> dict[str, object] | None:
|
||||
candidates: list[object] = [event.get("vllm_omni")]
|
||||
metadata = event.get("metadata")
|
||||
if isinstance(metadata, dict):
|
||||
candidates.extend((metadata, metadata.get("vllm_omni")))
|
||||
response = event.get("response")
|
||||
if isinstance(response, dict):
|
||||
response_metadata = response.get("metadata")
|
||||
if isinstance(response_metadata, dict):
|
||||
candidates.extend((response_metadata, response_metadata.get("vllm_omni")))
|
||||
for candidate in candidates:
|
||||
if not isinstance(candidate, dict):
|
||||
continue
|
||||
stage_metrics = candidate.get("stage_metrics")
|
||||
if isinstance(stage_metrics, dict):
|
||||
return stage_metrics
|
||||
return None
|
||||
|
||||
|
||||
def build_realtime_url(
|
||||
url: str,
|
||||
model: str,
|
||||
*,
|
||||
autostart: bool | None = None,
|
||||
session_id: str | None = None,
|
||||
) -> str:
|
||||
"""Add the explicit native-duplex query parameters to a Realtime URL."""
|
||||
parts = urlsplit(url)
|
||||
query = dict(parse_qsl(parts.query, keep_blank_values=True))
|
||||
query.setdefault("duplex", "1")
|
||||
query.setdefault("model", model)
|
||||
query.setdefault("minicpmo45_native_duplex", "1")
|
||||
if autostart is not None:
|
||||
query.setdefault("autostart", "1" if autostart else "0")
|
||||
if session_id:
|
||||
query.setdefault("session_id", session_id)
|
||||
return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
|
||||
|
||||
|
||||
def read_pcm16_wav(path: Path) -> bytes:
|
||||
"""Read a mono, uncompressed, 16 kHz PCM16 WAV file."""
|
||||
with wave.open(str(path), "rb") as wav_file:
|
||||
if wav_file.getnchannels() != 1:
|
||||
raise ValueError("input WAV must be mono")
|
||||
if wav_file.getsampwidth() != PCM16_BYTES_PER_SAMPLE:
|
||||
raise ValueError("input WAV must be 16-bit PCM")
|
||||
if wav_file.getframerate() != PCM16_SAMPLE_RATE:
|
||||
raise ValueError("input WAV must be 16 kHz")
|
||||
if wav_file.getcomptype() != "NONE":
|
||||
raise ValueError("input WAV must be uncompressed PCM")
|
||||
return wav_file.readframes(wav_file.getnframes())
|
||||
|
||||
|
||||
def write_pcm16_wav(path: Path, pcm16: bytes, *, sample_rate_hz: int) -> None:
|
||||
"""Write mono PCM16 bytes as a WAV artifact."""
|
||||
with wave.open(str(path), "wb") as wav_file:
|
||||
wav_file.setnchannels(1)
|
||||
wav_file.setsampwidth(PCM16_BYTES_PER_SAMPLE)
|
||||
wav_file.setframerate(sample_rate_hz)
|
||||
wav_file.writeframes(pcm16)
|
||||
|
||||
|
||||
async def wait_for(
|
||||
predicate: Callable[[], bool],
|
||||
*,
|
||||
timeout_s: float,
|
||||
label: str,
|
||||
) -> None:
|
||||
"""Wait for a collector predicate without coupling to a scenario runner."""
|
||||
deadline = time.monotonic() + timeout_s
|
||||
while time.monotonic() < deadline:
|
||||
if predicate():
|
||||
return
|
||||
await asyncio.sleep(0.02)
|
||||
raise TimeoutError(f"Timed out waiting for {label}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class RealtimeEventCollector:
|
||||
"""Collect server events and decode response audio by response identity."""
|
||||
|
||||
events: list[dict[str, object]] = field(default_factory=list)
|
||||
event_received_at_s: list[float] = field(default_factory=list)
|
||||
response_audio: dict[str, list[bytes]] = field(default_factory=dict)
|
||||
response_ids: list[str] = field(default_factory=list)
|
||||
output_sample_rate_hz: int = 24_000
|
||||
|
||||
@staticmethod
|
||||
def response_id(event: dict[str, object]) -> str | None:
|
||||
response_id = event.get("response_id")
|
||||
if isinstance(response_id, str):
|
||||
return response_id
|
||||
response = event.get("response")
|
||||
if isinstance(response, dict):
|
||||
response_id = response.get("id")
|
||||
if isinstance(response_id, str):
|
||||
return response_id
|
||||
return None
|
||||
|
||||
def add(self, event: dict[str, object], *, received_at_s: float | None = None) -> None:
|
||||
received_at = time.monotonic() if received_at_s is None else float(received_at_s)
|
||||
stored_event = dict(event)
|
||||
stored_event.setdefault("_client_received_at_s", received_at)
|
||||
self.events.append(stored_event)
|
||||
self.event_received_at_s.append(received_at)
|
||||
response_id = self.response_id(stored_event)
|
||||
event_type = stored_event.get("type")
|
||||
if event_type == "response.created" and response_id and response_id not in self.response_ids:
|
||||
self.response_ids.append(response_id)
|
||||
if event_type == "response.audio.delta":
|
||||
delta = stored_event.get("delta") or stored_event.get("audio")
|
||||
if isinstance(delta, str) and response_id:
|
||||
try:
|
||||
self.response_audio.setdefault(response_id, []).append(base64.b64decode(delta))
|
||||
except ValueError:
|
||||
pass
|
||||
sample_rate_hz = stored_event.get("sample_rate_hz")
|
||||
if isinstance(sample_rate_hz, int) and sample_rate_hz > 0:
|
||||
self.output_sample_rate_hz = sample_rate_hz
|
||||
|
||||
def count(self, event_type: str) -> int:
|
||||
return sum(event.get("type") == event_type for event in self.events)
|
||||
|
||||
def audio_bytes(self, response_id: str | None = None) -> bytes:
|
||||
if response_id is not None:
|
||||
return b"".join(self.response_audio.get(response_id, ()))
|
||||
return b"".join(
|
||||
chunk for response_id in self.response_ids for chunk in self.response_audio.get(response_id, ())
|
||||
)
|
||||
|
||||
def errors(self) -> list[dict[str, object]]:
|
||||
return [event for event in self.events if event.get("type") == "error"]
|
||||
|
||||
def first_received_at(
|
||||
self,
|
||||
*event_types: str,
|
||||
after_s: float = 0.0,
|
||||
) -> float | None:
|
||||
for event, received_at_s in zip(self.events, self.event_received_at_s, strict=True):
|
||||
if received_at_s >= after_s and event.get("type") in event_types:
|
||||
return received_at_s
|
||||
return None
|
||||
|
||||
def last_received_at(self, event_type: str) -> float | None:
|
||||
for event, received_at_s in zip(
|
||||
reversed(self.events),
|
||||
reversed(self.event_received_at_s),
|
||||
strict=True,
|
||||
):
|
||||
if event.get("type") == event_type:
|
||||
return received_at_s
|
||||
return None
|
||||
|
||||
def timing_summary(
|
||||
self,
|
||||
*,
|
||||
after_s: float,
|
||||
input_committed_at_s: float | None = None,
|
||||
response_id: str | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""Summarize engine token metrics and client-observed audio cadence."""
|
||||
stage0_metrics: dict[str, object] | None = None
|
||||
response_created_at_s: float | None = None
|
||||
audio_received_at_s: list[float] = []
|
||||
cumulative_audio_ms: list[float] = []
|
||||
for event, received_at_s in zip(self.events, self.event_received_at_s, strict=True):
|
||||
if received_at_s < after_s:
|
||||
continue
|
||||
event_response_id = self.response_id(event)
|
||||
if response_id is not None and event_response_id != response_id:
|
||||
continue
|
||||
if event.get("type") == "response.created" and response_created_at_s is None:
|
||||
response_created_at_s = received_at_s
|
||||
|
||||
stage_metrics = _event_stage_metrics(event)
|
||||
stage0 = stage_metrics.get("0") if isinstance(stage_metrics, dict) else None
|
||||
if isinstance(stage0, dict):
|
||||
stage0_metrics = stage0
|
||||
|
||||
if event.get("type") != "response.audio.delta" or (
|
||||
response_id is not None and event_response_id != response_id
|
||||
):
|
||||
continue
|
||||
delta = event.get("delta") or event.get("audio")
|
||||
if not isinstance(delta, str) or not delta:
|
||||
continue
|
||||
audio_received_at_s.append(received_at_s)
|
||||
metadata = event.get("metadata")
|
||||
duration_ms = metadata.get("audio_duration_ms") if isinstance(metadata, dict) else None
|
||||
if isinstance(duration_ms, int | float) and math.isfinite(float(duration_ms)):
|
||||
cumulative_audio_ms.append(max(0.0, float(duration_ms)))
|
||||
|
||||
result: dict[str, object] = {}
|
||||
if stage0_metrics is not None:
|
||||
raw_itls = stage0_metrics.get("vllm_itls_ms")
|
||||
itls = (
|
||||
[float(value) for value in raw_itls if isinstance(value, int | float)]
|
||||
if isinstance(raw_itls, list)
|
||||
else []
|
||||
)
|
||||
result["stage0_tokens"] = {
|
||||
"source": "engine_stage_metrics",
|
||||
"output_token_count": int(stage0_metrics.get("num_tokens_out") or 0),
|
||||
"ttft_ms": float(stage0_metrics.get("vllm_ttft_ms") or 0.0),
|
||||
"tpot_ms": float(stage0_metrics.get("vllm_tpot_ms") or 0.0),
|
||||
"inter_token_interval_ms": _interval_summary(itls),
|
||||
}
|
||||
|
||||
if audio_received_at_s:
|
||||
intervals_ms = [
|
||||
(current - previous) * 1000.0 for previous, current in zip(audio_received_at_s, audio_received_at_s[1:])
|
||||
]
|
||||
chunk_durations_ms: list[float] = []
|
||||
previous_duration_ms = 0.0
|
||||
for duration_ms in cumulative_audio_ms:
|
||||
chunk_durations_ms.append(
|
||||
duration_ms - previous_duration_ms if duration_ms >= previous_duration_ms else duration_ms
|
||||
)
|
||||
previous_duration_ms = duration_ms
|
||||
interval_summary = _interval_summary(intervals_ms)
|
||||
result["audio_output"] = {
|
||||
"source": "client_monotonic_receive",
|
||||
"chunk_count": len(audio_received_at_s),
|
||||
"response_created_to_first_audio_ms": (
|
||||
_rounded_ms((audio_received_at_s[0] - response_created_at_s) * 1000.0)
|
||||
if response_created_at_s is not None
|
||||
else None
|
||||
),
|
||||
"commit_to_first_audio_ms": (
|
||||
_rounded_ms((audio_received_at_s[0] - input_committed_at_s) * 1000.0)
|
||||
if input_committed_at_s is not None
|
||||
else None
|
||||
),
|
||||
"inter_chunk_interval_ms": interval_summary,
|
||||
"chunk_duration_ms": _interval_summary(chunk_durations_ms),
|
||||
"max_chunk_gap_ms": interval_summary["max"],
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
class RealtimeDuplexClient:
|
||||
"""Small async client used by the user demo and reusable smoke probes."""
|
||||
|
||||
def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None:
|
||||
self.url = url
|
||||
self.max_size = max_size
|
||||
self.events = RealtimeEventCollector()
|
||||
self._ws: Any = None
|
||||
self._reader_task: asyncio.Task[None] | None = None
|
||||
|
||||
async def __aenter__(self) -> RealtimeDuplexClient:
|
||||
self._ws = await websockets.connect(self.url, max_size=self.max_size)
|
||||
self._reader_task = asyncio.create_task(self._read_events())
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, traceback) -> None:
|
||||
if self._ws is not None:
|
||||
await self._ws.close()
|
||||
if self._reader_task is not None:
|
||||
self._reader_task.cancel()
|
||||
try:
|
||||
await self._reader_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def _read_events(self) -> None:
|
||||
try:
|
||||
while True:
|
||||
raw = await self._ws.recv()
|
||||
if not isinstance(raw, str):
|
||||
continue
|
||||
event = json.loads(raw)
|
||||
if isinstance(event, dict):
|
||||
self.events.add(event)
|
||||
except ConnectionClosed:
|
||||
return
|
||||
|
||||
async def send(self, event: dict[str, object]) -> None:
|
||||
await self._ws.send(json.dumps(event))
|
||||
|
||||
async def configure(
|
||||
self,
|
||||
model: str,
|
||||
*,
|
||||
output_audio_format: str = "pcm16",
|
||||
ref_audio: str | None = None,
|
||||
session_id: str | None = None,
|
||||
timeout_s: float = 20.0,
|
||||
) -> None:
|
||||
session: dict[str, object] = {
|
||||
"model": model,
|
||||
"modalities": ["audio", "text"],
|
||||
"input_audio_format": "pcm16",
|
||||
"output_audio_format": output_audio_format,
|
||||
"turn_detection": None,
|
||||
"overlap_policy": "listen_only",
|
||||
"playback_commit_policy": "ack_only",
|
||||
"extra_body": {
|
||||
"auto_response": True,
|
||||
"minicpmo45_native_duplex": True,
|
||||
"force_listen_count": 0,
|
||||
},
|
||||
}
|
||||
if ref_audio is not None:
|
||||
session["ref_audio"] = ref_audio
|
||||
if session_id:
|
||||
session["session_id"] = session_id
|
||||
await self.send({"type": "session.update", "session": session})
|
||||
await wait_for(
|
||||
lambda: self.events.count("session.created") > 0,
|
||||
timeout_s=timeout_s,
|
||||
label="session.created",
|
||||
)
|
||||
|
||||
async def stream_pcm16(
|
||||
self,
|
||||
pcm16: bytes,
|
||||
*,
|
||||
chunk_ms: int = 200,
|
||||
realtime: bool = True,
|
||||
) -> None:
|
||||
chunk_bytes = max(
|
||||
PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_ms // 1000,
|
||||
PCM16_BYTES_PER_SAMPLE,
|
||||
)
|
||||
audio_end_ms = 0
|
||||
for offset in range(0, len(pcm16), chunk_bytes):
|
||||
chunk = pcm16[offset : offset + chunk_bytes]
|
||||
duration_ms = len(chunk) * 1000 // (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
|
||||
audio_end_ms += duration_ms
|
||||
await self.send(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": base64.b64encode(chunk).decode("ascii"),
|
||||
"input_audio_format": "pcm16",
|
||||
"sample_rate_hz": PCM16_SAMPLE_RATE,
|
||||
"duration_ms": duration_ms,
|
||||
"audio_end_ms": audio_end_ms,
|
||||
}
|
||||
)
|
||||
if realtime:
|
||||
await asyncio.sleep(duration_ms / 1000)
|
||||
|
||||
async def commit(self) -> None:
|
||||
await self.send({"type": "input_audio_buffer.commit", "final": True})
|
||||
|
||||
async def acknowledge_playback(self) -> None:
|
||||
for response_id in self.events.response_ids:
|
||||
pcm16 = self.events.audio_bytes(response_id)
|
||||
if not pcm16:
|
||||
continue
|
||||
played_ms = len(pcm16) * 1000 // (self.events.output_sample_rate_hz * PCM16_BYTES_PER_SAMPLE)
|
||||
await self.send(
|
||||
{
|
||||
"type": "playback.ack",
|
||||
"response_id": response_id,
|
||||
"item_id": f"item_{response_id}",
|
||||
"played_ms": played_ms,
|
||||
"committed_ms": played_ms,
|
||||
}
|
||||
)
|
||||
|
||||
async def close_session(self, *, timeout_s: float = 20.0) -> None:
|
||||
await self.send({"type": "session.close"})
|
||||
await wait_for(
|
||||
lambda: self.events.count("session.closed") > 0,
|
||||
timeout_s=timeout_s,
|
||||
label="session.closed",
|
||||
)
|
||||
@@ -0,0 +1,4 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Engine-side helpers for the experimental full-duplex runtime."""
|
||||
@@ -0,0 +1,276 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Model-neutral contracts for the experimental duplex engine plugin.
|
||||
|
||||
This module contains only immutable data transfer objects and narrow protocols.
|
||||
Duplex control algorithms, session implementations, model policy, and Realtime
|
||||
serving remain in sibling experimental modules.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from collections.abc import Iterable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Protocol
|
||||
|
||||
from vllm_omni.engine.messages import EngineQueueMessage
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
|
||||
|
||||
class SessionMode(str, Enum):
|
||||
TURN = "turn"
|
||||
DUPLEX = "duplex"
|
||||
|
||||
|
||||
class DuplexInputMode(str, Enum):
|
||||
APPEND_TOKENS = "append_tokens"
|
||||
APPEND_AUDIO_CHUNK = "append_audio_chunk"
|
||||
REPLACE_LATEST_CHUNK = "replace_latest_chunk"
|
||||
REENCODE_CONTEXT = "reencode_context"
|
||||
ROLLBACK_TO_CHECKPOINT = "rollback_to_checkpoint"
|
||||
TURN_COMMIT_ONLY = "turn_commit_only"
|
||||
|
||||
|
||||
class DuplexOutputAction(str, Enum):
|
||||
DIRECT_RESPONSE = "direct_response"
|
||||
|
||||
|
||||
@dataclass
|
||||
class DuplexRuntimeCapabilities:
|
||||
input_modes: set[DuplexInputMode] = field(default_factory=lambda: {DuplexInputMode.TURN_COMMIT_ONLY})
|
||||
implementation_level: str = "serving_session_adapter"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexAppendPlan:
|
||||
prompt: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexOutputDecision:
|
||||
action: DuplexOutputAction
|
||||
metadata: Mapping[str, Any] = field(default_factory=dict)
|
||||
final_output_type: str = "text"
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata)))
|
||||
|
||||
|
||||
class DuplexRuntimeExtension(Protocol):
|
||||
"""Pure model policy invoked by the experimental duplex control plane."""
|
||||
|
||||
def configure_sampling_params(
|
||||
self,
|
||||
*,
|
||||
runtime_config: dict[str, Any],
|
||||
defaults: tuple[object, ...],
|
||||
) -> tuple[object, ...]: ...
|
||||
|
||||
def plan_append(
|
||||
self,
|
||||
*,
|
||||
request_id: str,
|
||||
fence: DuplexFence,
|
||||
session_config: dict[str, Any],
|
||||
runtime_config: dict[str, Any],
|
||||
seq: int,
|
||||
turn_seq: int,
|
||||
mode: DuplexInputMode,
|
||||
payload: object,
|
||||
final: bool,
|
||||
sampling_params: object,
|
||||
) -> DuplexAppendPlan: ...
|
||||
|
||||
def decide_output(
|
||||
self,
|
||||
*,
|
||||
stage_id: int,
|
||||
final_stage_id: int,
|
||||
segment_finished: bool,
|
||||
segment_token_ids: tuple[int, ...],
|
||||
segment_output_metadata: dict[str, Any],
|
||||
output: object,
|
||||
) -> DuplexOutputDecision | None: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexRequestIdentity:
|
||||
session_id: str
|
||||
fence: DuplexFence
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexStageRequestContext:
|
||||
request_id: str
|
||||
session_id: str
|
||||
fence: DuplexFence
|
||||
stage_id: int
|
||||
final_stage_id: int
|
||||
config_generation: int
|
||||
sampling_params: tuple[object, ...]
|
||||
session_config: Mapping[str, Any] = field(default_factory=dict)
|
||||
runtime_config: Mapping[str, Any] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "sampling_params", tuple(self.sampling_params))
|
||||
object.__setattr__(self, "session_config", MappingProxyType(dict(self.session_config)))
|
||||
object.__setattr__(self, "runtime_config", MappingProxyType(dict(self.runtime_config)))
|
||||
|
||||
@property
|
||||
def stage_sampling_params(self) -> object:
|
||||
return self.sampling_params[self.stage_id]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexStageSubmission:
|
||||
context: DuplexStageRequestContext
|
||||
prompt: Mapping[str, Any]
|
||||
already_submitted: bool
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "prompt", MappingProxyType(dict(self.prompt)))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexStageSubmissionResult:
|
||||
request_id: str
|
||||
stage_id: int
|
||||
replica_id: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexOutputContext:
|
||||
identity: DuplexRequestIdentity
|
||||
final_stage_id: int
|
||||
segment_finished: bool
|
||||
segment_token_ids: tuple[int, ...] = ()
|
||||
segment_output_metadata: Mapping[str, Any] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "segment_token_ids", tuple(self.segment_token_ids))
|
||||
object.__setattr__(
|
||||
self,
|
||||
"segment_output_metadata",
|
||||
MappingProxyType(dict(self.segment_output_metadata)),
|
||||
)
|
||||
|
||||
|
||||
class DuplexStagePort(Protocol):
|
||||
@property
|
||||
def stage_count(self) -> int: ...
|
||||
|
||||
def sampling_defaults(self) -> tuple[object, ...]: ...
|
||||
|
||||
def ensure_request(self, context: DuplexStageRequestContext) -> None: ...
|
||||
|
||||
async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: ...
|
||||
|
||||
async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: ...
|
||||
|
||||
|
||||
class DuplexControlPlanePort(Protocol):
|
||||
@property
|
||||
def sessions(self) -> object: ...
|
||||
|
||||
def accepts(self, message: object) -> bool: ...
|
||||
|
||||
def dispatch(self, message: object) -> None: ...
|
||||
|
||||
async def shutdown(self) -> None: ...
|
||||
|
||||
def close_sessions_for_request_ids(self, request_ids: list[str]) -> dict[str, list[str]]: ...
|
||||
|
||||
def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: ...
|
||||
|
||||
def session_for_identity(self, identity: DuplexRequestIdentity | None) -> object | None: ...
|
||||
|
||||
def decide_output(
|
||||
self,
|
||||
stage_id: int,
|
||||
output: object,
|
||||
context: DuplexOutputContext | None,
|
||||
) -> DuplexOutputDecision | None: ...
|
||||
|
||||
|
||||
class CorrelatedRpcTransport(Protocol):
|
||||
def execute(
|
||||
self,
|
||||
key: tuple[str, str],
|
||||
message: EngineQueueMessage,
|
||||
*,
|
||||
timeout: float | None,
|
||||
timeout_message: str,
|
||||
block_on_submit: bool = False,
|
||||
) -> EngineQueueMessage: ...
|
||||
|
||||
|
||||
def duplex_data_plane_request_info(result: dict[str, object]) -> tuple[str | None, int | None]:
|
||||
stage_results = result.get("stage_results")
|
||||
if not isinstance(stage_results, list):
|
||||
return None, None
|
||||
for item in stage_results:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
inner = item.get("result")
|
||||
if not isinstance(inner, dict) or inner.get("data_plane_append") is not True:
|
||||
continue
|
||||
request_id = inner.get("request_id")
|
||||
if not isinstance(request_id, str) or not request_id:
|
||||
continue
|
||||
response_stage_id = inner.get("response_stage_id")
|
||||
return request_id, response_stage_id if isinstance(response_stage_id, int) else None
|
||||
return None, None
|
||||
|
||||
|
||||
def duplex_resource_request_id(fence: DuplexFence, role: str) -> str:
|
||||
if not role or any(
|
||||
character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for character in role
|
||||
):
|
||||
raise ValueError(f"invalid duplex resource role: {role!r}")
|
||||
encoded_session_id = base64.urlsafe_b64encode(fence.session_id.encode("utf-8")).decode("ascii").rstrip("=")
|
||||
return f"duplex-s.{encoded_session_id}.i.{fence.incarnation}.e.{fence.epoch}.r.{role}"
|
||||
|
||||
|
||||
def duplex_resource_request_belongs_to_session(request_id: str, session_id: str) -> bool:
|
||||
"""Return whether a current-format resource request belongs to a session."""
|
||||
parts = request_id.split(".")
|
||||
if len(parts) != 8 or parts[0] != "duplex-s" or parts[2] != "i" or parts[4] != "e" or parts[6] != "r":
|
||||
return False
|
||||
try:
|
||||
int(parts[3])
|
||||
int(parts[5])
|
||||
except ValueError:
|
||||
return False
|
||||
role = parts[7]
|
||||
if not role or any(
|
||||
character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for character in role
|
||||
):
|
||||
return False
|
||||
encoded_session_id = base64.urlsafe_b64encode(session_id.encode("utf-8")).decode("ascii").rstrip("=")
|
||||
return parts[1] == encoded_session_id
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CorrelatedRpcTransport",
|
||||
"DuplexAppendPlan",
|
||||
"DuplexControlPlanePort",
|
||||
"DuplexInputMode",
|
||||
"DuplexOutputAction",
|
||||
"DuplexOutputContext",
|
||||
"DuplexOutputDecision",
|
||||
"DuplexRequestIdentity",
|
||||
"DuplexRuntimeCapabilities",
|
||||
"DuplexRuntimeExtension",
|
||||
"DuplexStagePort",
|
||||
"DuplexStageRequestContext",
|
||||
"DuplexStageSubmission",
|
||||
"DuplexStageSubmissionResult",
|
||||
"SessionMode",
|
||||
"duplex_data_plane_request_info",
|
||||
"duplex_resource_request_belongs_to_session",
|
||||
"duplex_resource_request_id",
|
||||
]
|
||||
@@ -0,0 +1,228 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import CorrelatedRpcTransport
|
||||
from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import (
|
||||
AppendDuplexInputMessage,
|
||||
CloseDuplexSessionMessage,
|
||||
DuplexControlResultMessage,
|
||||
DuplexFence,
|
||||
OpenDuplexSessionMessage,
|
||||
ResumeDuplexSessionMessage,
|
||||
SignalDuplexTurnMessage,
|
||||
TouchDuplexSessionMessage,
|
||||
)
|
||||
|
||||
DuplexControlMessage = (
|
||||
OpenDuplexSessionMessage
|
||||
| AppendDuplexInputMessage
|
||||
| SignalDuplexTurnMessage
|
||||
| CloseDuplexSessionMessage
|
||||
| TouchDuplexSessionMessage
|
||||
| ResumeDuplexSessionMessage
|
||||
)
|
||||
|
||||
|
||||
class DuplexControlRequestError(RuntimeError):
|
||||
"""Typed client-side failure returned by the duplex control plane."""
|
||||
|
||||
def __init__(self, result: dict[str, object]) -> None:
|
||||
error = result.get("error")
|
||||
error_data = error if isinstance(error, dict) else {}
|
||||
self.result = result
|
||||
self.code = str(error_data.get("code") or "internal_error")
|
||||
self.retryable = bool(error_data.get("retryable", False))
|
||||
accepted_fence = result.get("accepted_fence")
|
||||
self.accepted_fence = accepted_fence if isinstance(accepted_fence, DuplexFence) else None
|
||||
lease_generation = result.get("lease_generation")
|
||||
self.lease_generation = lease_generation if isinstance(lease_generation, int) else None
|
||||
super().__init__(f"duplex {result.get('operation')} failed: {result}")
|
||||
|
||||
|
||||
class DuplexControlClient:
|
||||
"""Builds duplex commands on top of an engine-owned RPC transport."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
transport: CorrelatedRpcTransport,
|
||||
*,
|
||||
control_id_factory: Callable[[], str] | None = None,
|
||||
) -> None:
|
||||
self._transport = transport
|
||||
self._control_id_factory = control_id_factory or (lambda: uuid.uuid4().hex)
|
||||
|
||||
def execute(self, message: DuplexControlMessage, *, timeout: float | None) -> dict[str, object]:
|
||||
result_message = self._transport.execute(
|
||||
("duplex", message.control_id),
|
||||
message,
|
||||
timeout=timeout,
|
||||
timeout_message=f"duplex control timed out after {timeout} seconds",
|
||||
)
|
||||
if not isinstance(result_message, DuplexControlResultMessage):
|
||||
raise RuntimeError(f"unexpected duplex control result type: {type(result_message).__name__}")
|
||||
if result_message.fence != message.fence:
|
||||
raise RuntimeError(
|
||||
f"duplex control fence mismatch: expected {message.fence!r}, got {result_message.fence!r}"
|
||||
)
|
||||
result = {
|
||||
"fence": result_message.fence,
|
||||
"operation": result_message.operation,
|
||||
"session_id": result_message.session_id,
|
||||
"ok": result_message.ok,
|
||||
"stage_results": list(result_message.stage_results),
|
||||
"unsupported_count": result_message.unsupported_count,
|
||||
"error_count": result_message.error_count,
|
||||
"accepted_fence": result_message.accepted_fence,
|
||||
"lease_generation": result_message.lease_generation,
|
||||
"error": (
|
||||
{
|
||||
"code": result_message.error.code,
|
||||
"message": result_message.error.message,
|
||||
"retryable": result_message.error.retryable,
|
||||
}
|
||||
if result_message.error is not None
|
||||
else None
|
||||
),
|
||||
}
|
||||
if not result_message.ok:
|
||||
raise DuplexControlRequestError(result)
|
||||
return result
|
||||
|
||||
def open(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
session_mode: str,
|
||||
capabilities: dict[str, object] | None,
|
||||
session_config: dict[str, object] | None,
|
||||
runtime_config: dict[str, object] | None,
|
||||
fence: DuplexFence,
|
||||
timeout: float | None,
|
||||
) -> dict[str, object]:
|
||||
return self.execute(
|
||||
OpenDuplexSessionMessage(
|
||||
control_id=self._control_id_factory(),
|
||||
fence=fence,
|
||||
session_id=session_id,
|
||||
session_mode=session_mode,
|
||||
capabilities=dict(capabilities or {}),
|
||||
session_config=dict(session_config or {}),
|
||||
runtime_config=dict(runtime_config or {}),
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def append(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
mode: str,
|
||||
payload: object,
|
||||
operation_id: str | None,
|
||||
final: bool,
|
||||
expected_epoch: int | None,
|
||||
fence: DuplexFence,
|
||||
timeout: float | None,
|
||||
) -> dict[str, object]:
|
||||
if expected_epoch is not None and expected_epoch != fence.epoch:
|
||||
raise ValueError("expected_epoch must match fence.epoch")
|
||||
return self.execute(
|
||||
AppendDuplexInputMessage(
|
||||
control_id=self._control_id_factory(),
|
||||
operation_id=operation_id,
|
||||
fence=fence,
|
||||
session_id=session_id,
|
||||
expected_epoch=expected_epoch,
|
||||
mode=mode,
|
||||
payload=payload,
|
||||
final=final,
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def signal(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
event: str,
|
||||
fence: DuplexFence,
|
||||
next_fence: DuplexFence | None,
|
||||
session_config: dict[str, object] | None,
|
||||
runtime_config: dict[str, object] | None,
|
||||
timeout: float | None,
|
||||
) -> dict[str, object]:
|
||||
return self.execute(
|
||||
SignalDuplexTurnMessage(
|
||||
control_id=self._control_id_factory(),
|
||||
fence=fence,
|
||||
session_id=session_id,
|
||||
event=event,
|
||||
next_fence=next_fence,
|
||||
session_config=dict(session_config) if session_config is not None else None,
|
||||
runtime_config=dict(runtime_config) if runtime_config is not None else None,
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def close(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
reason: str,
|
||||
fence: DuplexFence,
|
||||
timeout: float | None,
|
||||
) -> dict[str, object]:
|
||||
return self.execute(
|
||||
CloseDuplexSessionMessage(
|
||||
control_id=self._control_id_factory(),
|
||||
fence=fence,
|
||||
session_id=session_id,
|
||||
reason=reason,
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def touch(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
activity: DuplexLeaseActivity,
|
||||
timeout: float | None,
|
||||
) -> dict[str, object]:
|
||||
return self.execute(
|
||||
TouchDuplexSessionMessage(
|
||||
control_id=self._control_id_factory(),
|
||||
fence=fence,
|
||||
session_id=session_id,
|
||||
activity=activity.value,
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def resume(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
expected_lease_generation: int,
|
||||
timeout: float | None,
|
||||
) -> dict[str, object]:
|
||||
return self.execute(
|
||||
ResumeDuplexSessionMessage(
|
||||
control_id=self._control_id_factory(),
|
||||
fence=fence,
|
||||
session_id=session_id,
|
||||
expected_lease_generation=expected_lease_generation,
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["CorrelatedRpcTransport", "DuplexControlClient", "DuplexControlMessage"]
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,108 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Experimental duplex runtime loading and compatibility exports."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from importlib import import_module
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import (
|
||||
DuplexAppendPlan,
|
||||
DuplexInputMode,
|
||||
DuplexOutputAction,
|
||||
DuplexOutputDecision,
|
||||
DuplexRuntimeCapabilities,
|
||||
DuplexRuntimeExtension,
|
||||
SessionMode,
|
||||
duplex_data_plane_request_info,
|
||||
duplex_resource_request_belongs_to_session,
|
||||
duplex_resource_request_id,
|
||||
)
|
||||
|
||||
|
||||
def load_duplex_runtime_extension(path: str | None) -> DuplexRuntimeExtension | None:
|
||||
if not path:
|
||||
return None
|
||||
module_name, separator, attribute_name = path.rpartition(".")
|
||||
if not separator:
|
||||
raise ValueError(f"Invalid duplex runtime extension path: {path!r}")
|
||||
extension_type = getattr(import_module(module_name), attribute_name)
|
||||
return validate_duplex_runtime_extension(extension_type())
|
||||
|
||||
|
||||
def validate_duplex_runtime_extension(
|
||||
extension: object,
|
||||
*,
|
||||
sampling_defaults: tuple[object, ...] | None = None,
|
||||
) -> DuplexRuntimeExtension:
|
||||
required_methods = (
|
||||
"configure_sampling_params",
|
||||
"plan_append",
|
||||
"decide_output",
|
||||
)
|
||||
missing = [name for name in required_methods if not callable(getattr(extension, name, None))]
|
||||
if missing:
|
||||
raise TypeError(f"Duplex runtime extension is missing callable method(s): {', '.join(missing)}")
|
||||
typed_extension = extension # type: ignore[assignment]
|
||||
if sampling_defaults is not None:
|
||||
configured = typed_extension.configure_sampling_params(
|
||||
runtime_config={},
|
||||
defaults=sampling_defaults,
|
||||
)
|
||||
if not isinstance(configured, tuple):
|
||||
raise TypeError("Duplex runtime extension must return sampling parameters as a tuple")
|
||||
if len(configured) != len(sampling_defaults):
|
||||
raise ValueError("Duplex runtime extension must return one sampling parameter per stage")
|
||||
for stage_id, (value, default) in enumerate(zip(configured, sampling_defaults, strict=True)):
|
||||
if default is not None and not isinstance(value, type(default)):
|
||||
raise TypeError(
|
||||
"Duplex runtime extension sampling parameter type mismatch "
|
||||
f"for stage {stage_id}: expected {type(default).__name__}, got {type(value).__name__}"
|
||||
)
|
||||
return typed_extension
|
||||
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_session import ( # noqa: E402, F401
|
||||
DuplexAppendReservation,
|
||||
DuplexCompletedAppend,
|
||||
DuplexFenceMismatchError,
|
||||
DuplexInputAppend,
|
||||
DuplexRequestResource,
|
||||
DuplexSessionRuntimeManager,
|
||||
DuplexSessionRuntimeState,
|
||||
DuplexStageBinding,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.lease import ( # noqa: E402, F401
|
||||
DuplexLeaseActivity,
|
||||
DuplexLeaseConfig,
|
||||
DuplexLeaseState,
|
||||
DuplexSessionExpiry,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DuplexAppendPlan",
|
||||
"DuplexAppendReservation",
|
||||
"DuplexCompletedAppend",
|
||||
"DuplexFenceMismatchError",
|
||||
"DuplexInputAppend",
|
||||
"DuplexInputMode",
|
||||
"DuplexLeaseActivity",
|
||||
"DuplexLeaseConfig",
|
||||
"DuplexLeaseState",
|
||||
"DuplexOutputAction",
|
||||
"DuplexOutputDecision",
|
||||
"DuplexRequestResource",
|
||||
"DuplexRuntimeCapabilities",
|
||||
"DuplexRuntimeExtension",
|
||||
"DuplexSessionExpiry",
|
||||
"DuplexSessionRuntimeManager",
|
||||
"DuplexSessionRuntimeState",
|
||||
"DuplexStageBinding",
|
||||
"SessionMode",
|
||||
"duplex_data_plane_request_info",
|
||||
"duplex_resource_request_belongs_to_session",
|
||||
"duplex_resource_request_id",
|
||||
"load_duplex_runtime_extension",
|
||||
"validate_duplex_runtime_extension",
|
||||
]
|
||||
@@ -0,0 +1,534 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import (
|
||||
DuplexInputMode,
|
||||
DuplexRuntimeCapabilities,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.lease import (
|
||||
DuplexLeaseActivity,
|
||||
DuplexLeaseConfig,
|
||||
DuplexLeaseState,
|
||||
DuplexSessionExpiry,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
|
||||
|
||||
def _default_capabilities() -> DuplexRuntimeCapabilities:
|
||||
return DuplexRuntimeCapabilities()
|
||||
|
||||
|
||||
class DuplexFenceMismatchError(RuntimeError):
|
||||
def __init__(self, expected: DuplexFence, actual: DuplexFence) -> None:
|
||||
super().__init__(f"duplex fence mismatch: expected {expected!r}, got {actual!r}")
|
||||
self.expected = expected
|
||||
self.actual = actual
|
||||
|
||||
|
||||
@dataclass
|
||||
class DuplexStageBinding:
|
||||
request_id: str
|
||||
fence: DuplexFence
|
||||
|
||||
|
||||
@dataclass
|
||||
class DuplexRequestResource:
|
||||
stage_id: int
|
||||
request_id: str
|
||||
fence: DuplexFence
|
||||
submitted: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class DuplexInputAppend:
|
||||
seq: int
|
||||
turn_seq: int
|
||||
turn_id: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexAppendReservation:
|
||||
fence: DuplexFence
|
||||
mode: DuplexInputMode
|
||||
base_fence: DuplexFence
|
||||
base_input_seq: int
|
||||
base_input_turn_seq: int
|
||||
base_append_turn_key: tuple[int, int, int] | None
|
||||
update: DuplexInputAppend
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexCompletedAppend:
|
||||
fence: DuplexFence
|
||||
mode: DuplexInputMode
|
||||
final: bool
|
||||
stage_results: tuple[dict[str, object], ...]
|
||||
|
||||
|
||||
@dataclass
|
||||
class DuplexSessionRuntimeState:
|
||||
"""Engine resource handles associated with a core-owned identity fence."""
|
||||
|
||||
fence: DuplexFence
|
||||
lease: DuplexLeaseState
|
||||
_clock: Callable[[], float] = field(repr=False)
|
||||
capabilities: DuplexRuntimeCapabilities = field(default_factory=_default_capabilities)
|
||||
session_config: dict[str, Any] = field(default_factory=dict)
|
||||
runtime_config: dict[str, Any] = field(default_factory=dict)
|
||||
config_generation: int = 0
|
||||
stage_bindings: dict[int, DuplexStageBinding] = field(default_factory=dict)
|
||||
request_resources: dict[tuple[int, str], DuplexRequestResource] = field(default_factory=dict)
|
||||
input_seq: int = 0
|
||||
input_turn_seq: int = 0
|
||||
_append_turn_key: tuple[int, int, int] | None = None
|
||||
completed_append_limit: int = 256
|
||||
completed_appends: OrderedDict[str, DuplexCompletedAppend] = field(default_factory=OrderedDict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.completed_append_limit <= 0:
|
||||
raise ValueError("completed_append_limit must be positive")
|
||||
|
||||
@property
|
||||
def session_id(self) -> str:
|
||||
return self.fence.session_id
|
||||
|
||||
@property
|
||||
def epoch(self) -> int:
|
||||
return self.fence.epoch
|
||||
|
||||
@property
|
||||
def turn_id(self) -> int:
|
||||
return self.fence.turn_id
|
||||
|
||||
def _validate_fence(self, fence: DuplexFence) -> None:
|
||||
if fence.session_id != self.session_id or fence.incarnation != self.fence.incarnation:
|
||||
raise DuplexFenceMismatchError(self.fence, fence)
|
||||
current = self.fence
|
||||
if fence.epoch < current.epoch or (
|
||||
fence.epoch == current.epoch
|
||||
and (fence.turn_id < current.turn_id or fence.response_seq < current.response_seq)
|
||||
):
|
||||
raise DuplexFenceMismatchError(current, fence)
|
||||
|
||||
def accept_fence(self, fence: DuplexFence) -> None:
|
||||
self._validate_fence(fence)
|
||||
if fence.epoch != self.fence.epoch:
|
||||
self.input_seq = 0
|
||||
self.input_turn_seq = 0
|
||||
self._append_turn_key = None
|
||||
self.completed_appends.clear()
|
||||
self.fence = fence
|
||||
|
||||
def touch(self, fence: DuplexFence, activity: DuplexLeaseActivity) -> None:
|
||||
self._validate_fence(fence)
|
||||
self.lease.touch(self._clock(), activity)
|
||||
|
||||
def detach(self, fence: DuplexFence) -> None:
|
||||
self._validate_fence(fence)
|
||||
self.lease.detach(self._clock())
|
||||
|
||||
def resume(self, fence: DuplexFence, *, expected_lease_generation: int) -> int:
|
||||
self._validate_fence(fence)
|
||||
return self.lease.resume(
|
||||
self._clock(),
|
||||
expected_generation=expected_lease_generation,
|
||||
)
|
||||
|
||||
def begin_operation(self, fence: DuplexFence, operation_id: str) -> None:
|
||||
self._validate_fence(fence)
|
||||
self.lease.begin_operation(self._clock(), operation_id)
|
||||
|
||||
def end_operation(self, fence: DuplexFence, operation_id: str) -> None:
|
||||
self._validate_fence(fence)
|
||||
self.lease.end_operation(self._clock(), operation_id)
|
||||
|
||||
def replace_session_config(self, session_config: dict[str, Any]) -> None:
|
||||
self.session_config = dict(session_config)
|
||||
self.config_generation += 1
|
||||
|
||||
def replace_runtime_config(self, runtime_config: dict[str, Any]) -> None:
|
||||
self.runtime_config = dict(runtime_config)
|
||||
self.config_generation += 1
|
||||
|
||||
def replace_configs(
|
||||
self,
|
||||
*,
|
||||
session_config: dict[str, Any] | None = None,
|
||||
runtime_config: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Atomically publish one validated configuration generation."""
|
||||
if session_config is None and runtime_config is None:
|
||||
return
|
||||
next_session_config = self.session_config if session_config is None else dict(session_config)
|
||||
next_runtime_config = self.runtime_config if runtime_config is None else dict(runtime_config)
|
||||
self.session_config = next_session_config
|
||||
self.runtime_config = next_runtime_config
|
||||
self.config_generation += 1
|
||||
|
||||
def reserve_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None:
|
||||
self._validate_fence(fence)
|
||||
resource_key = (stage_id, request_id)
|
||||
existing = self.request_resources.get(resource_key)
|
||||
if existing is not None:
|
||||
if (
|
||||
existing.fence.session_id != fence.session_id
|
||||
or existing.fence.incarnation != fence.incarnation
|
||||
or existing.fence.epoch != fence.epoch
|
||||
):
|
||||
raise ValueError(f"Duplex request resource already reserved with different identity: {request_id}")
|
||||
return
|
||||
self.request_resources[resource_key] = DuplexRequestResource(
|
||||
stage_id=stage_id,
|
||||
request_id=request_id,
|
||||
fence=fence,
|
||||
)
|
||||
|
||||
def bind_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None:
|
||||
self.reserve_stage_request(stage_id, request_id, fence=fence)
|
||||
self.accept_fence(fence)
|
||||
resource = self.request_resources[(stage_id, request_id)]
|
||||
resource.fence = fence
|
||||
resource.submitted = True
|
||||
self.stage_bindings[stage_id] = DuplexStageBinding(request_id=request_id, fence=fence)
|
||||
|
||||
def stage_request_ids(self, fence: DuplexFence | None = None) -> list[str]:
|
||||
return [
|
||||
binding.request_id for binding in self.stage_bindings.values() if fence is None or binding.fence == fence
|
||||
]
|
||||
|
||||
def resource_request_ids(
|
||||
self,
|
||||
fence: DuplexFence | None = None,
|
||||
*,
|
||||
submitted: bool | None = None,
|
||||
) -> list[str]:
|
||||
return list(
|
||||
dict.fromkeys(
|
||||
resource.request_id
|
||||
for resource in self.request_resources.values()
|
||||
if (fence is None or resource.fence == fence) and (submitted is None or resource.submitted is submitted)
|
||||
)
|
||||
)
|
||||
|
||||
def release_request_ids(self, request_ids: list[str]) -> None:
|
||||
released = set(request_ids)
|
||||
if not released:
|
||||
return
|
||||
self.stage_bindings = {
|
||||
stage_id: binding for stage_id, binding in self.stage_bindings.items() if binding.request_id not in released
|
||||
}
|
||||
self.request_resources = {
|
||||
resource_key: resource
|
||||
for resource_key, resource in self.request_resources.items()
|
||||
if resource.request_id not in released
|
||||
}
|
||||
|
||||
def prepare_append(self, *, mode: DuplexInputMode, fence: DuplexFence) -> DuplexAppendReservation:
|
||||
if mode not in self.capabilities.input_modes:
|
||||
raise ValueError(f"Duplex input mode {mode.value!r} is not supported by session {self.session_id}")
|
||||
self._validate_fence(fence)
|
||||
input_seq = 0 if fence.epoch != self.fence.epoch else self.input_seq
|
||||
input_turn_seq = 0 if fence.epoch != self.fence.epoch else self.input_turn_seq
|
||||
append_turn_key = None if fence.epoch != self.fence.epoch else self._append_turn_key
|
||||
turn_key = (fence.epoch, fence.turn_id, fence.response_seq)
|
||||
turn_seq = input_turn_seq + 1 if turn_key == append_turn_key else 1
|
||||
return DuplexAppendReservation(
|
||||
fence=fence,
|
||||
mode=mode,
|
||||
base_fence=self.fence,
|
||||
base_input_seq=self.input_seq,
|
||||
base_input_turn_seq=self.input_turn_seq,
|
||||
base_append_turn_key=self._append_turn_key,
|
||||
update=DuplexInputAppend(
|
||||
seq=input_seq + 1,
|
||||
turn_seq=turn_seq,
|
||||
turn_id=fence.turn_id,
|
||||
),
|
||||
)
|
||||
|
||||
def commit_append(self, reservation: DuplexAppendReservation) -> DuplexInputAppend:
|
||||
if (
|
||||
self.fence != reservation.base_fence
|
||||
or self.input_seq != reservation.base_input_seq
|
||||
or self.input_turn_seq != reservation.base_input_turn_seq
|
||||
or self._append_turn_key != reservation.base_append_turn_key
|
||||
):
|
||||
raise RuntimeError("duplex append reservation is stale")
|
||||
self.accept_fence(reservation.fence)
|
||||
self.input_seq = reservation.update.seq
|
||||
self.input_turn_seq = reservation.update.turn_seq
|
||||
self._append_turn_key = (
|
||||
reservation.fence.epoch,
|
||||
reservation.fence.turn_id,
|
||||
reservation.fence.response_seq,
|
||||
)
|
||||
return reservation.update
|
||||
|
||||
def completed_append(
|
||||
self,
|
||||
operation_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
mode: DuplexInputMode,
|
||||
final: bool,
|
||||
) -> list[dict[str, object]] | None:
|
||||
self._validate_fence(fence)
|
||||
completed = self.completed_appends.get(operation_id)
|
||||
if completed is None:
|
||||
return None
|
||||
if completed.fence != fence or completed.mode != mode or completed.final != final:
|
||||
raise ValueError(f"duplex append operation {operation_id!r} was reused with different metadata")
|
||||
return [dict(result) for result in completed.stage_results]
|
||||
|
||||
def record_completed_append(
|
||||
self,
|
||||
operation_id: str,
|
||||
*,
|
||||
fence: DuplexFence,
|
||||
mode: DuplexInputMode,
|
||||
final: bool,
|
||||
stage_results: list[dict[str, object]],
|
||||
) -> None:
|
||||
self.completed_appends[operation_id] = DuplexCompletedAppend(
|
||||
fence=fence,
|
||||
mode=mode,
|
||||
final=final,
|
||||
stage_results=tuple(dict(result) for result in stage_results),
|
||||
)
|
||||
self.completed_appends.move_to_end(operation_id)
|
||||
while len(self.completed_appends) > self.completed_append_limit:
|
||||
self.completed_appends.popitem(last=False)
|
||||
|
||||
def append_input(self, *, mode: DuplexInputMode, fence: DuplexFence) -> DuplexInputAppend:
|
||||
return self.commit_append(self.prepare_append(mode=mode, fence=fence))
|
||||
|
||||
def release_fence(self, fence: DuplexFence) -> list[str]:
|
||||
stale = self.resource_request_ids(fence)
|
||||
self.stage_bindings = {
|
||||
stage_id: binding for stage_id, binding in self.stage_bindings.items() if binding.fence != fence
|
||||
}
|
||||
self.request_resources = {
|
||||
resource_key: resource
|
||||
for resource_key, resource in self.request_resources.items()
|
||||
if resource.fence != fence
|
||||
}
|
||||
return stale
|
||||
|
||||
def cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]:
|
||||
stale = self.prepare_cancel_fence(cancelled_fence, next_fence)
|
||||
self.release_fence(cancelled_fence)
|
||||
return stale
|
||||
|
||||
def prepare_cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]:
|
||||
"""Advance the cancellation fence without dropping cleanup records."""
|
||||
if cancelled_fence.session_id != self.session_id or cancelled_fence.incarnation != self.fence.incarnation:
|
||||
raise DuplexFenceMismatchError(self.fence, cancelled_fence)
|
||||
if (
|
||||
next_fence.session_id != self.session_id
|
||||
or next_fence.incarnation != self.fence.incarnation
|
||||
or next_fence.epoch <= cancelled_fence.epoch
|
||||
):
|
||||
raise DuplexFenceMismatchError(cancelled_fence, next_fence)
|
||||
current_key = (self.fence.epoch, self.fence.turn_id, self.fence.response_seq)
|
||||
cancelled_key = (cancelled_fence.epoch, cancelled_fence.turn_id, cancelled_fence.response_seq)
|
||||
next_key = (next_fence.epoch, next_fence.turn_id, next_fence.response_seq)
|
||||
if cancelled_key > current_key:
|
||||
raise DuplexFenceMismatchError(self.fence, cancelled_fence)
|
||||
if next_key > current_key:
|
||||
self.accept_fence(next_fence)
|
||||
return self.resource_request_ids(cancelled_fence)
|
||||
|
||||
def begin_close(self, fence: DuplexFence, *, reason: str) -> bool:
|
||||
"""Make close irreversible while retaining resources for cleanup retry."""
|
||||
self.accept_fence(fence)
|
||||
if self.lease.terminal_reason is not None:
|
||||
return True
|
||||
return self.lease.mark_terminal(reason)
|
||||
|
||||
def finalize_close(self) -> list[str]:
|
||||
return self.close()
|
||||
|
||||
def close(self, fence: DuplexFence | None = None) -> list[str]:
|
||||
if fence is not None:
|
||||
self.accept_fence(fence)
|
||||
stale = self.resource_request_ids()
|
||||
self.stage_bindings.clear()
|
||||
self.request_resources.clear()
|
||||
return stale
|
||||
|
||||
def terminate(self, fence: DuplexFence, *, reason: str) -> bool:
|
||||
if not self.begin_close(fence, reason=reason):
|
||||
return False
|
||||
self.finalize_close()
|
||||
return True
|
||||
|
||||
|
||||
class DuplexSessionRuntimeManager:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
clock: Callable[[], float] | None = None,
|
||||
max_sessions: int | None = None,
|
||||
completed_append_limit: int = 256,
|
||||
) -> None:
|
||||
if max_sessions is not None and max_sessions <= 0:
|
||||
raise ValueError("max_sessions must be positive or null")
|
||||
if completed_append_limit <= 0:
|
||||
raise ValueError("completed_append_limit must be positive")
|
||||
self._sessions: dict[str, DuplexSessionRuntimeState] = {}
|
||||
self._clock = clock or time.monotonic
|
||||
self._max_sessions = max_sessions
|
||||
self._completed_append_limit = completed_append_limit
|
||||
|
||||
def open_session(
|
||||
self,
|
||||
fence: DuplexFence,
|
||||
*,
|
||||
capabilities: DuplexRuntimeCapabilities | None = None,
|
||||
session_config: dict[str, Any] | None = None,
|
||||
runtime_config: dict[str, Any] | None = None,
|
||||
lease_config: DuplexLeaseConfig | None = None,
|
||||
) -> DuplexSessionRuntimeState:
|
||||
if not isinstance(fence, DuplexFence):
|
||||
raise TypeError("open_session requires DuplexFence")
|
||||
if fence.session_id in self._sessions:
|
||||
raise ValueError(f"Duplex session already exists: {fence.session_id}")
|
||||
if self._max_sessions is not None and len(self._sessions) >= self._max_sessions:
|
||||
raise RuntimeError(f"duplex_session_capacity_exhausted: limit={self._max_sessions}")
|
||||
session = DuplexSessionRuntimeState(
|
||||
fence=fence,
|
||||
lease=DuplexLeaseState(
|
||||
config=lease_config or DuplexLeaseConfig(),
|
||||
generation=0,
|
||||
last_activity=self._clock(),
|
||||
),
|
||||
_clock=self._clock,
|
||||
capabilities=capabilities or _default_capabilities(),
|
||||
session_config=dict(session_config or {}),
|
||||
runtime_config=dict(runtime_config or {}),
|
||||
completed_append_limit=self._completed_append_limit,
|
||||
)
|
||||
self._sessions[fence.session_id] = session
|
||||
return session
|
||||
|
||||
def get(self, session_id: str) -> DuplexSessionRuntimeState | None:
|
||||
return self._sessions.get(session_id)
|
||||
|
||||
def require(self, session_id: str) -> DuplexSessionRuntimeState:
|
||||
session = self.get(session_id)
|
||||
if session is None:
|
||||
raise KeyError(f"Unknown duplex session: {session_id}")
|
||||
return session
|
||||
|
||||
def close_session(
|
||||
self,
|
||||
fence: DuplexFence,
|
||||
*,
|
||||
reason: str = "explicit_close",
|
||||
) -> DuplexSessionRuntimeState | None:
|
||||
if not isinstance(fence, DuplexFence):
|
||||
raise TypeError("close_session requires DuplexFence")
|
||||
session = self._sessions.get(fence.session_id)
|
||||
if session is not None:
|
||||
if session.lease.terminal_reason is not None:
|
||||
return None
|
||||
if not session.terminate(fence, reason=reason):
|
||||
return None
|
||||
if self._sessions.get(fence.session_id) is session:
|
||||
self._sessions.pop(fence.session_id)
|
||||
return session
|
||||
|
||||
def begin_close_session(
|
||||
self,
|
||||
fence: DuplexFence,
|
||||
*,
|
||||
reason: str = "explicit_close",
|
||||
) -> DuplexSessionRuntimeState | None:
|
||||
if not isinstance(fence, DuplexFence):
|
||||
raise TypeError("begin_close_session requires DuplexFence")
|
||||
session = self._sessions.get(fence.session_id)
|
||||
if session is not None:
|
||||
session.begin_close(fence, reason=reason)
|
||||
return session
|
||||
|
||||
def finalize_close_session(self, session: DuplexSessionRuntimeState) -> None:
|
||||
session.finalize_close()
|
||||
if self._sessions.get(session.session_id) is session:
|
||||
self._sessions.pop(session.session_id, None)
|
||||
|
||||
def collect_expired(
|
||||
self,
|
||||
now: float | None = None,
|
||||
*,
|
||||
excluded_session_ids: set[str] | None = None,
|
||||
) -> list[DuplexSessionExpiry]:
|
||||
effective_now = self._clock() if now is None else now
|
||||
excluded = excluded_session_ids or set()
|
||||
expired: list[DuplexSessionExpiry] = []
|
||||
for session_id, session in list(self._sessions.items()):
|
||||
if session_id in excluded:
|
||||
continue
|
||||
if session.lease.disconnect_grace_expired(effective_now):
|
||||
reason = "disconnect_grace_expired"
|
||||
elif session.lease.idle_expired(effective_now):
|
||||
reason = "idle_ttl_expired"
|
||||
else:
|
||||
continue
|
||||
submitted = tuple(session.resource_request_ids(submitted=True))
|
||||
reserved = tuple(session.resource_request_ids(submitted=False))
|
||||
if not session.begin_close(session.fence, reason=reason):
|
||||
continue
|
||||
expired.append(
|
||||
DuplexSessionExpiry(
|
||||
session_id=session_id,
|
||||
fence=session.fence,
|
||||
lease_generation=session.lease.generation,
|
||||
reason=reason,
|
||||
submitted_request_ids=submitted,
|
||||
reserved_request_ids=reserved,
|
||||
)
|
||||
)
|
||||
return expired
|
||||
|
||||
def close_sessions_for_request_ids(self, request_ids: list[str]) -> dict[str, list[str]]:
|
||||
request_id_set = set(request_ids)
|
||||
closed: dict[str, list[str]] = {}
|
||||
for session_id, session in list(self._sessions.items()):
|
||||
stale = session.resource_request_ids()
|
||||
if request_id_set.isdisjoint(stale):
|
||||
continue
|
||||
if not session.begin_close(session.fence, reason="request_cleanup"):
|
||||
continue
|
||||
closed[session_id] = stale
|
||||
return closed
|
||||
|
||||
def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None:
|
||||
for session_id in session_ids:
|
||||
session = self._sessions.get(session_id)
|
||||
if session is None or session.lease.terminal_reason is None:
|
||||
continue
|
||||
self.finalize_close_session(session)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DuplexAppendReservation",
|
||||
"DuplexCompletedAppend",
|
||||
"DuplexFenceMismatchError",
|
||||
"DuplexInputAppend",
|
||||
"DuplexLeaseActivity",
|
||||
"DuplexLeaseConfig",
|
||||
"DuplexLeaseState",
|
||||
"DuplexRequestResource",
|
||||
"DuplexSessionExpiry",
|
||||
"DuplexSessionRuntimeManager",
|
||||
"DuplexSessionRuntimeState",
|
||||
"DuplexStageBinding",
|
||||
]
|
||||
@@ -0,0 +1,99 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TypedDict
|
||||
|
||||
|
||||
class DuplexIntermediateBuffer(TypedDict, total=False):
|
||||
"""Structured keys carried through ``model_intermediate_buffer``.
|
||||
|
||||
The buffer remains a dict for scheduler and msgspec compatibility, but
|
||||
duplex-specific producers and consumers should use the helpers in this
|
||||
module instead of scattering nested string keys across serving, runner, and
|
||||
model code.
|
||||
"""
|
||||
|
||||
request_id: str
|
||||
global_request_id: list[str]
|
||||
prompt_token_ids: list[int]
|
||||
llm_output_token_ids: list[int]
|
||||
llm_output_text: list[str]
|
||||
stream_output: bool
|
||||
native_duplex: bool
|
||||
ids: dict[str, object]
|
||||
hidden_states: dict[str, object]
|
||||
codes: dict[str, object]
|
||||
meta: dict[str, object]
|
||||
duplex: dict[str, object]
|
||||
omni_payload: object
|
||||
waveform: object
|
||||
mel_spec: object
|
||||
|
||||
|
||||
def build_duplex_intermediate_buffer(
|
||||
*,
|
||||
request_id: str,
|
||||
prompt_token_ids: list[int] | None = None,
|
||||
output_token_ids: list[int] | None = None,
|
||||
output_text: str | None = None,
|
||||
stream_output: bool = False,
|
||||
native_duplex: bool = False,
|
||||
) -> DuplexIntermediateBuffer:
|
||||
buffer: DuplexIntermediateBuffer = {
|
||||
"global_request_id": [str(request_id)],
|
||||
"ids": {},
|
||||
}
|
||||
if prompt_token_ids is not None:
|
||||
prompt_ids = [int(token_id) for token_id in prompt_token_ids]
|
||||
buffer["prompt_token_ids"] = prompt_ids
|
||||
buffer["ids"]["prompt"] = prompt_ids
|
||||
if output_token_ids is not None:
|
||||
output_ids = [int(token_id) for token_id in output_token_ids]
|
||||
buffer["llm_output_token_ids"] = output_ids
|
||||
buffer["ids"]["output"] = output_ids
|
||||
if output_text is not None:
|
||||
buffer["llm_output_text"] = [output_text]
|
||||
if stream_output:
|
||||
buffer["stream_output"] = True
|
||||
if native_duplex:
|
||||
buffer["native_duplex"] = True
|
||||
return buffer
|
||||
|
||||
|
||||
def set_ref_audio(buffer: dict[str, object], waveform: object, sample_rate_hz: int) -> None:
|
||||
buffer.setdefault("codes", {})["ref"] = waveform
|
||||
buffer.setdefault("meta", {})["ref_audio_sr"] = int(sample_rate_hz)
|
||||
|
||||
|
||||
def set_tts_handoff(buffer: dict[str, object], token_ids: object | None, hidden_states: object | None) -> None:
|
||||
"""Store the generic AR-stage token/hidden-state handoff consumed by TTS."""
|
||||
if token_ids is not None:
|
||||
buffer.setdefault("ids", {})["tts"] = token_ids
|
||||
if hidden_states is not None:
|
||||
buffer.setdefault("hidden_states", {})["tts"] = hidden_states
|
||||
|
||||
|
||||
def get_tts_handoff(info: dict[str, object]) -> tuple[object | None, object | None]:
|
||||
"""Read the generic AR-stage to TTS handoff, including legacy flat keys."""
|
||||
ids_info = info.get("ids")
|
||||
hidden_info = info.get("hidden_states")
|
||||
tts_token_ids = ids_info.get("tts") if isinstance(ids_info, dict) else None
|
||||
tts_hidden_states = hidden_info.get("tts") if isinstance(hidden_info, dict) else None
|
||||
if tts_token_ids is None:
|
||||
tts_token_ids = info.get("tts_token_ids")
|
||||
if tts_hidden_states is None:
|
||||
tts_hidden_states = info.get("tts_hidden_states")
|
||||
return tts_token_ids, tts_hidden_states
|
||||
|
||||
|
||||
def get_stream_request_key(info: dict[str, object]) -> str:
|
||||
key = info.get("global_request_id") or info.get("request_id") or info.get("_omni_req_id")
|
||||
if isinstance(key, (list, tuple)):
|
||||
key = key[0] if key else None
|
||||
if isinstance(key, bytes):
|
||||
key = key.decode("utf-8", errors="replace")
|
||||
if key is None:
|
||||
raise ValueError(
|
||||
"Duplex streaming handoff requires a stable request id; "
|
||||
"expected global_request_id, request_id, or _omni_req_id."
|
||||
)
|
||||
return str(key)
|
||||
@@ -0,0 +1,129 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Generic engine lease primitives shared by optional session runtimes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
|
||||
|
||||
class DuplexLeaseActivity(str, Enum):
|
||||
APPEND = "append"
|
||||
SIGNAL = "signal"
|
||||
PLAYBACK_ACK = "playback_ack"
|
||||
HEARTBEAT = "heartbeat"
|
||||
ATTACH = "attach"
|
||||
DETACH = "detach"
|
||||
RESUME = "resume"
|
||||
MODEL_OUTPUT = "model_output"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexLeaseConfig:
|
||||
idle_ttl_s: float | None = 300.0
|
||||
disconnect_grace_s: float = 30.0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.idle_ttl_s is not None and self.idle_ttl_s <= 0:
|
||||
raise ValueError("idle_ttl_s must be positive or null")
|
||||
if self.disconnect_grace_s <= 0:
|
||||
raise ValueError("disconnect_grace_s must be positive")
|
||||
|
||||
|
||||
@dataclass
|
||||
class DuplexLeaseState:
|
||||
config: DuplexLeaseConfig
|
||||
generation: int
|
||||
last_activity: float
|
||||
detached_at: float | None = None
|
||||
active_operations: set[str] = field(default_factory=set)
|
||||
terminal_reason: str | None = None
|
||||
|
||||
@property
|
||||
def expires_at(self) -> float | None:
|
||||
if self.config.idle_ttl_s is None:
|
||||
return None
|
||||
return self.last_activity + self.config.idle_ttl_s
|
||||
|
||||
def _require_open(self) -> None:
|
||||
if self.terminal_reason is not None:
|
||||
raise RuntimeError(f"duplex lease is terminal: {self.terminal_reason}")
|
||||
|
||||
def touch(self, now: float, activity: DuplexLeaseActivity) -> None:
|
||||
self._require_open()
|
||||
if not isinstance(activity, DuplexLeaseActivity):
|
||||
raise TypeError("lease activity must be DuplexLeaseActivity")
|
||||
self.last_activity = now
|
||||
|
||||
def detach(self, now: float) -> None:
|
||||
self.touch(now, DuplexLeaseActivity.DETACH)
|
||||
self.detached_at = now
|
||||
|
||||
def resume(self, now: float, *, expected_generation: int) -> int:
|
||||
self._require_open()
|
||||
if expected_generation != self.generation:
|
||||
raise ValueError(f"duplex lease generation mismatch: expected {self.generation}, got {expected_generation}")
|
||||
self.generation += 1
|
||||
self.detached_at = None
|
||||
self.last_activity = now
|
||||
return self.generation
|
||||
|
||||
def begin_operation(self, now: float, operation_id: str) -> None:
|
||||
self._require_open()
|
||||
if not operation_id:
|
||||
raise ValueError("duplex lease operation_id must not be empty")
|
||||
if operation_id in self.active_operations:
|
||||
raise ValueError(f"duplex lease operation already active: {operation_id}")
|
||||
self.active_operations.add(operation_id)
|
||||
self.last_activity = now
|
||||
|
||||
def end_operation(self, now: float, operation_id: str) -> None:
|
||||
self._require_open()
|
||||
if operation_id not in self.active_operations:
|
||||
raise KeyError(f"duplex lease operation is not active: {operation_id}")
|
||||
self.active_operations.remove(operation_id)
|
||||
self.last_activity = now
|
||||
|
||||
def disconnect_grace_expired(self, now: float) -> bool:
|
||||
return (
|
||||
self.terminal_reason is None
|
||||
and self.detached_at is not None
|
||||
and now >= self.detached_at + self.config.disconnect_grace_s
|
||||
)
|
||||
|
||||
def idle_expired(self, now: float) -> bool:
|
||||
expires_at = self.expires_at
|
||||
return (
|
||||
self.terminal_reason is None and not self.active_operations and expires_at is not None and now >= expires_at
|
||||
)
|
||||
|
||||
def mark_terminal(self, reason: str) -> bool:
|
||||
if self.terminal_reason is not None:
|
||||
return False
|
||||
self.terminal_reason = reason
|
||||
self.generation += 1
|
||||
self.detached_at = None
|
||||
self.active_operations.clear()
|
||||
return True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DuplexSessionExpiry:
|
||||
session_id: str
|
||||
fence: DuplexFence
|
||||
lease_generation: int
|
||||
reason: str
|
||||
submitted_request_ids: tuple[str, ...]
|
||||
reserved_request_ids: tuple[str, ...]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DuplexLeaseActivity",
|
||||
"DuplexLeaseConfig",
|
||||
"DuplexLeaseState",
|
||||
"DuplexSessionExpiry",
|
||||
]
|
||||
@@ -0,0 +1,130 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Queue contracts owned by the experimental duplex control plane."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
from vllm_omni.engine.messages import EngineQueueMessage
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DuplexFence:
|
||||
"""Identity fence carried by experimental duplex control messages."""
|
||||
|
||||
session_id: str
|
||||
epoch: int = 0
|
||||
turn_id: int = 0
|
||||
response_seq: int = 0
|
||||
incarnation: int = 0
|
||||
|
||||
|
||||
class OpenDuplexSessionMessage(EngineQueueMessage, kw_only=True):
|
||||
type: Literal["open_duplex_session"] = "open_duplex_session"
|
||||
control_id: str
|
||||
fence: DuplexFence
|
||||
session_id: str
|
||||
session_mode: str = "duplex"
|
||||
capabilities: dict[str, object]
|
||||
session_config: dict[str, object] | None = None
|
||||
runtime_config: dict[str, object] | None = None
|
||||
|
||||
|
||||
class AppendDuplexInputMessage(EngineQueueMessage, kw_only=True):
|
||||
type: Literal["append_duplex_input"] = "append_duplex_input"
|
||||
control_id: str
|
||||
operation_id: str | None = None
|
||||
fence: DuplexFence
|
||||
session_id: str
|
||||
expected_epoch: int | None = None
|
||||
mode: str
|
||||
payload: object
|
||||
final: bool = False
|
||||
|
||||
|
||||
class SignalDuplexTurnMessage(EngineQueueMessage, kw_only=True):
|
||||
type: Literal["signal_duplex_turn"] = "signal_duplex_turn"
|
||||
control_id: str
|
||||
fence: DuplexFence
|
||||
session_id: str
|
||||
event: str
|
||||
next_fence: DuplexFence | None = None
|
||||
session_config: dict[str, object] | None = None
|
||||
runtime_config: dict[str, object] | None = None
|
||||
|
||||
|
||||
class CloseDuplexSessionMessage(EngineQueueMessage, kw_only=True):
|
||||
type: Literal["close_duplex_session"] = "close_duplex_session"
|
||||
control_id: str
|
||||
fence: DuplexFence
|
||||
session_id: str
|
||||
reason: str = "client_close"
|
||||
|
||||
|
||||
class TouchDuplexSessionMessage(EngineQueueMessage, kw_only=True):
|
||||
type: Literal["touch_duplex_session"] = "touch_duplex_session"
|
||||
control_id: str
|
||||
fence: DuplexFence
|
||||
session_id: str
|
||||
activity: str
|
||||
|
||||
|
||||
class ResumeDuplexSessionMessage(EngineQueueMessage, kw_only=True):
|
||||
type: Literal["resume_duplex_session"] = "resume_duplex_session"
|
||||
control_id: str
|
||||
fence: DuplexFence
|
||||
session_id: str
|
||||
expected_lease_generation: int
|
||||
|
||||
|
||||
class DuplexSessionLifecycleMessage(EngineQueueMessage, kw_only=True):
|
||||
type: Literal["duplex_session_lifecycle"] = "duplex_session_lifecycle"
|
||||
fence: DuplexFence
|
||||
session_id: str
|
||||
event: str
|
||||
reason: str
|
||||
lease_generation: int
|
||||
submitted_request_ids: list[str]
|
||||
reserved_request_ids: list[str]
|
||||
|
||||
|
||||
class DuplexControlError(EngineQueueMessage, kw_only=True):
|
||||
code: str
|
||||
message: str
|
||||
retryable: bool = False
|
||||
|
||||
|
||||
class DuplexControlResultMessage(EngineQueueMessage, kw_only=True):
|
||||
type: Literal["duplex_control_result"] = "duplex_control_result"
|
||||
control_id: str
|
||||
fence: DuplexFence
|
||||
operation: str
|
||||
session_id: str
|
||||
ok: bool
|
||||
stage_results: list[dict[str, object]]
|
||||
unsupported_count: int = 0
|
||||
error_count: int = 0
|
||||
error: DuplexControlError | None = None
|
||||
accepted_fence: DuplexFence | None = None
|
||||
lease_generation: int | None = None
|
||||
|
||||
@property
|
||||
def rpc_correlation_key(self) -> tuple[str, str]:
|
||||
return ("duplex", self.control_id)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AppendDuplexInputMessage",
|
||||
"CloseDuplexSessionMessage",
|
||||
"DuplexControlError",
|
||||
"DuplexControlResultMessage",
|
||||
"DuplexFence",
|
||||
"DuplexSessionLifecycleMessage",
|
||||
"OpenDuplexSessionMessage",
|
||||
"ResumeDuplexSessionMessage",
|
||||
"SignalDuplexTurnMessage",
|
||||
"TouchDuplexSessionMessage",
|
||||
]
|
||||
Loaded 100 of 142 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user