Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
90c890ad6a |
No files matched your search
@@ -457,6 +457,7 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
|
||||
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
|
||||
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
|
||||
- Under vocoder-bound overload (single-stream `rtf_p99 ≥ 1` at the target concurrency), set `active_stream_window: 2` at the top of the deploy yaml to cap simultaneously active Stage 1 streams. Off by default; trades TTFP for streaming continuity. See [#3592](https://github.com/vllm-project/vllm-omni/pull/3592) for the mechanism and tradeoff numbers.
|
||||
- For a bounded active-stream deployment, `active_stream_policy: edf` can prioritize streams whose server-estimated playback buffer is closest to exhaustion. Configure it under the connector's `extra` block and use a `codec_chunk_ramp`; this is opt-in, estimates playback from server-side PCM emission (no client playback acknowledgements), and targets first-audio latency and continuity rather than single-request completion latency.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_generation_scheduler_records_audio_output_for_edf(mocker):
|
||||
scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
|
||||
scheduler.chunk_transfer_adapter = SimpleNamespace(
|
||||
playback_deadline_enabled=True,
|
||||
record_audio_output=mocker.Mock(),
|
||||
)
|
||||
multimodal_output = {
|
||||
"model_outputs": torch.zeros(24000),
|
||||
"sr": torch.tensor(24000),
|
||||
}
|
||||
|
||||
scheduler._record_playback_deadline_output("req", multimodal_output)
|
||||
|
||||
scheduler.chunk_transfer_adapter.record_audio_output.assert_called_once_with("req", 24000, 24000)
|
||||
|
||||
|
||||
def test_generation_scheduler_skips_audio_accounting_when_edf_is_disabled(mocker):
|
||||
scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
|
||||
scheduler.chunk_transfer_adapter = SimpleNamespace(
|
||||
playback_deadline_enabled=False,
|
||||
record_audio_output=mocker.Mock(),
|
||||
)
|
||||
|
||||
scheduler._record_playback_deadline_output("req", {"model_outputs": torch.zeros(24000)})
|
||||
|
||||
scheduler.chunk_transfer_adapter.record_audio_output.assert_not_called()
|
||||
@@ -113,6 +113,7 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
|
||||
*,
|
||||
stage_id: int = 1,
|
||||
model_mode: str = "ar",
|
||||
engine_output_type: str | None = None,
|
||||
max_num_seqs: int = 2,
|
||||
active_stream_window: int = 0,
|
||||
connector_extra: dict | None = None,
|
||||
@@ -143,6 +144,7 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
|
||||
|
||||
model_config = SimpleNamespace(
|
||||
worker_type=model_mode,
|
||||
engine_output_type=engine_output_type or ("audio" if model_mode == "generation" else None),
|
||||
max_num_seqs=max_num_seqs,
|
||||
active_stream_window=active_stream_window,
|
||||
stage_connector_config={
|
||||
@@ -843,6 +845,218 @@ def test_fifo_promotion(build_adapter):
|
||||
assert reqs[2].status == RequestStatus.WAITING_FOR_CHUNK
|
||||
|
||||
|
||||
def test_edf_requires_active_window(build_adapter):
|
||||
with pytest.raises(ValueError, match="requires active_stream_window"):
|
||||
build_adapter(
|
||||
stage_id=1,
|
||||
model_mode="generation",
|
||||
active_stream_window=0,
|
||||
connector_extra={"active_stream_policy": "edf"},
|
||||
)
|
||||
|
||||
|
||||
def test_edf_rejects_unknown_policy(build_adapter):
|
||||
with pytest.raises(ValueError, match="active_stream_policy"):
|
||||
build_adapter(
|
||||
stage_id=1,
|
||||
model_mode="generation",
|
||||
active_stream_window=1,
|
||||
connector_extra={"active_stream_policy": "unknown"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("key", "value", "message"),
|
||||
[
|
||||
("playback_deadline_first_chunk_ms", float("nan"), "first_chunk"),
|
||||
("playback_deadline_safety_margin_ms", float("inf"), "safety_margin"),
|
||||
],
|
||||
)
|
||||
def test_edf_rejects_non_finite_timing_config(build_adapter, key, value, message):
|
||||
with pytest.raises(ValueError, match=message):
|
||||
build_adapter(
|
||||
stage_id=1,
|
||||
model_mode="generation",
|
||||
active_stream_window=1,
|
||||
connector_extra={"active_stream_policy": "edf", key: value},
|
||||
)
|
||||
|
||||
|
||||
def test_edf_selects_earliest_playback_deadline(build_adapter, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic",
|
||||
lambda: 100.0,
|
||||
)
|
||||
adapter, _ = build_adapter(
|
||||
stage_id=1,
|
||||
model_mode="generation",
|
||||
max_num_seqs=2,
|
||||
active_stream_window=1,
|
||||
connector_extra={"active_stream_policy": "edf"},
|
||||
)
|
||||
later = _req("later", RequestStatus.WAITING)
|
||||
urgent = _req("urgent", RequestStatus.WAITING)
|
||||
adapter.record_audio_output(later.request_id, 5, 1, emitted_at=100.0)
|
||||
adapter.record_audio_output(urgent.request_id, 2, 1, emitted_at=100.0)
|
||||
adapter.requests_with_ready_chunks.update({later.request_id, urgent.request_id})
|
||||
waiting_queue = DummyWaitingQueue([later, urgent])
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, [])
|
||||
|
||||
assert list(adapter._active_streams) == [urgent.request_id]
|
||||
assert waiting_queue == [urgent]
|
||||
assert list(adapter.waiting_for_chunk_waiting_requests) == [later]
|
||||
|
||||
|
||||
def test_edf_balances_first_audio_against_playback_deadline(build_adapter, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic",
|
||||
lambda: 100.0,
|
||||
)
|
||||
adapter, _ = build_adapter(
|
||||
stage_id=1,
|
||||
model_mode="generation",
|
||||
max_num_seqs=2,
|
||||
active_stream_window=1,
|
||||
connector_extra={
|
||||
"active_stream_policy": "edf",
|
||||
"playback_deadline_first_chunk_ms": 100,
|
||||
"playback_deadline_safety_margin_ms": 0,
|
||||
},
|
||||
)
|
||||
new_request = _req("new", RequestStatus.WAITING)
|
||||
buffered = _req("buffered", RequestStatus.WAITING)
|
||||
adapter.record_audio_output(buffered.request_id, 5, 1, emitted_at=100.0)
|
||||
adapter.requests_with_ready_chunks.update({new_request.request_id, buffered.request_id})
|
||||
waiting_queue = DummyWaitingQueue([buffered, new_request])
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, [])
|
||||
|
||||
assert list(adapter._active_streams) == [new_request.request_id]
|
||||
|
||||
adapter.cleanup_receiver(new_request.request_id)
|
||||
adapter.restore_queues(waiting_queue, [])
|
||||
starving = _req("starving", RequestStatus.WAITING)
|
||||
adapter.record_audio_output(starving.request_id, 1, 1, emitted_at=98.0)
|
||||
adapter.requests_with_ready_chunks.update({new_request.request_id, starving.request_id})
|
||||
waiting_queue = DummyWaitingQueue([new_request, starving])
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, [])
|
||||
|
||||
assert list(adapter._active_streams) == [starving.request_id]
|
||||
|
||||
|
||||
def test_edf_protects_near_underrun_stream_from_overdue_first_audio(build_adapter, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic",
|
||||
lambda: 100.0,
|
||||
)
|
||||
adapter, _ = build_adapter(
|
||||
stage_id=1,
|
||||
model_mode="generation",
|
||||
max_num_seqs=2,
|
||||
active_stream_window=1,
|
||||
connector_extra={
|
||||
"active_stream_policy": "edf",
|
||||
"playback_deadline_first_chunk_ms": 100,
|
||||
"playback_deadline_safety_margin_ms": 250,
|
||||
},
|
||||
)
|
||||
overdue_startup = _req("startup", RequestStatus.WAITING)
|
||||
near_underrun = _req("playing", RequestStatus.WAITING)
|
||||
adapter._playback_deadline_key(overdue_startup, order=0, now=90.0)
|
||||
adapter.record_audio_output(near_underrun.request_id, 100, 1000, emitted_at=100.0)
|
||||
adapter.requests_with_ready_chunks.update({overdue_startup.request_id, near_underrun.request_id})
|
||||
waiting_queue = DummyWaitingQueue([overdue_startup, near_underrun])
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, [])
|
||||
|
||||
assert list(adapter._active_streams) == [near_underrun.request_id]
|
||||
|
||||
|
||||
def test_edf_reselects_at_chunk_boundaries(build_adapter, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic",
|
||||
lambda: 100.0,
|
||||
)
|
||||
adapter, _ = build_adapter(
|
||||
stage_id=1,
|
||||
model_mode="generation",
|
||||
max_num_seqs=2,
|
||||
active_stream_window=1,
|
||||
connector_extra={"active_stream_policy": "edf", "playback_deadline_safety_margin_ms": 0},
|
||||
)
|
||||
first = _req("first", RequestStatus.WAITING)
|
||||
second = _req("second", RequestStatus.WAITING)
|
||||
adapter.record_audio_output(first.request_id, 1, 1, emitted_at=100.0)
|
||||
adapter.record_audio_output(second.request_id, 2, 1, emitted_at=100.0)
|
||||
adapter.requests_with_ready_chunks.update({first.request_id, second.request_id})
|
||||
waiting_queue = DummyWaitingQueue([first, second])
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, [])
|
||||
assert list(adapter._active_streams) == [first.request_id]
|
||||
|
||||
scheduler_output = SimpleNamespace(
|
||||
scheduled_new_reqs=[SimpleNamespace(req_id=first.request_id)],
|
||||
scheduled_cached_reqs=None,
|
||||
)
|
||||
adapter.postprocess_scheduler_output(scheduler_output)
|
||||
adapter.restore_queues(waiting_queue, [])
|
||||
adapter.record_audio_output(first.request_id, 5, 1, emitted_at=100.0)
|
||||
adapter.requests_with_ready_chunks.add(first.request_id)
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, [])
|
||||
|
||||
assert list(adapter._active_streams) == [second.request_id]
|
||||
|
||||
|
||||
def test_edf_k2_orders_ready_running_and_waiting_requests(build_adapter, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic",
|
||||
lambda: 100.0,
|
||||
)
|
||||
adapter, _ = build_adapter(
|
||||
stage_id=1,
|
||||
model_mode="generation",
|
||||
max_num_seqs=4,
|
||||
active_stream_window=2,
|
||||
connector_extra={"active_stream_policy": "edf", "playback_deadline_safety_margin_ms": 0},
|
||||
)
|
||||
later_running = _req("later-running", RequestStatus.RUNNING)
|
||||
urgent_waiting = _req("urgent-waiting", RequestStatus.WAITING)
|
||||
middle_running = _req("middle-running", RequestStatus.RUNNING)
|
||||
adapter.record_audio_output(later_running.request_id, 5, 1, emitted_at=100.0)
|
||||
adapter.record_audio_output(urgent_waiting.request_id, 1, 1, emitted_at=100.0)
|
||||
adapter.record_audio_output(middle_running.request_id, 2, 1, emitted_at=100.0)
|
||||
adapter.requests_with_ready_chunks.update(
|
||||
{later_running.request_id, urgent_waiting.request_id, middle_running.request_id}
|
||||
)
|
||||
waiting_queue = DummyWaitingQueue([urgent_waiting])
|
||||
running_queue = [later_running, middle_running]
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
|
||||
assert list(adapter._active_streams) == [urgent_waiting.request_id, middle_running.request_id]
|
||||
assert waiting_queue == [urgent_waiting]
|
||||
assert running_queue == [middle_running]
|
||||
assert list(adapter._held_non_active) == [later_running]
|
||||
|
||||
|
||||
def test_edf_cleanup_drops_playback_state(build_adapter):
|
||||
adapter, _ = build_adapter(
|
||||
stage_id=1,
|
||||
model_mode="generation",
|
||||
active_stream_window=1,
|
||||
connector_extra={"active_stream_policy": "edf"},
|
||||
)
|
||||
adapter.record_audio_output("req", 24000, 24000, emitted_at=100.0)
|
||||
assert "req" in adapter._playback_deadlines
|
||||
|
||||
adapter.cleanup_receiver("req")
|
||||
|
||||
assert "req" not in adapter._playback_deadlines
|
||||
|
||||
|
||||
def test_non_active_waiting_request_is_held_off_scheduler(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=2, max_num_seqs=2, active_stream_window=1)
|
||||
active = _req("req-active", RequestStatus.WAITING)
|
||||
|
||||
@@ -105,6 +105,122 @@ def test_flush_on_finish():
|
||||
assert len(p.codes.audio) == _Q * 24
|
||||
|
||||
|
||||
def _record_successful_chunk(tm, rid):
|
||||
tm.put_req_chunk[rid] += 1
|
||||
tm.ramp_chunk_count[rid] += 1
|
||||
|
||||
|
||||
def test_fixed_chunk_terminal_after_exact_initial_boundary_emits_no_audio():
|
||||
tm = _tm(initial_chunk_frames=10)
|
||||
rid = "fixed-exact-initial"
|
||||
|
||||
first = _call(tm, rid, n_frames=10, req_ic=10)
|
||||
assert first is not None
|
||||
_record_successful_chunk(tm, rid)
|
||||
|
||||
terminal = _call(tm, rid, n_frames=10, finished=True, req_ic=10)
|
||||
|
||||
assert terminal is not None
|
||||
assert terminal.meta.finished.item() is True
|
||||
assert terminal.codes.audio.numel() == 0
|
||||
|
||||
|
||||
def test_fixed_chunk_terminal_after_exact_steady_boundary_emits_no_audio():
|
||||
tm = _tm(initial_chunk_frames=10)
|
||||
rid = "fixed-exact-steady"
|
||||
|
||||
first = _call(tm, rid, n_frames=10, req_ic=10)
|
||||
assert first is not None
|
||||
_record_successful_chunk(tm, rid)
|
||||
steady = _call(tm, rid, n_frames=35, req_ic=10)
|
||||
assert steady is not None
|
||||
_record_successful_chunk(tm, rid)
|
||||
|
||||
terminal = _call(tm, rid, n_frames=35, finished=True, req_ic=10)
|
||||
|
||||
assert terminal is not None
|
||||
assert terminal.meta.finished.item() is True
|
||||
assert terminal.codes.audio.numel() == 0
|
||||
|
||||
|
||||
def test_fixed_chunk_terminal_flushes_exact_boundary_not_previously_sent():
|
||||
tm = _tm(initial_chunk_frames=10)
|
||||
rid = "fixed-unsent-boundary"
|
||||
|
||||
first = _call(tm, rid, n_frames=10, req_ic=10)
|
||||
assert first is not None
|
||||
_record_successful_chunk(tm, rid)
|
||||
|
||||
terminal = _call(tm, rid, n_frames=35, finished=True, req_ic=10)
|
||||
|
||||
assert terminal is not None
|
||||
assert terminal.meta.finished.item() is True
|
||||
assert terminal.codes.audio.numel() == _Q * 25
|
||||
|
||||
|
||||
def test_fixed_chunk_without_initial_terminal_after_exact_boundary_emits_no_audio():
|
||||
tm = _tm(chunk_frames=4, initial_chunk_frames=4)
|
||||
rid = "fixed-no-initial-exact"
|
||||
|
||||
first = _call(tm, rid, n_frames=4, req_ic=4)
|
||||
assert first is not None
|
||||
_record_successful_chunk(tm, rid)
|
||||
second = _call(tm, rid, n_frames=8, req_ic=4)
|
||||
assert second is not None
|
||||
_record_successful_chunk(tm, rid)
|
||||
|
||||
terminal = _call(tm, rid, n_frames=8, finished=True, req_ic=4)
|
||||
|
||||
assert terminal is not None
|
||||
assert terminal.codes.audio.numel() == 0
|
||||
|
||||
|
||||
def test_fixed_chunk_terminal_after_exact_boundary_flushes_only_new_tail():
|
||||
tm = _tm(initial_chunk_frames=10)
|
||||
rid = "fixed-tail-after-boundary"
|
||||
|
||||
first = _call(tm, rid, n_frames=10, req_ic=10)
|
||||
assert first is not None
|
||||
_record_successful_chunk(tm, rid)
|
||||
steady = _call(tm, rid, n_frames=35, req_ic=10)
|
||||
assert steady is not None
|
||||
_record_successful_chunk(tm, rid)
|
||||
|
||||
terminal = _call(tm, rid, n_frames=40, finished=True, req_ic=10)
|
||||
|
||||
assert terminal is not None
|
||||
assert terminal.codes.audio.numel() == _Q * 5
|
||||
|
||||
|
||||
def test_fixed_chunk_stream_sends_every_codec_frame_exactly_once():
|
||||
tm = _tm(chunk_frames=4, initial_chunk_frames=2)
|
||||
rid = "fixed-unique"
|
||||
frames = [[frame * 10 + quantizer for quantizer in range(_Q)] for frame in range(10)]
|
||||
emitted = []
|
||||
|
||||
for length in (2, 6, 10):
|
||||
tm.code_prompt_token_ids[rid] = frames[:length]
|
||||
payload = talker2code2wav_async_chunk(
|
||||
transfer_manager=tm,
|
||||
multimodal_output={"codes": {"audio": torch.zeros((0,))}},
|
||||
request=_req(rid, finished=False, initial_codec_chunk_frames=2),
|
||||
)
|
||||
assert payload is not None
|
||||
emitted.extend(payload.codes.audio.reshape(_Q, -1).transpose(0, 1).tolist())
|
||||
_record_successful_chunk(tm, rid)
|
||||
|
||||
terminal = talker2code2wav_async_chunk(
|
||||
transfer_manager=tm,
|
||||
multimodal_output=None,
|
||||
request=_req(rid, finished=True, initial_codec_chunk_frames=2),
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert terminal is not None
|
||||
assert terminal.codes.audio.numel() == 0
|
||||
assert emitted == frames
|
||||
|
||||
|
||||
_CASES = [
|
||||
# ── IC boundary rule ──────────────────────────────────────────────
|
||||
# initial_codec_chunk_frames only controls the first emitted chunk.
|
||||
|
||||
@@ -24,6 +24,8 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
|
||||
from vllm_omni.core.sched.utils import omni_routed_experts_for_request
|
||||
from vllm_omni.metrics import definitions as metric_defs
|
||||
from vllm_omni.metrics.utils import count_audio_frames
|
||||
from vllm_omni.outputs import OmniModelRunnerOutput
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -70,6 +72,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
return False
|
||||
return super()._handle_stopped_request(request)
|
||||
|
||||
def _record_playback_deadline_output(self, request_id: str, multimodal_output: Any) -> None:
|
||||
adapter = self.chunk_transfer_adapter
|
||||
if adapter is None or not getattr(adapter, "playback_deadline_enabled", False):
|
||||
return
|
||||
audio_frames = count_audio_frames(multimodal_output)
|
||||
sample_rate = metric_defs.resolve_audio_sample_rate(multimodal_output)
|
||||
adapter.record_audio_output(request_id, audio_frames, sample_rate)
|
||||
|
||||
def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
|
||||
"""One-shot generation fast path:
|
||||
- Feed all input tokens of the request at once
|
||||
@@ -440,6 +450,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
kv_transfer_params = None
|
||||
pooler_output = pooler_outputs[req_index] if pooler_outputs else None
|
||||
mm_output = mm_outputs[req_index] if mm_outputs else None
|
||||
if mm_output is not None:
|
||||
self._record_playback_deadline_output(req_id, mm_output)
|
||||
status_before_stop = request.status
|
||||
finish_reason = None
|
||||
is_segment_finished = False
|
||||
|
||||
@@ -26,6 +26,14 @@ connectors:
|
||||
codec_left_context_frames: 72
|
||||
# Emit only the first audio chunk early, then return to codec_chunk_frames.
|
||||
initial_codec_chunk_frames: 1
|
||||
# Optional Stage-1 EDF scheduler. Requires top-level
|
||||
# active_stream_window > 0. It estimates each stream's playback deadline
|
||||
# from actual PCM emitted by Code2Wav and reselects the K most urgent
|
||||
# ready chunks at every scheduling turn. Use it with codec_chunk_ramp or
|
||||
# an initial chunk large enough to cover the first follow-up decode.
|
||||
# active_stream_policy: edf
|
||||
# playback_deadline_first_chunk_ms: 100
|
||||
# playback_deadline_safety_margin_ms: 250
|
||||
# Chunk ramp (opt-in): gradual size increase for first N chunks to
|
||||
# eliminate the IC→steady cliff (1→25). Uncomment to enable.
|
||||
# Each entry is the size of that chunk index; indices past the table
|
||||
|
||||
@@ -2,9 +2,11 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import importlib
|
||||
import math
|
||||
import time
|
||||
from collections import defaultdict, deque
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
@@ -22,6 +24,13 @@ from .base import OmniTransferAdapterBase
|
||||
logger = get_connector_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PlaybackDeadlineState:
|
||||
first_seen_at: float
|
||||
playback_deadline: float | None = None
|
||||
emitted_audio_s: float = 0.0
|
||||
|
||||
|
||||
class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
"""Chunk-level transfer adapter for Omni connector pipelines.
|
||||
|
||||
@@ -42,25 +51,54 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
|
||||
def __init__(self, vllm_config: Any):
|
||||
model_config = vllm_config.model_config
|
||||
self.model_mode = getattr(model_config, "worker_type", None) or "ar"
|
||||
self.engine_output_type = getattr(model_config, "engine_output_type", None)
|
||||
self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs
|
||||
active_stream_window = int(getattr(model_config, "active_stream_window", 0) or 0)
|
||||
model_max_num_seqs = int(getattr(model_config, "max_num_seqs", self.scheduler_max_num_seqs) or 0)
|
||||
if model_max_num_seqs <= 0:
|
||||
model_max_num_seqs = self.scheduler_max_num_seqs
|
||||
self._active_window = min(active_stream_window, model_max_num_seqs) if active_stream_window > 0 else 0
|
||||
connector_config = getattr(model_config, "stage_connector_config", None)
|
||||
raw_extra = (
|
||||
connector_config.get("extra", connector_config)
|
||||
if isinstance(connector_config, dict)
|
||||
else getattr(connector_config, "extra", None)
|
||||
)
|
||||
extra = raw_extra if isinstance(raw_extra, dict) else {}
|
||||
self._active_stream_policy = str(extra.get("active_stream_policy", "fifo")).strip().lower()
|
||||
if self._active_stream_policy not in {"fifo", "edf"}:
|
||||
raise ValueError(f"active_stream_policy must be 'fifo' or 'edf', got {self._active_stream_policy!r}")
|
||||
if self._active_stream_policy == "edf" and self._active_window <= 0:
|
||||
raise ValueError("active_stream_policy='edf' requires active_stream_window > 0")
|
||||
self._playback_first_chunk_slo_s = float(extra.get("playback_deadline_first_chunk_ms", 100.0)) / 1000.0
|
||||
self._playback_safety_margin_s = float(extra.get("playback_deadline_safety_margin_ms", 50.0)) / 1000.0
|
||||
if not math.isfinite(self._playback_first_chunk_slo_s) or self._playback_first_chunk_slo_s <= 0:
|
||||
raise ValueError("playback_deadline_first_chunk_ms must be positive")
|
||||
if not math.isfinite(self._playback_safety_margin_s) or self._playback_safety_margin_s < 0:
|
||||
raise ValueError("playback_deadline_safety_margin_ms must be non-negative")
|
||||
self._playback_deadlines: dict[str, _PlaybackDeadlineState] = {}
|
||||
if self._active_window > 0:
|
||||
logger.info(
|
||||
"Bounded active-stream window enabled: K=%d. "
|
||||
"Bounded active-stream window enabled: K=%d policy=%s. "
|
||||
"Multi-replica deployments require sticky per-stream routing across Stage 1 "
|
||||
"replicas (each replica owns an independent active-set; without sticky routing, "
|
||||
"a stream can be active on one replica and non-active on another and both will "
|
||||
"race to evict it).",
|
||||
self._active_window,
|
||||
self._active_stream_policy,
|
||||
)
|
||||
self.connector = self.create_connector(model_config)
|
||||
self.receives_chunks = stage_receives_chunks(model_config)
|
||||
if self.playback_deadline_enabled and not extra.get("codec_chunk_ramp"):
|
||||
initial_frames = int(extra.get("initial_codec_chunk_frames") or 0)
|
||||
if 0 < initial_frames < 4:
|
||||
logger.warning(
|
||||
"Playback EDF is enabled with initial_codec_chunk_frames=%d and no codec_chunk_ramp; "
|
||||
"the first audio buffer may be too short to avoid a chunk-0 underrun.",
|
||||
initial_frames,
|
||||
)
|
||||
super().__init__(model_config)
|
||||
self.model_mode = getattr(model_config, "worker_type", None) or "ar"
|
||||
# State specific to Chunk management
|
||||
self.custom_process_next_stage_input_func: Callable[..., OmniPayloadStruct | None] | None = None
|
||||
custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None)
|
||||
@@ -103,6 +141,51 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
# OmniSchedulingCoordinator._waiting_since on the full-payload path.
|
||||
self._waiting_since: dict[str, float] = {}
|
||||
|
||||
@property
|
||||
def playback_deadline_enabled(self) -> bool:
|
||||
return (
|
||||
self._active_stream_policy == "edf"
|
||||
and self.model_mode == "generation"
|
||||
and self.engine_output_type == "audio"
|
||||
)
|
||||
|
||||
def record_audio_output(
|
||||
self,
|
||||
request_id: str,
|
||||
audio_frames: int,
|
||||
sample_rate: int,
|
||||
*,
|
||||
emitted_at: float | None = None,
|
||||
) -> None:
|
||||
"""Advance the server-estimated playback deadline after one PCM chunk."""
|
||||
if not self.playback_deadline_enabled or audio_frames <= 0 or sample_rate <= 0:
|
||||
return
|
||||
now = time.monotonic() if emitted_at is None else float(emitted_at)
|
||||
state = self._playback_deadlines.setdefault(request_id, _PlaybackDeadlineState(first_seen_at=now))
|
||||
duration_s = float(audio_frames) / float(sample_rate)
|
||||
state.emitted_audio_s += duration_s
|
||||
# A late chunk restarts a starved player at its server emission time;
|
||||
# otherwise append its duration to the buffered playback horizon.
|
||||
playback_base = max(state.playback_deadline or now, now)
|
||||
state.playback_deadline = playback_base + duration_s
|
||||
|
||||
def _playback_deadline_key(self, request: Request, order: int, now: float) -> tuple[int, float, float, int]:
|
||||
state = self._playback_deadlines.setdefault(
|
||||
request.request_id,
|
||||
_PlaybackDeadlineState(first_seen_at=now),
|
||||
)
|
||||
if state.playback_deadline is None:
|
||||
deadline = state.first_seen_at + self._playback_first_chunk_slo_s
|
||||
protected_playback = False
|
||||
else:
|
||||
deadline = state.playback_deadline
|
||||
protected_playback = deadline - now <= self._playback_safety_margin_s
|
||||
# A stream whose buffered playback is close to exhaustion must not be
|
||||
# preempted by a large backlog of overdue first-audio deadlines. Once
|
||||
# established streams have enough headroom, ordinary EDF ordering lets
|
||||
# new requests consume the available slots.
|
||||
return (0 if protected_playback else 1), deadline, state.first_seen_at, order
|
||||
|
||||
@staticmethod
|
||||
def _is_truthy_scalar(value: Any) -> bool:
|
||||
if isinstance(value, torch.Tensor):
|
||||
@@ -465,6 +548,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
Idempotent: calling with an already-cleaned or unknown id is safe.
|
||||
"""
|
||||
self._active_streams.pop(request_id, None)
|
||||
self._playback_deadlines.pop(request_id, None)
|
||||
self.upstream_exhausted_requests.discard(request_id)
|
||||
self.segment_finished_requests.discard(request_id)
|
||||
self.get_req_chunk.pop(request_id, None)
|
||||
@@ -586,6 +670,19 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
waiting_queue.prepend_requests([request])
|
||||
return
|
||||
|
||||
if self.playback_deadline_enabled:
|
||||
self._process_chunk_queue_legacy(
|
||||
waiting_queue, self.waiting_for_chunk_waiting_requests, RequestStatus.WAITING, self._finished_load_reqs
|
||||
)
|
||||
self._process_chunk_queue_legacy(
|
||||
running_queue,
|
||||
self.waiting_for_chunk_running_requests,
|
||||
RequestStatus.RUNNING,
|
||||
self._finished_load_reqs,
|
||||
)
|
||||
self._select_playback_deadline_streams(waiting_queue, running_queue)
|
||||
return
|
||||
|
||||
self._promote_active_streams(running_queue)
|
||||
self._promote_active_streams(waiting_queue)
|
||||
self._process_chunk_queue(
|
||||
@@ -624,6 +721,45 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
# Iterating the existing queue preserves FIFO admission.
|
||||
self._active_streams[request_id] = request
|
||||
|
||||
def _select_playback_deadline_streams(self, waiting_queue: Any, running_queue: list[Request]) -> None:
|
||||
"""Select up to K ready chunks by server-estimated playback deadline."""
|
||||
now = time.monotonic()
|
||||
candidates = [
|
||||
request
|
||||
for request in (*running_queue, *list(waiting_queue))
|
||||
if request.request_id in self.requests_with_ready_chunks
|
||||
]
|
||||
ordered = sorted(
|
||||
enumerate(candidates),
|
||||
key=lambda item: self._playback_deadline_key(item[1], item[0], now),
|
||||
)
|
||||
selected = [request for _, request in ordered[: self._active_window]]
|
||||
selected_ids = {request.request_id for request in selected}
|
||||
self._active_streams = {request.request_id: request for request in selected}
|
||||
|
||||
for request in list(running_queue):
|
||||
if request.request_id in self.requests_with_ready_chunks and request.request_id not in selected_ids:
|
||||
running_queue.remove(request)
|
||||
self._held_non_active.append(request)
|
||||
|
||||
held_waiting = []
|
||||
for request in list(waiting_queue):
|
||||
if request.request_id in self.requests_with_ready_chunks and request.request_id not in selected_ids:
|
||||
waiting_queue.remove(request)
|
||||
self.requests_origin_status[request.request_id] = RequestStatus.WAITING
|
||||
self.waiting_for_chunk_waiting_requests.append(request)
|
||||
elif request.request_id in selected_ids:
|
||||
held_waiting.append(request)
|
||||
|
||||
# Keep selected running requests in EDF order so a tight token budget
|
||||
# still admits the most urgent chunk first.
|
||||
running_order = {request.request_id: index for index, request in enumerate(selected)}
|
||||
running_queue.sort(key=lambda request: running_order.get(request.request_id, len(running_order)))
|
||||
if len(held_waiting) > 1:
|
||||
held_waiting.sort(key=lambda request: running_order[request.request_id])
|
||||
waiting_queue.remove_requests(held_waiting)
|
||||
waiting_queue.prepend_requests(held_waiting)
|
||||
|
||||
def _ensure_active_stream(self, request: Request) -> bool:
|
||||
if self._active_window <= 0:
|
||||
return True
|
||||
|
||||
@@ -29,6 +29,17 @@ from vllm_omni.model_executor.stage_input_processors.tts_utils import (
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _empty_finished_payload(request_id: str) -> OmniPayloadStruct:
|
||||
return OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
|
||||
meta=MetaStruct(
|
||||
request_id=request_id,
|
||||
left_context_size=0,
|
||||
finished=torch.tensor(True, dtype=torch.bool),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _qwen3_tts_degenerate_finished_payload():
|
||||
"""Single-placeholder-frame finished payload for a degenerate talker take.
|
||||
|
||||
@@ -162,14 +173,7 @@ def talker2code2wav_async_chunk(
|
||||
|
||||
if length <= 0:
|
||||
if finished:
|
||||
return OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
|
||||
meta=MetaStruct(
|
||||
request_id=request_id,
|
||||
left_context_size=0,
|
||||
finished=torch.tensor(True, dtype=torch.bool),
|
||||
),
|
||||
)
|
||||
return _empty_finished_payload(request_id)
|
||||
return None
|
||||
|
||||
if ramp is not None:
|
||||
@@ -179,18 +183,28 @@ def talker2code2wav_async_chunk(
|
||||
if not emit:
|
||||
return None
|
||||
if context_length == 0:
|
||||
return OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
|
||||
meta=MetaStruct(
|
||||
request_id=request_id,
|
||||
left_context_size=0,
|
||||
finished=torch.tensor(True, dtype=torch.bool),
|
||||
),
|
||||
)
|
||||
return _empty_finished_payload(request_id)
|
||||
else:
|
||||
first_chunk = int(transfer_manager.put_req_chunk.get(request_id, 0)) <= 0
|
||||
chunk_index = int(transfer_manager.ramp_chunk_count.get(request_id, 0))
|
||||
first_chunk = chunk_index <= 0
|
||||
use_first_chunk = initial_chunk_size > 0 and initial_chunk_size < chunk_size
|
||||
|
||||
if use_first_chunk:
|
||||
emitted_frames = 0 if first_chunk else initial_chunk_size + (chunk_index - 1) * chunk_size
|
||||
else:
|
||||
emitted_frames = chunk_index * chunk_size
|
||||
if finished and length <= emitted_frames:
|
||||
# EOS is emitted on a separate talker step. If the preceding codec
|
||||
# frame completed an exact chunk boundary, that chunk has already
|
||||
# been sent; terminal must carry only an empty finish marker.
|
||||
logger.debug(
|
||||
"Qwen3-TTS terminal has no unsent codec frames: request=%s length=%d emitted=%d",
|
||||
request_id,
|
||||
length,
|
||||
emitted_frames,
|
||||
)
|
||||
return _empty_finished_payload(request_id)
|
||||
|
||||
if use_first_chunk and length <= initial_chunk_size:
|
||||
if not finished and length < initial_chunk_size:
|
||||
return None
|
||||
|
||||
Reference in new issue
Block a user