tzhouamandClaude Fable 5.1
c893b638ad
Measured after the DELTA fix: 53 engine outputs for 52 rows, but the client still received everything at once because each per-step payload is the newest row only and the row alignment treated it as a cumulative snapshot (one yield at the first row, one at the end). accumulate_codec_rows folds delta rows and the final consolidated matrix into one growing matrix; unit test drives the generator with per-step deltas and asserts one yield per step. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
283 lines
12 KiB
Python
283 lines
12 KiB
Python
"""Codec-token streaming: wire format, row alignment, and handler behaviour (CPU)."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from fastapi import FastAPI, WebSocket
|
|
|
|
from vllm_omni.entrypoints.openai import codec_stream as cs
|
|
from vllm_omni.entrypoints.openai.protocol.audio import StreamingSpeechSessionConfig
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_pack_unpack_roundtrip_and_flags():
|
|
codes = np.random.randint(0, 2048, size=(25, 16)).astype(np.int64)
|
|
buf = cs.pack_codec_frame(7, torch.from_numpy(codes), eos=False, segment_end=True)
|
|
assert len(buf) == cs.HEADER_SIZE + 25 * 16 * 2
|
|
seq, n, flags, arr = cs.unpack_codec_frame(buf, 16)
|
|
assert seq == 7 and n == 25 and flags == cs.FLAG_SEGMENT_END
|
|
assert arr.dtype == np.int16 and np.array_equal(arr.astype(np.int64), codes)
|
|
# EOS with zero frames
|
|
buf = cs.pack_codec_frame(8, None, eos=True, codebooks=16)
|
|
seq, n, flags, arr = cs.unpack_codec_frame(buf, 16)
|
|
assert (seq, n, flags) == (8, 0, cs.FLAG_EOS) and arr.shape == (0, 16)
|
|
buf = cs.pack_codec_frame(2**32 + 5, np.zeros((1, 16), dtype=np.int64), reset=True)
|
|
assert cs.unpack_codec_frame(buf, 16)[0] == 5
|
|
assert cs.unpack_codec_frame(buf, 16)[2] == cs.FLAG_RESET
|
|
|
|
|
|
def test_pack_rejects_bad_shapes_and_ranges():
|
|
with pytest.raises(ValueError):
|
|
cs.pack_codec_frame(0, np.zeros((2, 2, 2), dtype=np.int64))
|
|
with pytest.raises(ValueError):
|
|
cs.pack_codec_frame(0, np.full((1, 16), 40000, dtype=np.int64))
|
|
with pytest.raises(ValueError):
|
|
cs.unpack_codec_frame(b"\x00\x01", 16)
|
|
with pytest.raises(ValueError):
|
|
cs.unpack_codec_frame(cs.pack_codec_frame(0, np.zeros((1, 16), dtype=np.int64))[:-2], 16)
|
|
|
|
|
|
def test_align_codec_rows_drops_prefill_placeholder_and_eos():
|
|
# 1 prefill placeholder row (zeros, emitted before any token) + 3 real frames + EOS row
|
|
codes = torch.tensor([[0] * 16, [5] * 16, [6] * 16, [7] * 16, [0] * 16])
|
|
token_ids = [5, 6, 7, 2150] # codebook-0 ids then EOS
|
|
rows = cs.align_codec_rows(codes, token_ids, 2048)
|
|
assert rows.shape == (3, 16) and rows[:, 0].tolist() == [5, 6, 7]
|
|
# trailing -1 placeholders in token ids are ignored
|
|
rows = cs.align_codec_rows(codes, token_ids + [-1, -1], 2048)
|
|
assert rows.shape == (3, 16)
|
|
# out-of-range rows are dropped
|
|
bad = torch.tensor([[5] * 16, [3000] * 16])
|
|
assert cs.align_codec_rows(bad, [5, 6], 2048).shape == (1, 16)
|
|
assert cs.align_codec_rows(None, [], 2048).shape == (0, 0)
|
|
assert cs.align_codec_rows(torch.zeros(0, 16, dtype=torch.long), [], 2048).shape == (0, 16)
|
|
|
|
|
|
def test_session_config_codec_mode_validation():
|
|
cfg = StreamingSpeechSessionConfig(output_mode="codec_tokens", stream_audio=True, response_format="pcm")
|
|
assert cfg.output_mode == "codec_tokens"
|
|
with pytest.raises(ValueError):
|
|
StreamingSpeechSessionConfig(output_mode="codec_tokens", stream_audio=False, response_format="pcm")
|
|
with pytest.raises(ValueError):
|
|
StreamingSpeechSessionConfig(
|
|
output_mode="codec_tokens", stream_audio=True, response_format="pcm", word_timestamps=True
|
|
)
|
|
assert StreamingSpeechSessionConfig().output_mode == "pcm"
|
|
|
|
|
|
def test_codec_start_message_fields():
|
|
spec = {"codebooks": 16, "codebook_size": 2048, "frame_rate_hz": 12.5, "sample_rate": 24000, "decoder_id": "x"}
|
|
msg = cs.codec_start_message(utterance_index=1, sentence_index=0, sentence_text="hi", spec=spec, model="m")
|
|
assert msg["type"] == "codec.start" and msg["codebooks"] == 16 and msg["header_format"] == cs.HEADER_FMT
|
|
assert msg["decoder_id"] == "x" and msg["version"] == cs.CODEC_STREAM_VERSION
|
|
|
|
|
|
# --------------------------------------------------------------------- handler
|
|
def _codec_app(mocker, *, supports: bool = True, cancel_after: int | None = None):
|
|
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
|
from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
|
|
|
|
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
|
speech_service.forced_aligner_enabled = False
|
|
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
|
|
adapter = mocker.MagicMock()
|
|
adapter.supports_codec_stream = supports
|
|
adapter.codec_stream_spec = {
|
|
"codebooks": 16,
|
|
"codebook_size": 2048,
|
|
"frame_rate_hz": 12.5,
|
|
"sample_rate": 24000,
|
|
"decoder_id": "d",
|
|
}
|
|
speech_service._get_tts_adapter = mocker.MagicMock(return_value=adapter)
|
|
speech_service.engine_client = mocker.MagicMock()
|
|
speech_service.engine_client.abort = mocker.AsyncMock()
|
|
|
|
async def mock_codec_chunks(_generator, _request_id, *, tts_params=None, min_frames=1):
|
|
import asyncio
|
|
|
|
for i in range(3):
|
|
await asyncio.sleep(0.01)
|
|
yield np.full((1 + i, 16), i, dtype=np.int64), i == 2
|
|
|
|
speech_service._generate_codec_chunks = mock_codec_chunks
|
|
handler = OmniStreamingSpeechHandler(speech_service=speech_service, idle_timeout=5.0, config_timeout=5.0)
|
|
app = FastAPI()
|
|
|
|
@app.websocket("/v1/audio/speech/stream")
|
|
async def ws_endpoint(websocket: WebSocket):
|
|
# NOTE: no ``from __future__ import annotations`` in this module: FastAPI
|
|
# must resolve the WebSocket annotation, or it treats it as a query param.
|
|
await handler.handle_session(websocket)
|
|
|
|
return app, speech_service
|
|
|
|
|
|
def test_codec_stream_handler_frames_and_done(mocker):
|
|
from starlette.testclient import TestClient
|
|
|
|
app, svc = _codec_app(mocker)
|
|
with TestClient(app) as client, client.websocket_connect("/v1/audio/speech/stream") as ws:
|
|
ws.send_json(
|
|
{"type": "session.config", "output_mode": "codec_tokens", "stream_audio": True, "response_format": "pcm"}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "hello"})
|
|
ws.send_json({"type": "input.done"})
|
|
start = ws.receive_json()
|
|
assert start["type"] == "codec.start" and start["codebooks"] == 16
|
|
frames = []
|
|
while True:
|
|
msg = ws.receive()
|
|
if "bytes" in msg and msg["bytes"] is not None:
|
|
frames.append(cs.unpack_codec_frame(msg["bytes"], 16))
|
|
continue
|
|
done = __import__("json").loads(msg["text"])
|
|
break
|
|
assert [f[1] for f in frames] == [1, 2, 3]
|
|
assert [f[0] for f in frames] == [0, 1, 2]
|
|
assert frames[-1][2] & cs.FLAG_EOS
|
|
assert done["type"] == "codec.done" and done["total_frames"] == 6 and done["cancelled"] is False
|
|
assert ws.receive_json()["type"] == "session.done"
|
|
svc.engine_client.abort.assert_not_awaited()
|
|
|
|
|
|
def test_codec_stream_rejected_when_adapter_lacks_support(mocker):
|
|
from starlette.testclient import TestClient
|
|
|
|
app, _ = _codec_app(mocker, supports=False)
|
|
with TestClient(app) as client, client.websocket_connect("/v1/audio/speech/stream") as ws:
|
|
ws.send_json(
|
|
{"type": "session.config", "output_mode": "codec_tokens", "stream_audio": True, "response_format": "pcm"}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "hello"})
|
|
ws.send_json({"type": "input.done"})
|
|
err = ws.receive_json()
|
|
assert err["type"] == "error" and "not supported" in err["message"]
|
|
|
|
|
|
def test_input_cancel_mid_stream_aborts_request(mocker):
|
|
from starlette.testclient import TestClient
|
|
|
|
app, svc = _codec_app(mocker)
|
|
with TestClient(app) as client, client.websocket_connect("/v1/audio/speech/stream") as ws:
|
|
ws.send_json(
|
|
{"type": "session.config", "output_mode": "codec_tokens", "stream_audio": True, "response_format": "pcm"}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "hello"})
|
|
ws.send_json({"type": "input.done"})
|
|
assert ws.receive_json()["type"] == "codec.start"
|
|
ws.send_json({"type": "input.cancel"})
|
|
types = []
|
|
while True:
|
|
msg = ws.receive()
|
|
if "bytes" in msg and msg["bytes"] is not None:
|
|
continue
|
|
payload = __import__("json").loads(msg["text"])
|
|
types.append(payload["type"])
|
|
if payload["type"] == "session.done":
|
|
break
|
|
if payload["type"] == "codec.done":
|
|
assert payload["cancelled"] is True
|
|
assert "codec.done" in types and types[-1] == "session.done"
|
|
svc.engine_client.abort.assert_awaited()
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_talker_only_deployment_detection():
|
|
from types import SimpleNamespace
|
|
|
|
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
|
|
|
probe = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
|
|
probe._tts_stage = SimpleNamespace(final_output=True, final_output_type="latent")
|
|
assert probe._is_talker_only_deployment() is True
|
|
probe._tts_stage = SimpleNamespace(final_output=False, final_output_type="latent")
|
|
assert probe._is_talker_only_deployment() is False
|
|
probe._tts_stage = SimpleNamespace(final_output=True, final_output_type="audio")
|
|
assert probe._is_talker_only_deployment() is False
|
|
probe._tts_stage = None
|
|
assert probe._is_talker_only_deployment() is False
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_codec_snapshot_list_is_normalized():
|
|
import numpy as np
|
|
import torch
|
|
|
|
from vllm_omni.entrypoints.openai.serving_speech import _coerce_codec_snapshots
|
|
|
|
assert _coerce_codec_snapshots(None) is None
|
|
t = torch.arange(32).reshape(2, 16)
|
|
assert _coerce_codec_snapshots(t) is t
|
|
# DELTA-mode list of cumulative snapshots -> last one
|
|
snaps = [torch.arange(16).reshape(1, 16), torch.arange(32).reshape(2, 16), torch.arange(48).reshape(3, 16)]
|
|
out = _coerce_codec_snapshots(snaps)
|
|
assert out.shape == (3, 16) and int(out[2, 15]) == 47
|
|
# list of single-row deltas -> concatenation
|
|
deltas = [torch.full((1, 16), i) for i in range(4)]
|
|
out = _coerce_codec_snapshots(deltas)
|
|
assert out.shape == (4, 16) and list(out[:, 0]) == [0, 1, 2, 3]
|
|
assert _coerce_codec_snapshots([None]) is None
|
|
assert isinstance(_coerce_codec_snapshots([np.zeros(16)]), np.ndarray)
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_accumulate_codec_rows_handles_deltas_and_snapshots():
|
|
import numpy as np
|
|
|
|
from vllm_omni.entrypoints.openai.codec_stream import accumulate_codec_rows
|
|
|
|
r = lambda i: np.full((1, 16), i, dtype=np.int64) # noqa: E731
|
|
acc = accumulate_codec_rows(None, r(0))
|
|
acc = accumulate_codec_rows(acc, r(1))
|
|
acc = accumulate_codec_rows(acc, r(2))
|
|
assert acc.shape == (3, 16) and list(acc[:, 0]) == [0, 1, 2]
|
|
# duplicate of the last delta is not appended twice
|
|
assert accumulate_codec_rows(acc, acc).shape == (3, 16)
|
|
# final cumulative snapshot with a leading placeholder row replaces the matrix
|
|
final = np.concatenate([np.full((1, 16), -1), acc, r(3)], axis=0)
|
|
out = accumulate_codec_rows(acc, final)
|
|
assert out.shape == (5, 16) and out[0, 0] == -1 and out[-1, 0] == 3
|
|
# 1-D row payload
|
|
assert accumulate_codec_rows(None, np.arange(16)).shape == (1, 16)
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
@pytest.mark.asyncio
|
|
async def test_generate_codec_chunks_yields_per_step():
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
|
|
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
|
|
|
svc = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
|
|
svc._get_tts_adapter = lambda: SimpleNamespace(
|
|
codec_stream_spec={"codebook_size": 2048}, validates_generation=False
|
|
)
|
|
|
|
def out(token_ids, codes, finished):
|
|
return SimpleNamespace(
|
|
outputs=[SimpleNamespace(token_ids=token_ids, multimodal_output={"codes": {"audio": codes}})],
|
|
finished=finished,
|
|
)
|
|
|
|
rows = [np.full((1, 16), i + 10, dtype=np.int64) for i in range(4)]
|
|
|
|
async def gen():
|
|
# step k: token k and its codes row arrive together (DELTA mode: newest row only)
|
|
for k in range(4):
|
|
yield out([100 + k], rows[k], False)
|
|
# final: EOS token plus the consolidated cumulative matrix (EOS row is all zeros)
|
|
final = np.concatenate(rows + [np.zeros((1, 16), dtype=np.int64)], axis=0)
|
|
yield out([2150], final, True)
|
|
|
|
got = [(None if c is None else c.shape[0], eos) async for c, eos in svc._generate_codec_chunks(gen(), "r1")]
|
|
assert got[:4] == [(1, False), (1, False), (1, False), (1, False)]
|
|
assert got[-1][1] is True and sum(n for n, _ in got if n) == 4
|