anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
ed6bc5ae0cd22b6fb9b34a152337aed85126dc14
vllm-omni/tests/model_executor/models/personaplex/duplex/test_plugin.py
T

524 lines
21 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""PersonaPlex on the unified full-duplex framework: the plugin contract and the E2E driver helpers."""
from __future__ import annotations
import argparse
import base64
from pathlib import Path
from types import SimpleNamespace
from typing import cast
import numpy as np
import pytest
from vllm.sampling_params import SamplingParams
from tests.e2e.online_serving import personaplex_realtime_duplex as e2e_driver
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.engine.duplex.config import DuplexSessionConfig
from vllm_omni.engine.duplex.contracts import DuplexFence
from vllm_omni.engine.duplex.plugin import (
DefaultDuplexModelSessionState,
DuplexDataPlaneContext,
DuplexRuntimeConfigError,
load_duplex_plugin,
validate_duplex_plugin_sampling,
)
from vllm_omni.model_executor.models.personaplex.configuration_personaplex import PersonaPlexConfig
from vllm_omni.model_executor.models.personaplex.duplex import stage0
from vllm_omni.model_executor.models.personaplex.duplex.config import DEFAULT_PERSONA, FRAME_SIZE, SAMPLE_RATE
from vllm_omni.model_executor.models.personaplex.duplex.data_plane import PersonaPlexDataPlaneSession
from vllm_omni.model_executor.models.personaplex.duplex.input import PersonaPlexPcmAppendBuffer
from vllm_omni.model_executor.models.personaplex.duplex.plugin import (
PRIVATE_RUNTIME_CONFIG_KEYS,
PersonaPlexDuplexPlugin,
PersonaPlexSessionState,
)
from vllm_omni.model_executor.models.personaplex.pipeline import PERSONAPLEX_PIPELINE
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
DEPLOY_PATH = Path(__file__).parents[5] / "vllm_omni" / "deploy" / "personaplex.yaml"
def _pcm_payload(samples: np.ndarray, *, sample_rate_hz: int = SAMPLE_RATE) -> dict[str, object]:
samples = np.ascontiguousarray(samples, dtype="<f4")
return {
"type": "audio",
"format": "pcm_f32le",
"sample_rate_hz": sample_rate_hz,
"audio": base64.b64encode(samples.tobytes()).decode("ascii"),
}
def _encode_audio(audio, _sample_rate, _response_format, _speed):
if audio is None:
return None
size = int(np.asarray(audio, dtype=np.float32).size)
return f"audio-{size}" if size else None
def _plugin() -> PersonaPlexDuplexPlugin:
return PersonaPlexDuplexPlugin(_encode_audio)
# --------------------------------------------------------------------------- #
# Pipeline binding #
# --------------------------------------------------------------------------- #
def test_pipeline_declares_the_duplex_plugin_and_no_pre_framework_wiring() -> None:
assert PERSONAPLEX_PIPELINE.duplex_plugin == (
"vllm_omni.model_executor.models.personaplex.duplex.plugin.PersonaPlexDuplexPlugin"
)
assert PERSONAPLEX_PIPELINE.duplex_runtime_extension is None
assert PERSONAPLEX_PIPELINE.duplex_serving_adapter is None
assert PERSONAPLEX_PIPELINE.duplex_control_enabled is False
plugin = load_duplex_plugin(PERSONAPLEX_PIPELINE.duplex_plugin, _encode_audio)
assert isinstance(plugin, PersonaPlexDuplexPlugin)
assert plugin.plugin_id == "personaplex"
assert isinstance(plugin.data_plane, PersonaPlexDataPlaneSession)
validate_duplex_plugin_sampling(plugin, sampling_defaults=(SamplingParams(), SamplingParams()))
@pytest.mark.parametrize(
("engine_arg", "expected"),
[("skip_tokenizer_init", True), ("enable_prefix_caching", False)],
)
def test_personaplex_stage_engine_args(engine_arg: str, expected: bool) -> None:
stages = merge_pipeline_deploy(PERSONAPLEX_PIPELINE, load_deploy_config(DEPLOY_PATH))
assert [stage.yaml_engine_args.get(engine_arg) for stage in stages] == [expected, expected]
def test_personaplex_deploy_is_duplex_and_propagates_capacity_to_all_model_stages() -> None:
deploy = load_deploy_config(DEPLOY_PATH)
stages = merge_pipeline_deploy(PERSONAPLEX_PIPELINE, deploy)
assert deploy.session_mode == "duplex"
assert deploy.duplex_session.max_sessions == 2
assert [stage.yaml_engine_args.get("duplex_max_sessions") for stage in stages] == [2, 2]
# Lockstep frames of all sessions must share a step (see the deploy comment).
assert stages[0].yaml_engine_args.get("async_scheduling") is False
assert "personaplex_codec_max_sessions" not in deploy.connectors["connector_of_shared_memory"]["extra"]
def test_personaplex_deploy_enables_mimi_cuda_graphs_on_both_stage_configs(tmp_path: Path) -> None:
# The checkpoint ships an empty config.json; each stage's engine args build
# its PersonaPlexConfig from that plus the deploy YAML's hf_overrides.
(tmp_path / "config.json").write_text("{}")
stages = merge_pipeline_deploy(PERSONAPLEX_PIPELINE, load_deploy_config(DEPLOY_PATH))
def resolve(model_arch: str, hf_overrides: dict[str, object] | None) -> PersonaPlexConfig:
engine_args = OmniEngineArgs(
model=str(tmp_path),
model_arch=model_arch,
hf_overrides=hf_overrides,
skip_tokenizer_init=True,
trust_remote_code=True,
)
config = engine_args.create_model_config().hf_config
assert isinstance(config, PersonaPlexConfig)
return config
assert [stage.yaml_engine_args["model_arch"] for stage in stages] == [
"PersonaPlexTalkerForConditionalGeneration",
"PersonaPlexCode2Wav",
]
for stage in stages:
model_arch = stage.yaml_engine_args["model_arch"]
assert stage.yaml_engine_args["hf_overrides"] == {"mimi_cuda_graphs": True}
assert resolve(model_arch, dict(stage.yaml_engine_args["hf_overrides"])).mimi_cuda_graphs is True
# Without the deploy override the flag keeps its default.
assert resolve(model_arch, None).mimi_cuda_graphs is False
# --------------------------------------------------------------------------- #
# Session policy #
# --------------------------------------------------------------------------- #
def test_capabilities_are_honest() -> None:
single = _plugin().capabilities(max_sessions=1)
multi = _plugin().capabilities(max_sessions=2)
assert multi.as_dict()["implementation_level"] == "model_native_duplex"
assert multi.as_dict()["input_modes"] == ["append_audio_chunk"]
assert multi.chunk_period_ms == 80
assert single.supports_multi_session is False
assert multi.supports_multi_session is True
assert multi.supports_multi_session_same_replica is True
assert multi.supports_barge_in is False
assert multi.supports_client_commit is False
assert multi.supports_external_turn_signal is False
assert multi.supports_session_resume is False
assert multi.supports_audio_truncate is False
assert multi.supports_chat_completions is False
assert multi.session_admission_mode == "engine_managed"
def test_session_state_is_the_framework_default_with_the_personaplex_buffer() -> None:
state = _plugin().create_session_state()
assert isinstance(state, PersonaPlexSessionState)
assert isinstance(state, DefaultDuplexModelSessionState)
assert isinstance(state.audio_buffer, PersonaPlexPcmAppendBuffer)
state.retain_committed_audio({"audio": ""}, operation_id="op", reserved_bytes=8)
assert state.clear_committed_audio() == 8
assert state.committed_audio_payload is None
def test_private_runtime_keys_are_rejected_in_extra_body() -> None:
plugin = _plugin()
assert PRIVATE_RUNTIME_CONFIG_KEYS == {
"personaplex_prefill_slots",
"personaplex_model_path",
"personaplex_voice_prompt",
"personaplex_persona",
}
plugin.validate_client_extra_body({"auto_response": True})
plugin.validate_client_extra_body(None)
with pytest.raises(DuplexRuntimeConfigError, match="personaplex_persona") as excinfo:
plugin.validate_client_extra_body({"personaplex_persona": "x", "auto_response": True})
assert excinfo.value.code == "invalid_duplex_runtime_config"
@pytest.mark.asyncio
async def test_prepare_runtime_config_resolves_voice_and_prefill_slots_once(monkeypatch) -> None:
calls: list[tuple[str, str, str]] = []
def fake_prefill_slots(model_path: str, voice: str, persona: str) -> int:
calls.append((model_path, voice, persona))
return 42
monkeypatch.setattr(stage0, "personaplex_prefill_slots", fake_prefill_slots)
plugin = _plugin()
model_config = SimpleNamespace(model="/models/personaplex")
runtime_config = await plugin.prepare_runtime_config(
DuplexSessionConfig(voice="NATM1.pt", instructions="Be brief."),
model_config=model_config,
)
again = await plugin.prepare_runtime_config(
DuplexSessionConfig(voice="NATM1.pt", instructions="Be brief."),
model_config=model_config,
)
defaults = await plugin.prepare_runtime_config(DuplexSessionConfig(), model_config=model_config)
assert runtime_config == {
"personaplex_model_path": "/models/personaplex",
"personaplex_voice_prompt": "NATM1.pt",
"personaplex_persona": "Be brief.",
"personaplex_prefill_slots": 42,
}
assert again == runtime_config
assert defaults["personaplex_voice_prompt"] == "NATF2.pt"
assert defaults["personaplex_persona"] == DEFAULT_PERSONA
# The voice bundle and tokenizer are read once per (model, voice, persona).
assert calls == [
("/models/personaplex", "NATM1.pt", "Be brief."),
("/models/personaplex", "NATF2.pt", DEFAULT_PERSONA),
]
@pytest.mark.asyncio
async def test_prepare_runtime_config_reports_prefill_failures_as_runtime_config_errors(monkeypatch) -> None:
def failing_prefill_slots(model_path: str, voice: str, persona: str) -> int:
raise FileNotFoundError("voices.tgz missing")
monkeypatch.setattr(stage0, "personaplex_prefill_slots", failing_prefill_slots)
with pytest.raises(DuplexRuntimeConfigError, match="voices.tgz missing") as excinfo:
await _plugin().prepare_runtime_config(DuplexSessionConfig(), model_config=SimpleNamespace(model="/m"))
assert excinfo.value.code == "prefill_unavailable"
with pytest.raises(DuplexRuntimeConfigError, match="model path") as excinfo:
await _plugin().prepare_runtime_config(DuplexSessionConfig(), model_config=None)
assert excinfo.value.code == "model_path_unavailable"
@pytest.mark.parametrize("voice", ["../NATF2.pt", "voices/NATF2.pt", "NATF2.wav", "/abs/NATF2.pt"])
@pytest.mark.asyncio
async def test_prepare_runtime_config_rejects_voices_that_are_not_bundled_basenames(voice: str) -> None:
with pytest.raises(DuplexRuntimeConfigError, match="bundled .pt basename") as excinfo:
await _plugin().prepare_runtime_config(
DuplexSessionConfig(voice=voice), model_config=SimpleNamespace(model="/m")
)
assert excinfo.value.code == "invalid_voice"
def test_runtime_config_update_rejects_changed_persona() -> None:
plugin = _plugin()
short_persona = "You are a helpful assistant."
current = {"personaplex_persona": short_persona, "personaplex_voice_prompt": "NATF2.pt"}
with pytest.raises(DuplexRuntimeConfigError, match="persona") as excinfo:
plugin.runtime_config_for_update(DuplexSessionConfig(instructions="You are a pirate."), current)
assert excinfo.value.code == "persona_update_unsupported"
unchanged = plugin.runtime_config_for_update(DuplexSessionConfig(instructions=short_persona), current)
assert unchanged["personaplex_persona"] == short_persona
assert unchanged is not current
def test_runtime_config_update_rejects_changed_voice() -> None:
plugin = _plugin()
current = {"personaplex_persona": DEFAULT_PERSONA, "personaplex_voice_prompt": "NATF2.pt"}
with pytest.raises(DuplexRuntimeConfigError, match="voice") as excinfo:
plugin.runtime_config_for_update(DuplexSessionConfig(voice="NATM1.pt"), current)
assert excinfo.value.code == "voice_update_unsupported"
unchanged = plugin.runtime_config_for_update(DuplexSessionConfig(voice="NATF2.pt"), current)
assert unchanged["personaplex_voice_prompt"] == "NATF2.pt"
def test_data_plane_context_is_the_framework_default() -> None:
context = _plugin().data_plane_context(
epoch=1,
turn_id=2,
active_response_turn_id=None,
active_response_id="resp",
auto_responds=True,
response_format="pcm16",
speed=None,
modalities=("audio",),
)
assert isinstance(context, DuplexDataPlaneContext)
assert context.response_format == "pcm16" and context.epoch == 1
# --------------------------------------------------------------------------- #
# Engine policy #
# --------------------------------------------------------------------------- #
def test_sampling_params_are_greedy_one_token_on_stage0_only() -> None:
defaults = (SamplingParams(temperature=0.8, top_k=10, max_tokens=10), SamplingParams(max_tokens=1024))
configured = _plugin().configure_sampling_params(runtime_config={}, defaults=defaults)
assert configured[0].temperature == 0.0
assert configured[0].top_k == 1
assert configured[0].max_tokens == 1
assert configured[1] is defaults[1]
assert defaults[0].max_tokens == 10
assert _plugin().configure_sampling_params(runtime_config={}, defaults=()) == ()
def test_plan_append_reserves_one_slot_plus_prefill_on_the_first_seq() -> None:
plugin = _plugin()
fence = DuplexFence("session", epoch=3, turn_id=1)
common = {
"request_id": "req",
"fence": fence,
"session_config": {"instructions": "Be concise."},
"runtime_config": {"personaplex_prefill_slots": 4},
"turn_seq": 1,
"payload": _pcm_payload(np.zeros(FRAME_SIZE, np.float32)),
"final": False,
"sampling_params": SamplingParams(),
}
first = plugin.plan_append(seq=1, **common)
second = plugin.plan_append(seq=2, **common)
assert len(first.prompt["prompt_token_ids"]) == 5
assert len(second.prompt["prompt_token_ids"]) == 1
duplex = first.prompt["model_intermediate_buffer"]["duplex"]
assert duplex["session_id"] == "session"
assert duplex["epoch"] == 3
assert duplex["turn_id"] == 1
assert duplex["seq"] == 1
assert duplex["data_plane"] is True
assert duplex["fence"] == fence
assert duplex["scheduler_token_budget"] == 5
assert duplex["runtime_config"] == {"personaplex_prefill_slots": 4}
assert "incarnation" not in duplex
assert first.prompt["model_intermediate_buffer"]["global_request_id"] == ["session"]
@pytest.mark.parametrize(
("payload", "match"),
[
(_pcm_payload(np.zeros(FRAME_SIZE - 1, np.float32)), "exactly 1920 samples"),
(_pcm_payload(np.zeros(FRAME_SIZE, np.float32), sample_rate_hz=16000), "sample_rate_hz must be 24000"),
({"format": "pcm16", "sample_rate_hz": SAMPLE_RATE, "audio": ""}, "format must be pcm_f32le"),
("frame", "must be a mapping"),
],
)
def test_plan_append_rejects_anything_but_one_24k_frame(payload: object, match: str) -> None:
with pytest.raises(ValueError, match=match):
_plugin().plan_append(
request_id="req",
fence=DuplexFence("session"),
session_config={},
runtime_config={},
seq=1,
turn_seq=1,
payload=payload,
final=False,
sampling_params=None,
)
def test_plan_append_rejects_a_malformed_prefill_slot_count() -> None:
with pytest.raises(ValueError, match="personaplex_prefill_slots"):
_plugin().plan_append(
request_id="req",
fence=DuplexFence("session"),
session_config={},
runtime_config={"personaplex_prefill_slots": "many"},
seq=1,
turn_seq=1,
payload=_pcm_payload(np.zeros(FRAME_SIZE, np.float32)),
final=False,
sampling_params=None,
)
def test_decide_output_never_decides() -> None:
assert (
_plugin().decide_output(
stage_id=0,
final_stage_id=1,
segment_finished=True,
segment_token_ids=(1,),
segment_output_metadata={},
output=SimpleNamespace(outputs=[]),
)
is None
)
def test_silence_unit_is_one_24k_frame_that_plan_append_accepts() -> None:
plugin = _plugin()
unit = plugin.silence_unit_payload()
assert unit["format"] == "pcm_f32le"
assert unit["sample_rate_hz"] == SAMPLE_RATE
assert len(base64.b64decode(unit["audio"])) == FRAME_SIZE * 4
plan = plugin.plan_append(
request_id="req",
fence=DuplexFence("session"),
session_config={},
runtime_config={},
seq=7,
turn_seq=7,
payload=unit,
final=False,
sampling_params=None,
)
assert len(plan.prompt["prompt_token_ids"]) == 1
# --------------------------------------------------------------------------- #
# E2E driver helpers (CPU-checked parts of the GPU driver) #
# --------------------------------------------------------------------------- #
def _audio_client(frame_count: int = 10, *, voiced_frames: int = 0) -> e2e_driver.RawRealtimeProbe:
"""A probe double carrying only the collected events the driver's result builder reads."""
frames = np.full((frame_count, e2e_driver.FRAME_SAMPLES), 7, dtype="<i2")
frames[:voiced_frames] = 1000
raw = frames.tobytes()
event = {
"type": "response.output_audio.delta",
"response_id": "response-1",
"sample_rate_hz": e2e_driver.SAMPLE_RATE_HZ,
"metadata": {
"vllm_omni": {
"runtime_impl": "scheduler_data_plane",
"uses_model_runner_scheduler": True,
"runner_kv_backed": True,
}
},
}
events = SimpleNamespace(events=[event], response_audio={"response-1": [raw]}, audio_bytes=lambda: raw)
return cast(e2e_driver.RawRealtimeProbe, SimpleNamespace(events=events))
def test_realtime_audio_frame_stats_separate_voiced_and_silent_frames() -> None:
silent = np.zeros(e2e_driver.FRAME_SAMPLES, dtype="<i2")
voiced = np.full(e2e_driver.FRAME_SAMPLES, 1000, dtype="<i2")
stats = e2e_driver._audio_frame_stats(
silent.tobytes() + voiced.tobytes(),
input_frames=3,
voiced_frame_rms_threshold=1e-4,
)
assert stats == {
"output_frames": 2,
"voiced_frames": 1,
"silent_frames": 1,
"frame_deficit": 1,
"frame_coverage_ratio": pytest.approx(2 / 3),
}
def test_realtime_audio_frame_stats_reject_partial_codec_frame() -> None:
partial_frame = np.zeros(e2e_driver.FRAME_SAMPLES - 1, dtype="<i2")
with pytest.raises(AssertionError, match="whole 80 ms codec frames"):
e2e_driver._audio_frame_stats(partial_frame.tobytes(), input_frames=1, voiced_frame_rms_threshold=1e-4)
@pytest.mark.parametrize("frame_count", [10, 40])
def test_e2e_driver_rejects_inaudible_sessions(frame_count: int) -> None:
args = argparse.Namespace(max_frame_deficit=4, voiced_frame_rms_threshold=1e-3, min_voiced_frames=5)
with pytest.raises(AssertionError, match="audible speech"):
e2e_driver._session_result(_audio_client(frame_count), input_frames=frame_count, args=args, minimum_chunks=1)
def test_e2e_driver_uses_absolute_audible_floor() -> None:
args = argparse.Namespace(max_frame_deficit=4, voiced_frame_rms_threshold=1e-3, min_voiced_frames=5)
e2e_driver._session_result(_audio_client(200, voiced_frames=5), input_frames=200, args=args, minimum_chunks=1)
def test_e2e_driver_defaults_replacement_to_full_input() -> None:
args = e2e_driver.parse_args(["--model", "/model", "--input-wav", "/input.wav"])
assert args.replacement_frames == 0
assert args.min_voiced_frames == 5
def test_e2e_driver_records_and_validates_input_sha256(tmp_path: Path) -> None:
input_wav = tmp_path / "input.wav"
input_wav.write_bytes(b"fixture-wav")
identity = e2e_driver._input_identity(
input_wav,
expected_sha256="bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
)
assert identity == {
"path": str(input_wav.resolve()),
"sha256": "bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
}
def test_e2e_driver_rejects_unexpected_input_sha256(tmp_path: Path) -> None:
input_wav = tmp_path / "input.wav"
input_wav.write_bytes(b"fixture-wav")
with pytest.raises(ValueError, match="input WAV SHA-256 mismatch"):
e2e_driver._input_identity(input_wav, expected_sha256="0" * 64)
def test_e2e_driver_lets_the_server_allocate_the_session_id() -> None:
url = e2e_driver._realtime_url("ws://127.0.0.1:8099/v1/realtime?duplex=1", "nvidia/personaplex-7b-v1")
assert "session_id=" not in url
assert "duplex=1" in url and "autostart=0" in url
Reference in new issue View Git Blame Copy Permalink