anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
59db6a4285e9086391efc85a6eff6030a6294878
vllm-omni/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
T
Zhou TaichangClaude Opus 5Gao Han
 59db6a4285 [Rebase] Rebase to vLLM 0.30.0 (#7820) 
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-09-22 15:52:01 +00:00

3531 lines
137 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import threading
from collections import deque
from types import MethodType, SimpleNamespace
from unittest.mock import Mock, call, patch
import pytest
import torch
from pytest_mock import MockerFixture
from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
from vllm.v1.request import Request, RequestStatus
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt
from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
OmniChunkTransferAdapter,
)
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class DummyWaitingQueue(list):
def prepend_requests(self, requests):
self[:0] = list(requests)
def add_request(self, request):
self.append(request)
def remove_requests(self, requests):
remove = set(requests)
self[:] = [request for request in self if request not in remove]
def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None):
request = Mock(
client_index=0,
request_id=req_id,
external_req_id=external_req_id or req_id,
status=status,
prompt_token_ids=[],
num_prompt_tokens=0,
num_computed_tokens=0,
num_output_placeholders=0,
prefill_stats=None,
additional_information=None,
resumable=False,
)
request.is_finished = lambda: RequestStatus.is_finished(request.status)
return request
def _streaming_request(mocker: MockerFixture, num_computed_tokens: int) -> SimpleNamespace:
prompt_token_ids = [0] * num_computed_tokens
return SimpleNamespace(
_all_token_ids=prompt_token_ids.copy(),
_output_token_ids=[],
prompt_token_ids=prompt_token_ids,
num_computed_tokens=num_computed_tokens,
num_prompt_tokens=num_computed_tokens,
num_output_placeholders=0,
update_block_hashes=mocker.Mock(),
)
def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture) -> None:
request = SimpleNamespace(
_all_token_ids=[0, 0, 7, 8],
_output_token_ids=[7, 8],
prompt_token_ids=[0, 0],
num_computed_tokens=4,
num_prompt_tokens=2,
update_block_hashes=mocker.Mock(),
)
payload = {
"ids": {"prompt": [1, 2, 3]},
"meta": {
"replace_streaming_prompt": True,
"next_stage_prompt_len": 7,
},
}
construct_next_stage_streaming_input_prompt(payload, request)
assert request.prompt_token_ids == [0] * 7
assert request._all_token_ids == [0] * 7
assert request._output_token_ids == []
assert request.num_computed_tokens == 0
assert request.num_prompt_tokens == 7
request.update_block_hashes.assert_called_once_with()
def test_turn_start_replacement_ignores_accumulated_prompt_capacity(mocker: MockerFixture) -> None:
request = _streaming_request(mocker, num_computed_tokens=4064)
payload = {
"ids": {"prompt": [1]},
"meta": {
"replace_streaming_prompt": True,
"next_stage_prompt_len": 10,
"next_stage_generation_tokens": 26,
},
}
replaced = construct_next_stage_streaming_input_prompt(
payload,
request,
max_model_len=4096,
previous_condition_len=10,
previous_condition_seq=7,
condition_seq=8,
recompute_previous_chunks=1,
)
assert replaced is True
assert request.prompt_token_ids == [0] * 10
assert payload["meta"]["streaming_prompt_recompute"] is False
assert "streaming_prompt_previous_codes" not in payload["ids"]
@pytest.mark.parametrize("reserve", [True, 0, -1, 1.5, "26", None])
def test_streaming_prompt_rejects_invalid_generation_reserve(
mocker: MockerFixture,
reserve: object,
) -> None:
request = _streaming_request(mocker, num_computed_tokens=10)
payload = {
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 10,
"next_stage_generation_tokens": reserve,
},
}
with pytest.raises(ValueError, match="next_stage_generation_tokens must be a positive integer"):
construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)
def test_capacity_managed_streaming_window_at_limit_still_appends(mocker: MockerFixture) -> None:
request = _streaming_request(mocker, num_computed_tokens=4060)
payload = {
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 10,
"next_stage_generation_tokens": 26,
},
}
replaced = construct_next_stage_streaming_input_prompt(
payload,
request,
max_model_len=4096,
previous_condition_len=10,
previous_condition_seq=0,
condition_seq=1,
recompute_previous_chunks=1,
recompute_on_capacity=True,
)
assert replaced is False
assert request.num_computed_tokens == 4060
assert request.num_prompt_tokens == 4070
def test_capacity_managed_streaming_prompt_appends_from_declared_length_without_ids_prompt(
mocker: MockerFixture,
) -> None:
request = _streaming_request(mocker, num_computed_tokens=20)
payload = {
"ids": {"tts": [1, 2, 3]},
"hidden_states": {"tts": [[0.1], [0.2], [0.3]]},
"meta": {
"next_stage_prompt_len": 4,
"next_stage_generation_tokens": 26,
},
}
replaced = construct_next_stage_streaming_input_prompt(
payload,
request,
max_model_len=4096,
previous_condition_len=4,
previous_condition_seq=0,
condition_seq=1,
recompute_previous_chunks=1,
recompute_on_capacity=True,
)
assert replaced is False
assert request.num_computed_tokens == 20
assert request.num_prompt_tokens == 24
assert request.prompt_token_ids == [0] * 24
request.update_block_hashes.assert_called_once_with()
def test_streaming_window_builds_one_chunk_recompute_recipe(mocker: MockerFixture) -> None:
previous_condition_len = 10
previous_codes = list(range(25))
request = SimpleNamespace(
_all_token_ids=[0] * 4039 + previous_codes + [999],
_output_token_ids=previous_codes + [999],
prompt_token_ids=[0] * 4039,
num_computed_tokens=4065,
num_prompt_tokens=4039,
num_output_placeholders=1,
update_block_hashes=mocker.Mock(),
)
payload = {
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 10,
"next_stage_generation_tokens": 26,
},
}
replaced = construct_next_stage_streaming_input_prompt(
payload,
request,
max_model_len=4096,
previous_condition_len=previous_condition_len,
previous_condition_seq=7,
condition_seq=8,
recompute_previous_chunks=1,
)
assert replaced is True
assert request.prompt_token_ids == [0] * 45
assert request.num_prompt_tokens == 45
assert request.num_computed_tokens == 0
assert payload["ids"]["streaming_prompt_previous_codes"] == previous_codes
assert payload["meta"] == {
"next_stage_prompt_len": 10,
"next_stage_generation_tokens": 26,
"streaming_prompt_recompute": True,
"streaming_condition_seq": 8,
}
request.update_block_hashes.assert_called_once_with()
def test_capacity_policy_turn_start_replacement_does_not_require_generation_reserve(
mocker: MockerFixture,
) -> None:
request = _streaming_request(mocker, num_computed_tokens=4064)
payload = {
"ids": {"prompt": [1]},
"meta": {
"replace_streaming_prompt": True,
"next_stage_prompt_len": 10,
},
}
replaced = construct_next_stage_streaming_input_prompt(
payload,
request,
max_model_len=4096,
previous_condition_len=10,
previous_condition_seq=7,
condition_seq=8,
recompute_previous_chunks=1,
recompute_on_capacity=True,
)
assert replaced is True
assert request.num_computed_tokens == 0
assert request.num_prompt_tokens == 10
def test_streaming_window_appends_until_capacity_then_recomputes(mocker: MockerFixture) -> None:
request = SimpleNamespace(
_all_token_ids=[0] * 4000 + [101, 102],
_output_token_ids=[101, 102],
prompt_token_ids=[0] * 4000,
num_computed_tokens=4002,
num_prompt_tokens=4000,
num_output_placeholders=0,
update_block_hashes=mocker.Mock(),
)
second = {
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 12,
"next_stage_generation_tokens": 26,
},
}
assert not construct_next_stage_streaming_input_prompt(
second,
request,
max_model_len=4096,
previous_condition_len=10,
previous_condition_seq=0,
condition_seq=1,
recompute_previous_chunks=1,
recompute_on_capacity=True,
)
assert request.num_computed_tokens == 4002
assert request.num_prompt_tokens == 4014
assert "streaming_prompt_previous_codes" not in second["ids"]
assert second["meta"]["streaming_prompt_recompute"] is False
request._all_token_ids.extend([201, 202, 203])
request._output_token_ids.extend([201, 202, 203])
request.num_computed_tokens = 4017
third = {
"ids": {"prompt": [2]},
"meta": {
"next_stage_prompt_len": 60,
"next_stage_generation_tokens": 26,
},
}
assert construct_next_stage_streaming_input_prompt(
third,
request,
max_model_len=4096,
previous_condition_len=12,
previous_condition_seq=1,
condition_seq=2,
recompute_previous_chunks=1,
recompute_on_capacity=True,
)
assert request.num_prompt_tokens == 75
assert request.num_computed_tokens == 0
assert third["ids"]["streaming_prompt_previous_codes"] == [201, 202, 203]
assert third["meta"]["streaming_condition_seq"] == 2
request._all_token_ids.extend([301])
request._output_token_ids.extend([301])
request.num_computed_tokens = 76
fourth = {
"ids": {"prompt": [3]},
"meta": {
"next_stage_prompt_len": 8,
"next_stage_generation_tokens": 26,
"streaming_condition_seq": 3,
},
}
assert not construct_next_stage_streaming_input_prompt(
fourth,
request,
max_model_len=4096,
previous_condition_len=60,
previous_condition_seq=2,
condition_seq=3,
recompute_previous_chunks=1,
recompute_on_capacity=True,
)
assert request.num_computed_tokens == 76
assert request.num_prompt_tokens == 84
assert fourth["meta"]["streaming_prompt_recompute"] is False
def test_capacity_rollover_requires_explicit_window_contract(mocker: MockerFixture) -> None:
request = _streaming_request(mocker, num_computed_tokens=4064)
payload = {
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 10,
"next_stage_generation_tokens": 26,
},
}
with pytest.raises(ValueError, match="requires window_size=1"):
construct_next_stage_streaming_input_prompt(
payload,
request,
max_model_len=4096,
previous_condition_len=10,
previous_condition_seq=0,
condition_seq=1,
)
@pytest.mark.parametrize("next_stage_prompt_len", [0, -1])
def test_capacity_managed_streaming_prompt_rejects_nonpositive_length(
mocker: MockerFixture, next_stage_prompt_len: int
) -> None:
request = _streaming_request(mocker, num_computed_tokens=0)
payload = {
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": next_stage_prompt_len,
"next_stage_generation_tokens": 26,
},
}
with pytest.raises(ValueError, match="positive next_stage_prompt_len"):
construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)
def test_capacity_managed_streaming_prompt_rejects_fresh_overflow(mocker: MockerFixture) -> None:
request = _streaming_request(mocker, num_computed_tokens=0)
payload = {
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 4071,
"next_stage_generation_tokens": 26,
},
}
with pytest.raises(ValueError, match="exceeds max_model_len"):
construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)
@pytest.fixture
def build_adapter(monkeypatch, mocker: MockerFixture):
def _build(
*,
stage_id: int = 1,
model_mode: str = "ar",
max_num_seqs: int = 2,
max_model_len: int = 0,
tts_max_model_len: int = 0,
tts_attention_type: str = "full_attention",
flat_tts_config: bool = False,
active_stream_window: int = 0,
connector_extra: dict | None = None,
):
connector = mocker.MagicMock()
connector.stage_id = stage_id
connector.config = {"extra": connector_extra or {}}
connector.get.return_value = None
connector.put.return_value = (True, 1, {})
def _fake_base_init(self, config):
self.config = config
self._pending_load_reqs = deque()
self._finished_load_reqs = set()
self._cancelled_load_reqs = set()
self._pending_save_reqs = deque()
self._finished_save_reqs = set()
self._send_failures = {}
self._send_failure_lock = threading.Lock()
self._receive_failures = {}
self._receive_failure_lock = threading.Lock()
self.stop_event = threading.Event()
self._recv_cond = threading.Condition()
self._save_cond = threading.Condition()
monkeypatch.setattr(OmniTransferAdapterBase, "__init__", _fake_base_init)
monkeypatch.setattr(
OmniChunkTransferAdapter,
"create_connector",
classmethod(lambda cls, _model_config: connector),
)
hf_config = SimpleNamespace(
model_type="conditional_chattts",
max_position_embeddings=tts_max_model_len,
attention_type=tts_attention_type,
)
if not flat_tts_config:
hf_config = SimpleNamespace(
tts_config=SimpleNamespace(
model_type="minicpmtts",
max_position_embeddings=tts_max_model_len,
attention_type=tts_attention_type,
)
)
model_config = SimpleNamespace(
worker_type=model_mode,
max_num_seqs=max_num_seqs,
max_model_len=max_model_len,
hf_config=hf_config,
hf_config_name="tts_config" if stage_id == 1 else None,
active_stream_window=active_stream_window,
stage_connector_config={
"name": "SharedMemoryConnector",
"extra": connector_extra or {},
},
)
scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs)
adapter = OmniChunkTransferAdapter(
SimpleNamespace(model_config=model_config, scheduler_config=scheduler_config)
)
return adapter, connector
return _build
def _dequeue_load_entry(adapter, request):
"""Register and take the entry as the background recv loop would."""
adapter.load_async(request)
entry = adapter._registered_load_entries[request.request_id]
try:
adapter._pending_load_reqs.remove(entry)
except ValueError:
pass
return entry
@pytest.mark.parametrize(
("attention_type", "expected_previous_chunks", "expected_capacity_trigger"),
[("full_attention", 1, True), ("sliding_recompute", 1, False), ("other", 0, False)],
)
def test_talker_attention_policy_controls_streaming_recompute(
build_adapter, attention_type, expected_previous_chunks, expected_capacity_trigger
):
adapter, _ = build_adapter(tts_attention_type=attention_type)
assert adapter._streaming_prompt_previous_chunks == expected_previous_chunks
assert adapter._streaming_prompt_recompute_on_capacity is expected_capacity_trigger
@pytest.mark.parametrize("stage_id", [0, 2])
def test_talker_streaming_policy_is_scoped_to_tts_stage(build_adapter, stage_id: int) -> None:
adapter, _ = build_adapter(stage_id=stage_id, tts_attention_type="full_attention")
assert adapter._streaming_prompt_previous_chunks == 0
assert adapter._streaming_prompt_recompute_on_capacity is False
@pytest.mark.parametrize(
("raw_cfg", "expected_name", "expected_extra"),
[
(None, "SharedMemoryConnector", {}),
(SimpleNamespace(name="YuanrongConnector", extra={"k": "v"}), "YuanrongConnector", {"k": "v"}),
],
)
def test_create_connector_config_parsing(monkeypatch, raw_cfg, expected_name, expected_extra):
captured = {}
def _fake_create(spec):
captured["spec"] = spec
return "ok"
monkeypatch.setattr(
"vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter"
".OmniConnectorFactory.create_connector",
_fake_create,
)
model_config = SimpleNamespace(stage_connector_config=raw_cfg) if raw_cfg is not None else SimpleNamespace()
connector = OmniChunkTransferAdapter.create_connector(model_config)
assert connector == "ok"
assert isinstance(captured["spec"], ConnectorSpec)
assert captured["spec"].name == expected_name
assert captured["spec"].extra == {**expected_extra, "stage_id": 0}
def test_load_poll(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="ar")
request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
adapter.load_async(request)
payload: OmniPayload = {
"codes": {"audio": [[1]]},
"hidden_states": {"output": torch.tensor([[2.0]])},
"meta": {"finished": torch.tensor(True, dtype=torch.bool)},
}
connector.get.return_value = (payload, 16)
adapter._poll_single_request(_dequeue_load_entry(adapter, request))
assert request.additional_information is None
assert "req-1" in adapter._finished_load_reqs
request.status = RequestStatus.WAITING_FOR_CHUNK
adapter.process_pending_chunks(
DummyWaitingQueue(),
[request],
scheduler_requests={request.request_id: request},
)
assert request.additional_information == payload
assert adapter.get_req_chunk["req-1"] == 1
assert "req-1" not in adapter._finished_load_reqs
def test_load_poll_uses_request_scoped_payload_sender_endpoint(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="ar")
request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
request.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
connector.get.return_value = None
adapter._poll_single_request(_dequeue_load_entry(adapter, request))
connector.get.assert_called_once_with(
"1",
"2",
"req-1_1_0",
{"source_host": "10.0.0.1", "source_port": 50051},
)
def test_load_poll_keeps_concurrent_payload_sender_endpoints_distinct(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="ar")
first = _req("req-1", RequestStatus.WAITING)
first.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
second = _req("req-2", RequestStatus.WAITING)
second.payload_sender_info = {"host": "10.0.0.2", "zmq_port": 51051}
connector.get.return_value = None
second_entry = _dequeue_load_entry(adapter, second)
first_entry = _dequeue_load_entry(adapter, first)
adapter._poll_single_request(second_entry)
adapter._poll_single_request(first_entry)
assert connector.get.call_args_list == [
call(
"1",
"2",
"req-2_1_0",
{"source_host": "10.0.0.2", "source_port": 51051},
),
call(
"1",
"2",
"req-1_1_0",
{"source_host": "10.0.0.1", "source_port": 50051},
),
]
def test_load_async_does_not_requeue_registered_inflight_request(build_adapter):
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
request = _req("req-inflight", RequestStatus.RUNNING, external_req_id="ext-inflight")
entry = _dequeue_load_entry(adapter, request)
# The recv loop owns the popped request until polling completes. A
# scheduler-side duplicate registration must not append a second copy.
adapter.load_async(request)
assert adapter._pending_load_reqs == deque()
connector.get.return_value = (
{
"ids": {"prompt": [1]},
"meta": {"finished": True},
},
1,
)
assert adapter._poll_single_request(entry) is True
# Successful polling releases the registration for the next chunk.
adapter.load_async(request)
assert [queued.request for queued in adapter._pending_load_reqs] == [request]
def test_cleanup_then_reregister_drops_stale_queued_entry(build_adapter):
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
request = _req("req-resume", RequestStatus.WAITING, external_req_id="ext-resume")
adapter.load_async(request)
stale_entry = adapter._pending_load_reqs[0]
adapter.cleanup_receiver(request.request_id)
adapter.load_async(request)
fresh_entry = adapter._pending_load_reqs[-1]
assert len(adapter._pending_load_reqs) == 2
assert adapter._pending_load_reqs.popleft() is stale_entry
assert adapter._pending_load_reqs.popleft() is fresh_entry
connector.get.return_value = (
{
"ids": {"prompt": [1]},
"meta": {"finished": True},
},
1,
)
assert adapter._poll_single_request(stale_entry) is True
connector.get.assert_not_called()
assert adapter._registered_load_entries[request.request_id] is fresh_entry
assert adapter._poll_single_request(fresh_entry) is True
connector.get.assert_called_once()
assert request.request_id not in adapter._registered_load_entries
assert adapter.get_req_chunk[request.request_id] == 1
def test_load_poll_ar_requeues_explicitly_replaced_running_prompt(build_adapter):
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
request = _req("req-replace", RequestStatus.RUNNING, external_req_id="external-replace")
request.resumable = True
request.prompt_token_ids = [0] * 4
request._all_token_ids = [0] * 4
request._output_token_ids = []
request.num_prompt_tokens = 4
request.num_computed_tokens = 4
request.update_block_hashes = Mock()
adapter.get_req_chunk[request.request_id] = 1
adapter.requests_num_chunks_sent[request.external_req_id] = 4
adapter.request_ids_mapping[request.request_id] = request.external_req_id
connector.get.return_value = (
{
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 10,
"replace_streaming_prompt": True,
"finished": False,
},
},
1,
)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
request.status = RequestStatus.WAITING_FOR_CHUNK
running_queue = [request]
waiting_queue = DummyWaitingQueue()
adapter.process_pending_chunks(
waiting_queue,
running_queue,
scheduler_requests={request.request_id: request},
)
assert request.num_computed_tokens == 0
assert request.prompt_token_ids == [0] * 10
assert request.request_id in adapter.replaced_streaming_prompt_ids
assert running_queue == []
assert waiting_queue == [request]
assert request.status == RequestStatus.WAITING
adapter.requests_num_chunks_sent[request.external_req_id] = 9
adapter.postprocess_scheduler_output(
SimpleNamespace(
scheduled_new_reqs=[SimpleNamespace(req_id=request.request_id)],
scheduled_cached_reqs=SimpleNamespace(req_ids=[]),
)
)
assert request.request_id not in adapter.replaced_streaming_prompt_ids
assert request.external_req_id not in adapter.requests_num_chunks_sent
@pytest.mark.parametrize("flat_tts_config", [False, True])
def test_load_poll_ar_recomputes_native_duplex_prompt_over_capacity(build_adapter, flat_tts_config):
adapter, connector = build_adapter(
stage_id=1,
model_mode="ar",
max_model_len=8192,
tts_max_model_len=4096,
tts_attention_type="full_attention",
flat_tts_config=flat_tts_config,
)
request = _req("req-rollover", RequestStatus.RUNNING, external_req_id="external-rollover")
request.resumable = True
previous_codes = list(range(25))
request.prompt_token_ids = [0] * 4039
request._all_token_ids = [0] * 4039 + previous_codes
request._output_token_ids = previous_codes.copy()
request.num_prompt_tokens = 4039
request.num_computed_tokens = 4064
request.num_output_placeholders = 0
request.update_block_hashes = Mock()
adapter.get_req_chunk[request.request_id] = 1
adapter.request_ids_mapping[request.request_id] = request.external_req_id
adapter._streaming_condition_lengths[request.request_id] = 10
adapter._streaming_condition_seqs[request.request_id] = 0
connector.get.return_value = (
{
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 10,
"next_stage_generation_tokens": 26,
"finished": True,
},
},
1,
)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
request.status = RequestStatus.WAITING_FOR_CHUNK
running_queue = [request]
waiting_queue = DummyWaitingQueue()
adapter.process_pending_chunks(
waiting_queue,
running_queue,
scheduler_requests={request.request_id: request},
)
assert request.num_computed_tokens == 0
assert request.prompt_token_ids == [0] * 45
assert request.additional_information["ids"]["streaming_prompt_previous_codes"] == previous_codes
assert request.additional_information["meta"]["streaming_prompt_recompute"] is True
assert request.resumable is False
assert request.request_id in adapter.replaced_streaming_prompt_ids
assert adapter._max_model_len == 4096
assert running_queue == []
assert waiting_queue == [request]
def test_load_poll_ar_keeps_generic_prompt_extension_for_non_tts_stage(build_adapter) -> None:
adapter, connector = build_adapter(stage_id=2, model_mode="ar")
request = _req("req-generic-extension", RequestStatus.RUNNING)
request.resumable = True
request.prompt_token_ids = [1, 2, 3]
request._all_token_ids = [1, 2, 3, 7]
request._output_token_ids = [7]
request.num_prompt_tokens = 3
request.num_computed_tokens = 4
request.update_block_hashes = Mock()
adapter.get_req_chunk[request.request_id] = 1
connector.get.return_value = (
{
"ids": {"prompt": [1]},
"meta": {"next_stage_prompt_len": 2},
},
1,
)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
adapter.requests_with_ready_chunks.add(request.request_id)
adapter._apply_pending_ar_prompt_updates({request.request_id: request})
assert request.num_computed_tokens == 4
assert request.prompt_token_ids == [1, 2, 3, 7, 0, 0]
assert request._all_token_ids == [1, 2, 3, 7, 0, 0]
assert request.num_prompt_tokens == 6
def test_window_condition_sequence_ignores_control_only_chunks(build_adapter) -> None:
adapter, connector = build_adapter(
stage_id=1,
model_mode="ar",
tts_attention_type="sliding_recompute",
)
request = _req("req-condition-seq", RequestStatus.RUNNING)
def receive(payload):
connector.get.return_value = (payload, 1)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
adapter.requests_with_ready_chunks.add(request.request_id)
adapter._apply_pending_ar_prompt_updates({request.request_id: request})
adapter.requests_with_ready_chunks.discard(request.request_id)
return request.additional_information
first = receive(
{
"native_duplex": True,
"ids": {"prompt": [1]},
"meta": {"next_stage_prompt_len": 10},
}
)
assert first["meta"]["streaming_condition_seq"] == 0
assert adapter._streaming_condition_seqs[request.request_id] == 0
control = receive({"meta": {"is_segment_finished": True}})
assert "streaming_condition_seq" not in control["meta"]
assert adapter._streaming_condition_seqs[request.request_id] == 0
second = receive(
{
"native_duplex": True,
"ids": {"prompt": [2]},
"meta": {"next_stage_prompt_len": 12},
}
)
assert second["meta"]["streaming_condition_seq"] == 1
assert adapter._streaming_condition_lengths[request.request_id] == 12
def test_load_poll_ar_reports_invalid_capacity_metadata(build_adapter):
adapter, connector = build_adapter(stage_id=1, model_mode="ar", max_model_len=4096)
request = _req("req-invalid-capacity", RequestStatus.RUNNING)
request.resumable = True
adapter.get_req_chunk[request.request_id] = 1
connector.get.return_value = (
{
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 4096,
"next_stage_generation_tokens": 26,
"replace_streaming_prompt": True,
"finished": False,
},
},
1,
)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
request.status = RequestStatus.WAITING_FOR_CHUNK
running_queue = [request]
adapter.process_pending_chunks(
DummyWaitingQueue(),
running_queue,
scheduler_requests={request.request_id: request},
)
assert adapter.collect_failed_receive_request_ids() == {
request.request_id: "fresh streaming prompt plus generation reserve exceeds max_model_len: "
"prompt=4096, reserve=26, limit=4096"
}
assert request.request_id not in adapter._finished_load_reqs
def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter):
adapter, connector = build_adapter(stage_id=1, model_mode="generation")
request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor")
codes = torch.tensor([[1, 2], [3, 4]], dtype=torch.long)
payload: OmniPayload = {
"codes": {"audio": codes},
"meta": {
"left_context_size": 1,
"finished": torch.tensor(False, dtype=torch.bool),
},
}
connector.get.return_value = (payload, 16)
adapter._poll_single_request(_dequeue_load_entry(adapter, request))
assert request.prompt_token_ids == [0]
assert request.num_computed_tokens == 0
assert torch.equal(request.additional_information["codes"]["audio"], codes)
assert request.additional_information["meta"]["left_context_size"] == 1
assert "finished" not in request.additional_information["meta"]
assert "req-tensor" in adapter._finished_load_reqs
def test_load_poll_generation_empty_nonterminal_chunk_keeps_polling(build_adapter):
adapter, connector = build_adapter(stage_id=1, model_mode="generation")
request = _req("req-empty-tensor", RequestStatus.WAITING, external_req_id="external-empty")
empty_payload: OmniPayload = {
"codes": {"audio": torch.empty((4, 0), dtype=torch.long)},
"meta": {
"left_context_size": 0,
"finished": torch.tensor(False, dtype=torch.bool),
},
}
ready_payload: OmniPayload = {
"codes": {"audio": torch.tensor([[1, 2]], dtype=torch.long)},
"meta": {
"left_context_size": 0,
"finished": torch.tensor(False, dtype=torch.bool),
},
}
connector.get.side_effect = [(empty_payload, 16), (ready_payload, 16)]
entry = _dequeue_load_entry(adapter, request)
assert adapter._poll_single_request(entry) is False
assert request.request_id not in adapter._finished_load_reqs
assert request.request_id not in adapter.requests_with_ready_chunks
assert adapter.get_req_chunk[request.request_id] == 1
assert adapter._poll_single_request(entry) is True
assert request.request_id in adapter._finished_load_reqs
assert torch.equal(request.additional_information["codes"]["audio"], ready_payload["codes"]["audio"])
assert adapter.get_req_chunk[request.request_id] == 2
def test_save_async(build_adapter):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": False}
adapter.save_async(multimodal_output=None, request=request)
adapter.custom_process_next_stage_input_func = lambda **kwargs: {}
adapter.save_async(multimodal_output=None, request=request)
task = adapter._pending_save_reqs.popleft()
assert task["is_finished"] is False
def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_adapter):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-async", RequestStatus.WAITING, external_req_id="external-async")
request.num_computed_tokens = 10
request.num_output_placeholders = 2
adapter.save_async(multimodal_output=None, request=request)
assert adapter.requests_num_chunks_sent["external-async"] == 8
assert len(adapter._pending_save_reqs) == 1
def test_segment_boundary_starts_new_send_watermark_before_background_flush(build_adapter):
"""A queued boundary owns the end of its deduplication generation.
The next segment can start before the save thread sends the old boundary.
Sending that old task later must not erase the new segment's watermark.
"""
adapter, _ = build_adapter(stage_id=1)
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
request.resumable = True
request._omni_segment_generation = 1
request.num_computed_tokens = 0
adapter.requests_num_chunks_sent["ext-stream"] = 26
adapter.save_async(
multimodal_output=None,
request=request,
is_segment_finished=True,
confirmed_num_computed_tokens=26,
segment_generation=0,
)
assert len(adapter._pending_save_reqs) == 1
boundary_task = adapter._pending_save_reqs.popleft()
assert boundary_task["segment_generation"] == 0
assert "ext-stream" not in adapter.requests_num_chunks_sent
request.num_computed_tokens = 3
request._omni_segment_generation = 1
adapter.save_async(
multimodal_output=None,
request=request,
is_segment_finished=False,
)
assert len(adapter._pending_save_reqs) == 1
assert adapter.requests_num_chunks_sent["ext-stream"] == 3
adapter._send_single_request(boundary_task)
assert adapter.requests_num_chunks_sent["ext-stream"] == 3
def test_background_send_uses_enqueued_request_snapshot(build_adapter):
"""A queued segment must not observe later in-place request mutations."""
adapter, _ = build_adapter(stage_id=1)
request = Request(
request_id="req-stream",
prompt_token_ids=[1, 2],
sampling_params=SamplingParams(max_tokens=8),
pooling_params=None,
resumable=True,
)
request.external_req_id = "ext-stream"
ref_audio = torch.tensor([0.1, -0.1])
request.additional_information = {"codes": {"ref": ref_audio}, "meta": {"segment": "old"}}
request.append_output_token_ids([7])
seen_requests = []
def recording_processor(**kwargs):
queued = kwargs["request"]
seen_requests.append(
(
queued.additional_information["meta"]["segment"],
list(queued.prompt_token_ids),
list(queued.output_token_ids),
list(queued.all_token_ids),
queued.additional_information["codes"]["ref"] is ref_audio,
)
)
return OmniPayloadStruct()
adapter.custom_process_next_stage_input_func = recording_processor
adapter.save_async(
multimodal_output=None,
request=request,
is_segment_finished=True,
)
request.additional_information["meta"]["segment"] = "next"
request.prompt_token_ids.append(3)
request.append_output_token_ids([8])
request._omni_segment_generation = 1
adapter.save_async(
multimodal_output=None,
request=request,
is_segment_finished=True,
)
request.additional_information["meta"]["segment"] = "later"
request.prompt_token_ids.append(4)
request.append_output_token_ids([9])
assert len(adapter._pending_save_reqs) == 2
first_task = adapter._pending_save_reqs.popleft()
second_task = adapter._pending_save_reqs.popleft()
assert first_task["request"].additional_information is not request.additional_information
assert first_task["request"].additional_information["meta"] is not request.additional_information["meta"]
adapter._send_single_request(first_task)
adapter._send_single_request(second_task)
assert seen_requests == [
("old", [1, 2], [7], [1, 2, 7], True),
("next", [1, 2, 3], [7, 8], [1, 2, 7, 8], True),
]
def test_save_without_custom_processor_does_not_snapshot_request(build_adapter, mocker):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-direct", RequestStatus.RUNNING, external_req_id="ext-direct")
snapshot = mocker.patch.object(
adapter,
"_snapshot_processor_request",
side_effect=AssertionError("unexpected snapshot"),
)
adapter.save_async(multimodal_output=None, request=request)
snapshot.assert_not_called()
assert adapter._pending_save_reqs.popleft()["request"] is request
def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch):
"""A terminal stop is not a segment boundary (#5383), but the producer-side
processor must still receive the flush signal on the terminal chunk.
Passing only ``is_segment_finished`` starved processors of their final
accumulated payload once terminal stops stopped setting it (#5413: the
downstream stage got ``meta.finished`` with the tail data missing).
"""
adapter, connector = build_adapter(stage_id=0)
request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal")
seen_flush_flags = []
def recording_processor(**kwargs):
seen_flush_flags.append(kwargs["is_finished"])
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.long)),
)
adapter.custom_process_next_stage_input_func = recording_processor
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
adapter._send_single_request(
{"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": False}
)
assert seen_flush_flags == [True]
sent_payload = connector.put.call_args.kwargs["data"]
assert bool(sent_payload.meta.finished.item()) is True
assert bool(sent_payload.meta.is_segment_finished.item()) is False
def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch):
"""Producer may return a struct with ``meta=None`` (e.g. payload that
carries only ``embed`` or ``codes``). The sender's ``meta is not None``
guard handles this without AttributeError; ``finished_flag`` is None and
the cleanup path is not triggered.
"""
adapter, _ = build_adapter(stage_id=1)
request = _req("req-no-meta", RequestStatus.WAITING, external_req_id="ext-no-meta")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
codes=CodesStruct(audio=torch.tensor([1, 2], dtype=torch.long)),
)
cleanup_calls = []
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
adapter._send_single_request(
{"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
)
assert cleanup_calls == [] # no terminal cleanup; meta.finished is false
def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatch):
"""Pin the contract: an explicitly empty ``OmniPayloadStruct()`` passes
the ``payload_data is None`` check and gets sent. To skip a chunk, the
producer must return ``None``, not an empty struct. (Filtering empty
structs at the adapter would require introspecting all struct fields on
every send and was rejected for cost vs. value.)
"""
adapter, connector = build_adapter(stage_id=1)
request = _req("req-empty", RequestStatus.WAITING, external_req_id="ext-empty")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
adapter._send_single_request(
{"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
)
assert connector.put.called
sent_payload = connector.put.call_args.kwargs["data"]
assert isinstance(sent_payload, OmniPayloadStruct)
assert sent_payload.meta.finished.item() is False
assert sent_payload.meta.is_segment_finished.item() is False
def test_send_single_request_struct_preserves_segment_finished(build_adapter, monkeypatch):
adapter, connector = build_adapter(stage_id=1)
request = _req("req-segment", RequestStatus.WAITING, external_req_id="ext-segment")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
adapter._send_single_request(
{"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
)
sent_payload = connector.put.call_args.kwargs["data"]
assert sent_payload.meta.finished.item() is False
assert sent_payload.meta.is_segment_finished.item() is True
def test_send_single_request_respects_processor_receiver_boundary(build_adapter, monkeypatch):
adapter, connector = build_adapter(stage_id=1)
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
meta=MetaStruct(is_segment_finished=torch.tensor(False, dtype=torch.bool))
)
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
adapter._send_single_request(
{"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
)
sent_payload = connector.put.call_args.kwargs["data"]
assert sent_payload.meta.is_segment_finished.item() is False
def test_send_single_request_personaplex_pending_frame_is_not_segment_boundary(
build_adapter,
):
from vllm_omni.model_executor.stage_input_processors.personaplex import (
talker2code2wav_async_chunk,
)
adapter, connector = build_adapter(
stage_id=0,
connector_extra={
"initial_codec_chunk_frames": 1,
"codec_chunk_frames": 5,
},
)
request = _req(
"req-personaplex",
RequestStatus.WAITING,
external_req_id="ext-personaplex",
)
request.resumable = True
request.additional_information = {
"codes": {
"audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
}
}
adapter.custom_process_next_stage_input_func = talker2code2wav_async_chunk
adapter._send_single_request(
{
"multimodal_output": None,
"request": request,
"is_finished": False,
"is_segment_finished": True,
}
)
sent_payload = connector.put.call_args.kwargs["data"]
assert sent_payload.codes is None
assert sent_payload.meta.finished.item() is False
assert sent_payload.meta.is_segment_finished.item() is False
def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter):
from vllm_omni.model_executor.stage_input_processors.personaplex import (
talker2code2wav_async_chunk,
)
adapter, _ = build_adapter(
stage_id=0,
connector_extra={
"initial_codec_chunk_frames": 1,
"codec_chunk_frames": 5,
},
)
request = _req(
"req-personaplex-first",
RequestStatus.WAITING,
external_req_id="ext-personaplex-reused",
)
request.resumable = True
request.additional_information = {
"codes": {
"audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
}
}
first = talker2code2wav_async_chunk(
adapter,
multimodal_output=None,
request=request,
is_finished=True,
)
assert first is not None
assert first.codes is None
adapter.cleanup_sender(request.external_req_id)
replacement = _req(
"req-personaplex-replacement",
RequestStatus.WAITING,
external_req_id=request.external_req_id,
)
replacement.resumable = True
replacement.additional_information = {
"codes": {
"audio": torch.arange(8, 16, dtype=torch.long).reshape(1, 8),
}
}
second = talker2code2wav_async_chunk(
adapter,
multimodal_output=None,
request=replacement,
is_finished=True,
)
assert second is not None
assert second.codes is None
def test_save_async_skips_stale_resumable_chunk_within_segment(build_adapter):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
request.resumable = True
request.num_computed_tokens = 0
adapter.requests_num_chunks_sent["ext-stream"] = 111
adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
assert len(adapter._pending_save_reqs) == 0
assert adapter.requests_num_chunks_sent["ext-stream"] == 111
adapter.requests_num_chunks_sent.pop("ext-stream")
adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
assert len(adapter._pending_save_reqs) == 1
assert adapter.requests_num_chunks_sent["ext-stream"] == 0
def test_save_async_drops_late_previous_segment_after_boundary_reset(build_adapter):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
request.resumable = True
request._omni_segment_generation = 0
adapter.requests_num_chunks_sent["ext-stream"] = 111
adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
late = _req("req-stream-late", RequestStatus.WAITING, external_req_id="ext-stream")
late.resumable = True
late.num_computed_tokens = 111
late._omni_segment_generation = 0
adapter.save_async(multimodal_output=None, request=late, is_segment_finished=False)
next_segment = _req("req-stream-next", RequestStatus.WAITING, external_req_id="ext-stream")
next_segment.resumable = True
next_segment.num_computed_tokens = 3
next_segment._omni_segment_generation = 1
adapter.save_async(multimodal_output=None, request=next_segment, is_segment_finished=False)
assert len(adapter._pending_save_reqs) == 2
assert adapter.requests_num_chunks_sent["ext-stream"] == 3
assert adapter._pending_save_reqs[-1]["request"] is next_segment
def test_send_single_request_preserves_queued_preboundary_chunk(build_adapter):
adapter, connector = build_adapter(stage_id=1)
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(meta=MetaStruct())
request = _req("req-queued-old", RequestStatus.WAITING, external_req_id="ext-queued-old")
request.resumable = True
request._omni_segment_generation = 0
adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
old_task = adapter._pending_save_reqs.popleft()
adapter._send_single_request(old_task)
connector.put.assert_called_once()
boundary_task = adapter._pending_save_reqs.popleft()
adapter._send_single_request(boundary_task)
assert connector.put.call_count == 2
def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool))
)
cleanup_calls = []
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
adapter._send_single_request(
{"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
)
assert len(cleanup_calls) == 1
args, _ = cleanup_calls[0]
assert args[0] == "req-finished"
assert args[1] == "ext-finished"
def test_old_turn_terminal_send_preserves_new_turn_state_and_sender(build_adapter):
adapter, connector = build_adapter(
stage_id=1,
model_mode="ar",
max_model_len=4096,
tts_attention_type="sliding_recompute",
)
request = _req("req-live", RequestStatus.RUNNING, external_req_id="ext-live")
request.resumable = True
request._all_token_ids = []
request._output_token_ids = []
request.update_block_hashes = Mock()
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
meta=MetaStruct(
finished=torch.tensor(True, dtype=torch.bool),
codec_streaming=True,
)
)
put_started = threading.Event()
release_put = threading.Event()
put_keys = []
def blocking_first_put(*_args, **kwargs):
put_keys.append(kwargs["put_key"])
if len(put_keys) == 1:
put_started.set()
assert release_put.wait(timeout=5)
return True, 1, {}
connector.put.side_effect = blocking_first_put
adapter.save_async(multimodal_output=None, request=request)
first_task = adapter._pending_save_reqs.popleft()
first_sender_token = first_task["sender_token"]
send_thread = threading.Thread(target=lambda: adapter._send_single_request(first_task))
send_thread.start()
assert put_started.wait(timeout=5)
entry = _dequeue_load_entry(adapter, request)
connector.get.return_value = (
{
"ids": {"prompt": [1]},
"meta": {
"replace_streaming_prompt": True,
"next_stage_prompt_len": 10,
"next_stage_generation_tokens": 26,
},
},
1,
)
assert adapter._poll_single_request(entry) is True
adapter.requests_with_ready_chunks.add(request.request_id)
adapter._apply_pending_ar_prompt_updates({request.request_id: request})
assert adapter._streaming_condition_seqs[request.request_id] == 0
request.num_computed_tokens = 1
adapter.save_async(multimodal_output=None, request=request)
second_task = adapter._pending_save_reqs.popleft()
assert second_task["sender_token"] is first_sender_token
release_put.set()
send_thread.join(timeout=5)
assert not send_thread.is_alive()
assert adapter._streaming_condition_seqs[request.request_id] == 0
assert adapter._sender_tokens[request.external_req_id] is first_sender_token
assert first_sender_token.cancelled is False
adapter._send_single_request(second_task)
assert put_keys == ["ext-live_1_0", "ext-live_1_1"]
def test_load_poll_non_ar_merges_into_existing_additional_information(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="diffusion")
request = _req("req-non-ar", RequestStatus.WAITING, external_req_id="ext-non-ar")
request.additional_information = {
"hidden_states": {"output": torch.tensor([[1.0]])},
"ids": {"prompt": [11, 12]},
"meta": {"finished": torch.tensor(False, dtype=torch.bool), "step": 1},
}
request.num_computed_tokens = 9
payload: OmniPayload = {
"hidden_states": {"output": torch.tensor([[2.0]])},
"ids": {"all": [21, 22]},
"codes": {"audio": torch.tensor([7, 8], dtype=torch.long)},
"meta": {"finished": torch.tensor(True, dtype=torch.bool), "phase": "decode"},
"kv_metadata": {"foo": "bar"},
}
connector.get.return_value = (payload, 8)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
assert request.prompt_token_ids == [7, 8]
assert request.num_computed_tokens == 0
assert torch.equal(
request.additional_information["hidden_states"]["output"],
torch.tensor([[2.0]]),
)
assert request.additional_information["ids"]["prompt"] == [11, 12]
assert request.additional_information["ids"]["all"] == [21, 22]
# non-ar merge path intentionally doesn't overwrite meta.finished.
assert request.additional_information["meta"]["finished"].item() is False
assert request.additional_information["meta"]["phase"] == "decode"
assert request.additional_information["kv_metadata"] == {"foo": "bar"}
assert "req-non-ar" in adapter._finished_load_reqs
assert "req-non-ar" in adapter.upstream_exhausted_requests
def test_load_poll_generation_segment_marker_replaces_previous_chunk(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="generation")
request = _req("req-marker", RequestStatus.WAITING, external_req_id="external-marker")
request.additional_information = {
"codes": {"audio": torch.tensor([1, 2])},
"meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
}
connector.get.return_value = (
{
"codes": {
"audio": torch.tensor([7, 8], dtype=torch.long),
"ref": torch.tensor([0.1, -0.1]),
},
"meta": {
"finished": torch.tensor(False, dtype=torch.bool),
"is_segment_finished": torch.tensor(True, dtype=torch.bool),
"request_finished": torch.tensor(False, dtype=torch.bool),
"replace_runtime_additional_information": True,
},
},
1,
)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
assert request.prompt_token_ids == [7, 8]
assert "audio" not in request.additional_information["codes"]
torch.testing.assert_close(
request.additional_information["codes"]["ref"],
torch.tensor([0.1, -0.1]),
)
assert not {"cache_epoch", "chunk_seq", "last_chunk"}.intersection(request.additional_information["meta"])
assert request.request_id in adapter.segment_finished_requests
def test_load_poll_generation_empty_replacement_snapshot_is_ready(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="generation")
request = _req("req-empty-marker", RequestStatus.WAITING, external_req_id="external-empty-marker")
request.additional_information = {
"codes": {"audio": torch.tensor([1, 2])},
"meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
}
connector.get.return_value = (
{
"meta": {
"is_segment_finished": torch.tensor(True, dtype=torch.bool),
"replace_runtime_additional_information": True,
}
},
1,
)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
assert request.prompt_token_ids == [0]
assert "codes" not in request.additional_information
assert request.additional_information["meta"]["replace_runtime_additional_information"] is True
assert request.request_id in adapter.segment_finished_requests
assert request.request_id in adapter._finished_load_reqs
def test_load_poll_generation_without_snapshot_marker_keeps_incremental_state(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="generation")
request = _req("req-incremental", RequestStatus.WAITING, external_req_id="external-incremental")
request.additional_information = {
"codes": {"audio": torch.tensor([1, 2])},
"meta": {"cache_epoch": 3, "chunk_seq": 2},
}
connector.get.return_value = (
{
"meta": {
"finished": torch.tensor(False, dtype=torch.bool),
"phase": "decode",
}
},
1,
)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is False
assert torch.equal(request.additional_information["codes"]["audio"], torch.tensor([1, 2]))
assert request.additional_information["meta"]["cache_epoch"] == 3
assert request.additional_information["meta"]["chunk_seq"] == 2
assert request.additional_information["meta"]["phase"] == "decode"
def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="ar")
request = _req("req-merged", RequestStatus.WAITING, external_req_id="ext-merged")
request.additional_information = {
"hidden_states": {"output": torch.tensor([[1.0]])},
"ids": {"prompt": [11, 12]},
"meta": {"finished": torch.tensor(False, dtype=torch.bool)},
}
adapter.request_ids_mapping["req-merged"] = "ext-merged"
payload: OmniPayload = {
"hidden_states": {"output": torch.tensor([[2.0]])},
"meta": {"finished": torch.tensor(True, dtype=torch.bool)},
}
connector.get.return_value = (payload, 8)
adapter._poll_single_request(_dequeue_load_entry(adapter, request))
# The recv thread records the latest payload; scheduler-side queue
# restoration publishes it onto the Request.
assert torch.equal(request.additional_information["hidden_states"]["output"], torch.tensor([[1.0]]))
request.status = RequestStatus.WAITING_FOR_CHUNK
adapter.process_pending_chunks(
DummyWaitingQueue(),
[request],
scheduler_requests={request.request_id: request},
)
assert request.additional_information == payload
assert request.additional_information["meta"]["finished"].item() is True
def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="generation")
request = _req("req-later-chunk", RequestStatus.WAITING, external_req_id="ext-later-chunk")
request.prefill_stats = PrefillStats()
adapter.request_ids_mapping[request.request_id] = request.external_req_id
connector.get.return_value = (
{
"codes": {"audio": torch.tensor([7, 8], dtype=torch.long)},
"meta": {"finished": torch.tensor(False, dtype=torch.bool)},
},
8,
)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
OmniGenerationScheduler._record_prefill_stats(request)
first_chunk_stats = request.prefill_stats
request.prefill_stats = None
connector.get.return_value = (
{
"codes": {"audio": torch.tensor([9, 10, 11], dtype=torch.long)},
"meta": {"finished": torch.tensor(True, dtype=torch.bool)},
},
8,
)
assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
assert isinstance(request.prefill_stats, PrefillStats)
OmniGenerationScheduler._record_prefill_stats(request)
second_chunk_stats = request.prefill_stats
assert first_chunk_stats is not None
assert first_chunk_stats.num_prompt_tokens == 2
assert second_chunk_stats is not None
assert second_chunk_stats.num_prompt_tokens == 3
prompt_token_stats = PromptTokenStats()
prompt_token_stats.update_from_output(first_chunk_stats)
prompt_token_stats.update_from_output(second_chunk_stats)
assert prompt_token_stats.total == 5
assert prompt_token_stats.computed == 5
def test_sender_only_adapter_does_not_park_or_clear_requests(build_adapter):
adapter, _ = build_adapter(stage_id=1, connector_extra={"role": "sender"})
request = _req("req-1", RequestStatus.WAITING)
request.additional_information = {"tts_token_ids": torch.tensor([1])}
waiting_queue = DummyWaitingQueue([request])
running_queue = []
adapter.load_async(request)
adapter.process_pending_chunks(
waiting_queue,
running_queue,
scheduler_requests={request.request_id: request},
)
assert waiting_queue == [request]
assert request.status == RequestStatus.WAITING
assert request.additional_information["tts_token_ids"].item() == 1
assert adapter._pending_load_reqs == deque()
def test_process_and_restore_queues(build_adapter):
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
waiting_req = _req("w1", RequestStatus.WAITING)
running_req = _req("r1", RequestStatus.RUNNING)
waiting_queue = DummyWaitingQueue([waiting_req])
running_queue = [running_req]
scheduler_requests = {waiting_req.request_id: waiting_req, running_req.request_id: running_req}
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK
assert running_req.status == RequestStatus.WAITING_FOR_CHUNK
assert waiting_queue == []
assert running_queue == []
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert waiting_queue == [waiting_req]
assert running_queue == [running_req]
assert adapter.waiting_for_chunk_waiting_requests == deque()
assert adapter.waiting_for_chunk_running_requests == deque()
def test_fifo_promotion(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
waiting_queue = DummyWaitingQueue(reqs)
running_queue = []
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-1", "req-2"]
assert waiting_queue == []
assert reqs[0].status == RequestStatus.WAITING_FOR_CHUNK
assert reqs[1].status == RequestStatus.WAITING_FOR_CHUNK
assert reqs[2].status == RequestStatus.WAITING
# Ordinary completion calls cleanup_receiver(), not finish_requests()
# (the abort path) -- see test_omni_ar_scheduler_free_request_cleanup.py.
# Eviction is deferred to postprocess_scheduler_output in the runtime path
# (commit c4d95fd9 -- otherwise the terminal chunk deadlocks at c=8 K=2).
# Simulate the restore step here so promotion can pick up the freed slot.
adapter.cleanup_receiver("req-1")
adapter.restore_queues(waiting_queue, running_queue)
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-2", "req-3"]
# Promotion + chunk processing happen in the same call, so req-3 is
# already WAITING_FOR_CHUNK by the time we check.
assert reqs[2].status == RequestStatus.WAITING_FOR_CHUNK
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)
non_active = _req("req-non-active", RequestStatus.WAITING)
waiting_queue = DummyWaitingQueue([active, non_active])
running_queue = []
adapter.process_pending_chunks(waiting_queue, running_queue)
assert waiting_queue == []
assert active.status == RequestStatus.WAITING_FOR_CHUNK
assert non_active.status == RequestStatus.WAITING
assert [entry.request for entry in adapter._pending_load_reqs] == [active]
assert list(adapter.waiting_for_chunk_waiting_requests) == [active, non_active]
def test_legacy_k0(build_adapter):
adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=0)
waiting_req = _req("waiting", RequestStatus.WAITING)
running_req_1 = _req("running-1", RequestStatus.RUNNING)
running_req_2 = _req("running-2", RequestStatus.RUNNING)
waiting_queue = DummyWaitingQueue([waiting_req])
running_queue = [running_req_1, running_req_2]
adapter.requests_with_ready_chunks.update({"running-1", "running-2"})
adapter.process_pending_chunks(waiting_queue, running_queue)
assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK
assert adapter.waiting_for_chunk_waiting_requests == deque([waiting_req])
assert running_queue == [running_req_1]
assert waiting_queue == [running_req_2]
assert running_req_2.status == RequestStatus.PREEMPTED
assert adapter._active_streams == {}
def test_finished_releases_slot(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=1, active_stream_window=1)
req_1 = _req("req-1", RequestStatus.WAITING)
req_2 = _req("req-2", RequestStatus.WAITING)
waiting_queue = DummyWaitingQueue([req_1, req_2])
running_queue = []
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-1"]
assert waiting_queue == []
# req-1 finishes ordinarily (see test_fifo_promotion for why this is
# cleanup_receiver(), not finish_requests()). Eviction is deferred to
# postprocess_scheduler_output in the runtime path (commit c4d95fd9);
# simulate the restore step so promotion can pick up the freed slot.
adapter.cleanup_receiver("req-1")
adapter.restore_queues(waiting_queue, running_queue)
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-2"]
# Promotion + chunk processing happen in the same call, so req-2 is
# already WAITING_FOR_CHUNK by the time we check.
assert req_2.status == RequestStatus.WAITING_FOR_CHUNK
def test_postprocess_scheduler_output(build_adapter):
adapter, _ = build_adapter()
adapter.requests_with_ready_chunks = {"new-ready", "cached-ready", "leftover"}
scheduler_output = SimpleNamespace(
scheduled_new_reqs=[SimpleNamespace(req_id="new-ready")],
scheduled_cached_reqs=SimpleNamespace(req_ids=["cached-ready", "missing"]),
)
requests = {"cached-ready": SimpleNamespace(additional_information={"k": "v"})}
adapter.postprocess_scheduler_output(scheduler_output, requests)
cached_info = scheduler_output.scheduled_cached_reqs.additional_information
assert cached_info["cached-ready"] == {"k": "v"}
assert cached_info["missing"] is None
assert adapter.requests_with_ready_chunks == {"leftover"}
@pytest.mark.parametrize("model_mode", ["ar", "generation"])
def test_active_stream_window_stalls_lone_running_request_after_upstream_finishes(build_adapter, model_mode):
"""Regression test for vllm-project/vllm-omni#5349.
upstream_exhausted_requests means the upstream sent its terminal chunk,
not that this stage's own generation is done. A downstream stage can
still have a long decode ahead after its upstream finishes. Evicting on
that signal and then permanently denying re-admission wedges the
request out of running_queue forever, with no error.
requests_with_ready_chunks is set up front so the request stays in
running_queue this tick instead of legitimately parking in
waiting_for_chunk_running_requests (a separate mechanism this test must
not conflate with the bug).
Parametrized over both worker types: the real issue hit an "ar" stage
(Qwen3-Omni's talker); must not regress "generation" (Qwen3-TTS's
Code2Wav), which this feature was designed for.
"""
adapter, _ = build_adapter(stage_id=1, model_mode=model_mode, max_num_seqs=2, active_stream_window=2)
running_req = _req("req-1", RequestStatus.RUNNING)
running_queue = [running_req]
waiting_queue = DummyWaitingQueue([])
adapter._active_streams["req-1"] = running_req
adapter.requests_with_ready_chunks.add("req-1")
# Upstream terminal chunk arrives; this stage hasn't finished.
adapter.upstream_exhausted_requests.add("req-1")
scheduler_output = SimpleNamespace(
scheduled_new_reqs=[],
scheduled_cached_reqs=SimpleNamespace(req_ids=["req-1"]),
)
adapter.postprocess_scheduler_output(scheduler_output)
assert not running_req.is_finished() # sanity: not done
# No competition for the slot; req-1 must stay schedulable.
adapter.process_pending_chunks(waiting_queue, running_queue)
assert running_queue == [running_req], "req-1 must still be schedulable -- losing it here is #5349"
assert not adapter.waiting_for_chunk_running_requests
assert running_req not in adapter._held_non_active
assert list(adapter._active_streams) == ["req-1"]
def test_cleanup_receiver_releases_multiple_slots_in_sequence(build_adapter):
"""Companion to test_fifo_promotion: after EACH of the K active streams
finishes (not just one), promotion of a new K-sized batch must still
work -- the window must not stay exhausted by stale entries. See
test_omni_ar_scheduler_free_request_cleanup.py for the scheduler-level
proof that ordinary completion calls cleanup_receiver().
"""
adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
waiting_queue = DummyWaitingQueue(reqs)
running_queue = []
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-1", "req-2"]
adapter.cleanup_receiver("req-1")
adapter.cleanup_receiver("req-2")
# Eviction is deferred to postprocess_scheduler_output in the runtime
# path; simulate the restore step so promotion can pick up the freed
# slots (see test_fifo_promotion).
adapter.restore_queues(waiting_queue, running_queue)
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-3", "req-4"], (
"both freed slots must be reusable after ordinary completion"
)
# ---------------------------------------------------------------
# Cleanup tests
# ---------------------------------------------------------------
def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"):
"""Fill every per-request structure so cleanup can be verified."""
adapter.upstream_exhausted_requests.add(req_id)
adapter._active_streams[req_id] = SimpleNamespace(request_id=req_id)
adapter.get_req_chunk[req_id] = 3
adapter.requests_with_ready_chunks.add(req_id)
adapter.request_ids_mapping[req_id] = ext_id
entry = SimpleNamespace(request_id=req_id)
adapter._pending_load_reqs.append(entry)
adapter._registered_load_entries[req_id] = entry
adapter._finished_load_reqs.add(req_id)
adapter._pending_ar_prompt_updates[req_id] = (
SimpleNamespace(request_id=req_id),
{},
False,
False,
False,
)
adapter._streaming_condition_lengths[req_id] = 10
adapter._streaming_condition_seqs[req_id] = 2
adapter.put_req_chunk[ext_id] = 5
adapter.request_payload[ext_id] = {"hidden": [1, 2]}
adapter.code_prompt_token_ids[ext_id] = [[10, 20]]
def test_cleanup_clears_all_state(build_adapter):
"""After cleanup, no per-request key should remain in any dict/set."""
adapter, _ = build_adapter(stage_id=1)
req_id, ext_id = "req-1", "ext-1"
_populate_adapter_state(adapter, req_id, ext_id)
adapter.cleanup(req_id, ext_id)
assert req_id not in adapter.upstream_exhausted_requests
assert req_id not in adapter._active_streams
assert req_id not in adapter.get_req_chunk
assert req_id not in adapter.requests_with_ready_chunks
assert req_id not in adapter.request_ids_mapping
assert req_id in adapter._cancelled_load_reqs
assert req_id not in adapter._finished_load_reqs
assert req_id not in adapter._pending_ar_prompt_updates
assert req_id not in adapter._streaming_condition_lengths
assert req_id not in adapter._streaming_condition_seqs
assert req_id not in adapter._registered_load_entries
assert ext_id not in adapter.put_req_chunk
assert ext_id not in adapter.request_payload
assert ext_id not in adapter.code_prompt_token_ids
def test_cleanup_and_reregister_discards_inflight_old_segment_chunk(build_adapter):
"""A late old-segment chunk cannot overwrite a resumed registration."""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
request = _req("req-late", RequestStatus.RUNNING, external_req_id="ext-late")
request.resumable = True
entry = _dequeue_load_entry(adapter, request)
get_started = threading.Event()
release_get = threading.Event()
def blocking_get(*_args):
get_started.set()
assert release_get.wait(timeout=5)
return (
{
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": 10,
},
},
1,
)
connector.get.side_effect = blocking_get
poll_results = []
poll_thread = threading.Thread(target=lambda: poll_results.append(adapter._poll_single_request(entry)))
poll_thread.start()
assert get_started.wait(timeout=5)
adapter.cleanup_receiver(request.request_id)
fresh_entry = _dequeue_load_entry(adapter, request)
release_get.set()
poll_thread.join(timeout=5)
assert not poll_thread.is_alive()
assert poll_results == [True]
assert request.additional_information is None
assert request.request_id not in adapter._pending_ar_prompt_updates
assert request.request_id not in adapter._streaming_condition_seqs
assert adapter.get_req_chunk[request.request_id] == 0
assert adapter._registered_load_entries[request.request_id] is fresh_entry
connector.get.side_effect = None
connector.get.return_value = (
{
"ids": {"prompt": [2]},
"meta": {"finished": True},
},
1,
)
assert adapter._poll_single_request(fresh_entry) is True
assert adapter.get_req_chunk[request.request_id] == 1
assert request.request_id not in adapter._registered_load_entries
def test_cleanup_infers_external_id(build_adapter):
"""When external_req_id is None, cleanup should look it up from the mapping."""
adapter, _ = build_adapter(stage_id=1)
req_id, ext_id = "req-2", "ext-2"
_populate_adapter_state(adapter, req_id, ext_id)
adapter.cleanup(req_id)
assert ext_id not in adapter.put_req_chunk
assert ext_id not in adapter.request_payload
def test_cleanup_idempotent(build_adapter):
"""Calling cleanup multiple times for the same (or nonexistent) request must not raise."""
adapter, _ = build_adapter(stage_id=1)
try:
adapter.cleanup("nonexistent")
adapter.cleanup("nonexistent")
except Exception as e:
pytest.fail(f"cleanup should be idempotent: {e}")
req_id, ext_id = "req-3", "ext-3"
_populate_adapter_state(adapter, req_id, ext_id)
adapter.cleanup(req_id, ext_id)
try:
adapter.cleanup(req_id, ext_id)
except Exception as e:
pytest.fail(f"second cleanup should be idempotent: {e}")
def test_cleanup_request_id_reuse_not_polluted(build_adapter):
"""After cleanup, reusing the same request_id must not be treated as finished."""
adapter, _ = build_adapter(stage_id=1)
req_id, ext_id = "req-reuse", "ext-reuse"
_populate_adapter_state(adapter, req_id, ext_id)
adapter.cleanup(req_id, ext_id)
assert req_id not in adapter.upstream_exhausted_requests
assert req_id not in adapter.get_req_chunk
def test_cleanup_preserves_pending_save(build_adapter):
"""Cleanup must NOT remove _pending_save_reqs to avoid losing unsent chunks."""
adapter, _ = build_adapter(stage_id=1)
req_id, ext_id = "req-4", "ext-4"
_populate_adapter_state(adapter, req_id, ext_id)
pending_task = {"put_key": f"{ext_id}_1_0", "data": {"x": 1}}
adapter._pending_save_reqs.append(pending_task)
adapter.cleanup(req_id, ext_id)
assert len(adapter._pending_save_reqs) == 1
def test_abort_invalidates_queued_sender_task_before_external_id_reuse(build_adapter):
adapter, connector = build_adapter(stage_id=1)
request = _req("req-old", RequestStatus.WAITING, external_req_id="ext-reused")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
)
adapter.save_async(multimodal_output=None, request=request)
stale_task = adapter._pending_save_reqs.popleft()
adapter.cleanup(request.request_id, request.external_req_id)
adapter._send_single_request(stale_task)
connector.put.assert_not_called()
assert request.external_req_id not in adapter.put_req_chunk
replacement = _req("req-new", RequestStatus.WAITING, external_req_id=request.external_req_id)
adapter.save_async(multimodal_output=None, request=replacement)
current_task = adapter._pending_save_reqs.popleft()
adapter._send_single_request(current_task)
connector.put.assert_called_once()
assert connector.put.call_args.kwargs["put_key"] == "ext-reused_1_0"
def test_finish_requests_does_not_wait_for_inflight_send(build_adapter):
adapter, connector = build_adapter(stage_id=1)
first = _req("req-first", RequestStatus.WAITING, external_req_id="ext-first")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
put_started = threading.Event()
release_put = threading.Event()
def blocking_put(**kwargs):
put_started.set()
release_put.wait(timeout=2)
return True, 1, {}
connector.put.side_effect = blocking_put
adapter.save_async(multimodal_output=None, request=first)
task = adapter._pending_save_reqs.popleft()
sender = threading.Thread(target=adapter._send_single_request, args=(task,))
sender.start()
assert put_started.wait(timeout=1)
cleaner = threading.Thread(
target=adapter.finish_requests,
args=([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first}),
)
cleaner.start()
try:
cleaner.join(timeout=0.5)
assert not cleaner.is_alive()
assert task["sender_token"].cancelled
assert first.external_req_id in adapter._sender_tokens
finally:
release_put.set()
sender.join(timeout=1)
cleaner.join(timeout=1)
assert not sender.is_alive()
assert first.external_req_id not in adapter._sender_tokens
assert first.external_req_id not in adapter.put_req_chunk
def test_cleanup_only_affects_target_request(build_adapter):
"""Cleanup for one request must not affect another request's state."""
adapter, _ = build_adapter(stage_id=1)
_populate_adapter_state(adapter, "req-a", "ext-a")
_populate_adapter_state(adapter, "req-b", "ext-b")
adapter.cleanup("req-a", "ext-a")
assert "req-b" in adapter.upstream_exhausted_requests
assert "req-b" in adapter.get_req_chunk
assert "ext-b" in adapter.put_req_chunk
assert "ext-b" in adapter.request_payload
assert "ext-b" in adapter.code_prompt_token_ids
assert "req-b" in adapter.request_ids_mapping
def test_cleanup_after_poll_flow(build_adapter):
"""Simulate full load_async -> poll -> finished -> cleanup cycle."""
adapter, connector = build_adapter(stage_id=2, model_mode="ar")
request = _req("req-flow", RequestStatus.WAITING, external_req_id="ext-flow")
adapter.load_async(request)
adapter.request_ids_mapping["req-flow"] = "ext-flow"
payload: OmniPayload = {
"hidden_states": {"output": torch.tensor([[1.0]])},
"meta": {"finished": torch.tensor(True, dtype=torch.bool)},
}
connector.get.return_value = (payload, 8)
adapter._poll_single_request(_dequeue_load_entry(adapter, request))
assert "req-flow" in adapter.upstream_exhausted_requests
assert adapter.get_req_chunk["req-flow"] == 1
assert "req-flow" in adapter.request_ids_mapping
adapter.cleanup("req-flow", "ext-flow")
assert "req-flow" not in adapter.upstream_exhausted_requests
assert "req-flow" not in adapter.get_req_chunk
assert "req-flow" not in adapter.request_ids_mapping
assert "ext-flow" not in adapter.request_payload
def test_finish_requests_restores_status(build_adapter):
"""Abort path must pop ``requests_origin_status`` and restore pre-wait status.
While ``process_pending_chunks`` holds a request off the scheduler queues, the
adapter records the prior status (WAITING or RUNNING). ``finish_requests`` must
put that status back on the live ``Request`` so base ``Scheduler.finish_requests``
can finish bookkeeping without inconsistent state / crashes.
"""
adapter, _ = build_adapter(stage_id=1)
req_id = "req-abort-during-chunk"
prior = RequestStatus.RUNNING
request = _req(req_id, RequestStatus.WAITING_FOR_CHUNK)
adapter.requests_origin_status[req_id] = prior
adapter.waiting_for_chunk_running_requests.append(request)
requests_map = {req_id: request}
adapter.finish_requests([req_id], RequestStatus.FINISHED_ABORTED, requests_map)
assert request.status == prior
assert req_id not in adapter.requests_origin_status
assert not adapter.waiting_for_chunk_running_requests
def test_finish_requests_does_not_restore_stale_status_without_connector_ownership(build_adapter):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-restored", RequestStatus.WAITING_FOR_STREAMING_REQ)
request.resumable = True
adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING
adapter.finish_requests(
[request.request_id],
RequestStatus.FINISHED_ABORTED,
{request.request_id: request},
)
assert request.status == RequestStatus.WAITING_FOR_STREAMING_REQ
assert request.request_id not in adapter.requests_origin_status
def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter):
adapter, _ = build_adapter(stage_id=1)
zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
other = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
adapter.waiting_for_chunk_waiting_requests = deque([zombie, other])
adapter.waiting_for_chunk_running_requests = deque([other, zombie])
adapter.requests_with_ready_chunks.add("req-zombie")
adapter.upstream_exhausted_requests.add("req-zombie")
requests_map = {
"req-zombie": zombie,
"req-live": other,
}
adapter.finish_requests(
["req-zombie"],
RequestStatus.FINISHED_ABORTED,
requests_map,
)
assert [req.request_id for req in adapter.waiting_for_chunk_waiting_requests] == ["req-live"]
assert [req.request_id for req in adapter.waiting_for_chunk_running_requests] == ["req-live"]
assert "req-zombie" not in adapter.requests_with_ready_chunks
assert "req-zombie" not in adapter.upstream_exhausted_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
@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler])
@pytest.mark.parametrize(
("placement", "origin_status", "initial_status", "streaming_counter_owned"),
[
("hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, False),
("running", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False),
("waiting", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False),
("skipped", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, True),
("skipped_streaming", RequestStatus.RUNNING, RequestStatus.WAITING_FOR_STREAMING_REQ, True),
("skipped_hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, True),
("skipped_non_streaming", RequestStatus.WAITING, RequestStatus.WAITING, False),
],
)
def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
build_adapter,
scheduler_cls,
placement,
origin_status,
initial_status,
streaming_counter_owned,
):
adapter, _ = build_adapter(stage_id=1, active_stream_window=2)
request = _req(
"req-segment-stop",
initial_status,
external_req_id="ext-segment-stop",
)
request.resumable = True
adapter.requests_origin_status[request.request_id] = origin_status
if placement in {"hidden", "skipped_hidden"}:
adapter.waiting_for_chunk_running_requests.append(request)
adapter._active_streams[request.request_id] = request
scheduler = scheduler_cls.__new__(scheduler_cls)
scheduler.chunk_transfer_adapter = adapter
scheduler.input_coordinator = None
scheduler.requests = {request.request_id: request}
scheduler.running = [request] if placement == "running" else []
scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else [])
scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)
freed = []
def fake_free_request(self, live, delay_free_blocks=False):
del delay_free_blocks
freed.append(live.request_id)
self.requests.pop(live.request_id)
return None, None
scheduler._free_request = MethodType(fake_free_request, scheduler)
first = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED)
second = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED)
assert len(first) == 1
assert second == []
assert freed == [request.request_id]
assert request.request_id not in scheduler.requests
assert request.request_id not in adapter._active_streams
assert not adapter.waiting_for_chunk_running_requests
assert scheduler.running == []
assert list(scheduler.waiting) == []
assert list(scheduler.skipped_waiting) == []
assert scheduler.num_waiting_for_streaming_input == 0
fresh_a = _req("req-fresh-a", RequestStatus.WAITING)
fresh_b = _req("req-fresh-b", RequestStatus.WAITING)
assert adapter._ensure_active_stream(fresh_a)
assert adapter._ensure_active_stream(fresh_b)
assert set(adapter._active_streams) == {fresh_a.request_id, fresh_b.request_id}
@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler])
def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapter, scheduler_cls):
adapter, _ = build_adapter(stage_id=1, active_stream_window=2)
request = _req("req-off-queue", RequestStatus.FINISHED_STOPPED)
request.resumable = True
adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING
scheduler = scheduler_cls.__new__(scheduler_cls)
scheduler.chunk_transfer_adapter = adapter
scheduler.input_coordinator = None
scheduler.requests = {request.request_id: request}
scheduler.running = []
scheduler.waiting = DummyWaitingQueue()
scheduler.skipped_waiting = DummyWaitingQueue()
scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
assert (
scheduler_cls.finish_requests(
scheduler,
[request.request_id],
RequestStatus.FINISHED_ABORTED,
)
== []
)
assert scheduler.requests == {request.request_id: request}
assert request.status == RequestStatus.FINISHED_STOPPED
assert request.request_id not in adapter.requests_origin_status
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)
live = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
waiting_queue = DummyWaitingQueue()
running_queue = []
adapter.waiting_for_chunk_waiting_requests = deque([zombie, live])
adapter.waiting_for_chunk_running_requests = deque([zombie, live])
adapter.restore_queues(
waiting_queue,
running_queue,
scheduler_requests={"req-live": live},
)
assert [req.request_id for req in waiting_queue] == ["req-live"]
assert [req.request_id for req in running_queue] == ["req-live"]
assert not adapter.waiting_for_chunk_waiting_requests
assert not adapter.waiting_for_chunk_running_requests
# ---------------------------------------------------------------
# Scheduler trigger tests
# ---------------------------------------------------------------
class _HashableRequest(SimpleNamespace):
"""SimpleNamespace that can be added to a set (needed by scheduler internals)."""
# vLLM 0.26: update_from_output settles this counter for every scheduled
# request; real Requests initialise it to 0 (vllm/v1/request.py).
num_in_flight_tokens = 0
# vLLM 0.27 (a0c092ee72): the stale-output counter that replaced
# async_tokens_to_discard. update_from_output reads it for every scheduled
# request, and real Requests initialise it to 0 (vllm/v1/request.py), so the
# double needs it too.
num_stale_output_tokens = 0
def __hash__(self):
return hash(self.request_id)
def __eq__(self, other):
return getattr(other, "request_id", None) == self.request_id
def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: MockerFixture):
"""OmniGenerationScheduler must call adapter.cleanup when request finishes."""
cleanup_calls = []
adapter_mock = mocker.MagicMock()
adapter_mock.upstream_exhausted_requests = {"req-s1"}
adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
scheduler = mocker.MagicMock()
scheduler.chunk_transfer_adapter = adapter_mock
scheduler.connector = None
scheduler.ec_connector = None
scheduler.perf_metrics = None
scheduler.log_stats = False
scheduler.recompute_kv_load_failures = False
scheduler.structured_output_manager = mocker.MagicMock()
scheduler.structured_output_manager.accept_tokens.return_value = True
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager.take_events.return_value = None
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = mocker.MagicMock()
request = _HashableRequest(
request_id="req-s1",
external_req_id="ext-s1",
status=RequestStatus.RUNNING,
is_finished=lambda: False,
num_computed_tokens=10,
num_prompt_tokens=10,
prompt_token_ids=list(range(10)),
num_output_placeholders=0,
sampling_params=None,
pooling_params=None,
stop_reason=None,
client_index=0,
take_events=lambda: [],
trace_headers=None,
has_encoder_inputs=False,
take_prefill_stats=lambda: None,
num_nans_in_logits=0,
get_finished_reason=lambda: "stop",
)
scheduler.requests = {"req-s1": request}
scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
scheduler._free_request = mocker.MagicMock(return_value=(None, None))
scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
scheduler.running = [request]
scheduler.waiting = mocker.MagicMock()
scheduler.waiting.remove_requests = mocker.MagicMock()
scheduler.make_stats = mocker.MagicMock(return_value=None)
scheduler_output = SimpleNamespace(
num_scheduled_tokens={"req-s1": 10},
scheduled_spec_decode_tokens={},
num_invalid_spec_tokens=0,
)
model_runner_output = SimpleNamespace(
sampled_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
pooler_output=None,
num_nans_in_logits=None,
kv_connector_output=None,
cudagraph_stats=None,
req_id_to_index={"req-s1": 0},
routed_experts=None,
routed_experts_dict=None,
)
OmniGenerationScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
assert len(cleanup_calls) == 1
args, _ = cleanup_calls[0]
assert args[0] == "req-s1"
assert args[1] == "ext-s1"
def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerFixture):
"""OmniARScheduler should enqueue save; adapter cleanup is handled in save thread."""
cleanup_calls = []
save_calls = []
adapter_mock = mocker.MagicMock()
adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
scheduler = mocker.MagicMock()
scheduler.chunk_transfer_adapter = adapter_mock
scheduler.connector = None
scheduler.perf_metrics = None
scheduler.log_stats = False
scheduler.recompute_kv_load_failures = False
scheduler.structured_output_manager = mocker.MagicMock()
scheduler.structured_output_manager.accept_tokens.return_value = True
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager = mocker.MagicMock()
scheduler.kv_cache_manager.take_events.return_value = None
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = mocker.MagicMock()
scheduler.waiting_for_transfer_free = set()
scheduler.transfer_triggered_requests = set()
scheduler.active_kv_transfers = set()
request = _HashableRequest(
request_id="req-ar",
external_req_id="ext-ar",
status=RequestStatus.RUNNING,
is_finished=lambda: False,
num_computed_tokens=1,
num_prompt_tokens=1,
prompt_token_ids=[1],
num_output_placeholders=0,
sampling_params=None,
pooling_params=None,
stop_reason=None,
client_index=0,
take_events=lambda: [],
trace_headers=None,
has_encoder_inputs=False,
take_prefill_stats=lambda: None,
num_nans_in_logits=0,
get_finished_reason=lambda: "stop",
)
scheduler.requests = {"req-ar": request}
scheduler._update_request_with_output = mocker.MagicMock(return_value=([], True))
scheduler._process_kv_transfer_trigger = mocker.MagicMock(return_value=False)
scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
scheduler._free_request = mocker.MagicMock(return_value=(None, None))
scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
scheduler.running = [request]
scheduler.waiting = mocker.MagicMock()
scheduler.waiting.remove_requests = mocker.MagicMock()
scheduler.make_spec_decoding_stats = mocker.MagicMock(return_value=None)
scheduler.make_stats = mocker.MagicMock(return_value=None)
scheduler_output = SimpleNamespace(
num_scheduled_tokens={"req-ar": 1},
scheduled_spec_decode_tokens={},
num_invalid_spec_tokens=0,
)
model_runner_output = SimpleNamespace(
sampled_token_ids=[[123]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=None,
num_nans_in_logits=None,
kv_connector_output=None,
cudagraph_stats=None,
req_id_to_index={"req-ar": 0},
kv_extracted_req_ids=None,
routed_experts=None,
routed_experts_dict=None,
)
OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
assert len(cleanup_calls) == 0
assert len(save_calls) == 1
def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
"""``OmniARScheduler.finish_requests`` must run chunk adapter hook before vLLM base."""
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
order: list[str] = []
adapter = mocker.MagicMock()
def _adapter_finish(request_ids, finished_status, requests):
order.append("adapter")
return []
adapter.finish_requests.side_effect = _adapter_finish
def _super_finish(_self, request_ids, finished_status):
order.append("super")
return []
sched = OmniARScheduler.__new__(OmniARScheduler)
sched.chunk_transfer_adapter = adapter
sched.requests = {}
sched.running = []
sched.waiting = []
with patch.object(VLLMScheduler, "finish_requests", _super_finish):
OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
assert order == ["adapter", "super"]
def _parked_sender_scheduler(adapter, session):
"""A minimal OmniARScheduler holding one parked streaming session."""
scheduler = OmniARScheduler.__new__(OmniARScheduler)
scheduler.chunk_transfer_adapter = adapter
scheduler.input_coordinator = None
scheduler.requests = {session.request_id: session}
scheduler.running = []
scheduler.waiting = DummyWaitingQueue()
scheduler.skipped_waiting = DummyWaitingQueue([session])
scheduler.num_waiting_for_streaming_input = 1
scheduler.finished_req_ids_dict = {}
def _free_request(self, request, delay_free_blocks=False):
del delay_free_blocks
self.requests.pop(request.request_id, None)
return None, None
scheduler._free_request = MethodType(_free_request, scheduler)
return scheduler
@pytest.mark.parametrize("send_first_chunk", [False, True])
def test_parked_streaming_final_update_emits_downstream_terminal(build_adapter, send_first_chunk):
"""#6670: a final update on a parked sender must terminate the receiver.
Upstream turns "final update while parked in WAITING_FOR_STREAMING_REQ"
into a silent local abort. The downstream stage only learns a stream ended
from a chunk carrying ``finished``, so the abort used to leave it in
WAITING_FOR_CHUNK until the 600s input deadline.
"""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
session = _req(
"req-parked-final",
RequestStatus.WAITING_FOR_STREAMING_REQ,
external_req_id="ext-parked-final",
)
session.resumable = True
session._omni_segment_generation = 0
# A segment stop queues its boundary before the final update arrives.
# The downstream request is already prewarmed even if the sender has not
# dequeued this first task, so it still needs a terminal in both cases.
adapter.save_async(None, session, is_segment_finished=True)
assert session.external_req_id not in adapter.put_req_chunk
if send_first_chunk:
adapter._send_single_request(adapter._pending_save_reqs.popleft())
connector.put.reset_mock()
scheduler = _parked_sender_scheduler(adapter, session)
final_update = _req(session.request_id, RequestStatus.WAITING)
final_update.resumable = False
OmniARScheduler.add_request(scheduler, final_update)
# Local teardown happened...
assert session.request_id not in scheduler.requests
# ...but only after a terminal chunk was queued for the next stage.
assert session.status == RequestStatus.FINISHED_STOPPED
assert any(task["is_finished"] for task in adapter._pending_save_reqs)
while adapter._pending_save_reqs:
adapter._send_single_request(adapter._pending_save_reqs.popleft())
connector.put.assert_called_once()
assert connector.put.call_args.kwargs["put_key"] == f"ext-parked-final_1_{int(send_first_chunk)}"
payload = connector.put.call_args.kwargs["data"]
assert bool(payload.meta.finished.item()) is True
assert bool(payload.meta.is_segment_finished.item()) is False
def test_queued_terminal_chunk_survives_sender_cleanup(build_adapter):
"""#6670: ``cleanup_sender`` must not discard an already-queued terminal."""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal")
request.resumable = False
adapter.save_async(None, request, is_segment_finished=False)
assert len(adapter._pending_save_reqs) == 1
# The scheduler tears the request down before the save_loop drains.
adapter.cleanup_sender("ext-terminal")
adapter._send_single_request(adapter._pending_save_reqs.popleft())
connector.put.assert_called_once()
assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True
# The terminal put reclaims the generation it kept alive.
assert "ext-terminal" not in adapter._sender_tokens
def test_non_terminal_chunk_is_still_dropped_by_sender_cleanup(build_adapter):
"""The #6670 fence is terminal-only: ordinary queued chunks still drop."""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
request = _req("req-mid", RequestStatus.RUNNING, external_req_id="ext-mid")
request.resumable = True
adapter.save_async(None, request, is_segment_finished=False)
adapter.cleanup_sender("ext-mid")
adapter._send_single_request(adapter._pending_save_reqs.popleft())
connector.put.assert_not_called()
def test_final_update_on_non_sender_stage_keeps_upstream_abort(build_adapter):
"""A final stage never put a chunk, so it must keep the upstream path."""
adapter, _connector = build_adapter(stage_id=2, model_mode="generation")
session = _req(
"req-final-stage",
RequestStatus.WAITING_FOR_STREAMING_REQ,
external_req_id="ext-final-stage",
)
session.resumable = True
scheduler = _parked_sender_scheduler(adapter, session)
final_update = _req(session.request_id, RequestStatus.WAITING)
final_update.resumable = False
with patch.object(VLLMScheduler, "add_request") as base_add:
OmniARScheduler.add_request(scheduler, final_update)
base_add.assert_called_once()
assert not adapter._pending_save_reqs
def _queue_terminal(
adapter,
external_req_id: str,
req_id: str = "req-terminal",
segment_generation: int = 0,
) -> None:
"""Queue one request-terminal chunk on the adapter's save queue.
``segment_generation`` must match the adapter's dedup watermark, which a
preceding segment stop advances; otherwise ``save_async`` drops the chunk
as a late duplicate.
"""
request = _req(req_id, RequestStatus.FINISHED_STOPPED, external_req_id=external_req_id)
request.resumable = False
request._omni_segment_generation = segment_generation
adapter.save_async(None, request, is_segment_finished=False)
def _assert_terminal_put(connector) -> None:
connector.put.assert_called_once()
assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True
def test_terminal_chunk_survives_cleanup_after_dequeue(build_adapter):
"""#6670: the fence must hold across the connector handoff, not just the queue.
``_finish_parked_streaming_session`` queues the terminal and then calls
``finish_requests``, so ``cleanup_sender`` can land *after* the save_loop
has already dequeued the terminal. Releasing the fence at dequeue left
that interleaving cancelling the chunk before ``put`` -- the same lost
terminal, and the same downstream WAITING_FOR_CHUNK hang.
"""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
_queue_terminal(adapter, "ext-late-cleanup")
task = adapter._pending_save_reqs.popleft()
real_send = adapter._send_single_request_for_generation
def send_with_cleanup_in_flight(inner_task, sender_token=None):
# The scheduler thread reaches cleanup_sender here: the save_loop owns
# the terminal but has not handed it to the connector yet.
adapter.cleanup_sender("ext-late-cleanup")
return real_send(inner_task, sender_token)
adapter._send_single_request_for_generation = send_with_cleanup_in_flight
adapter._send_single_request(task)
_assert_terminal_put(connector)
# The generation the fence kept alive is reclaimed, not leaked.
assert "ext-late-cleanup" not in adapter._sender_tokens
assert "ext-late-cleanup" not in adapter.put_req_chunk
def test_terminal_chunk_survives_cleanup_racing_an_inflight_sibling(build_adapter):
"""#6670: an in-flight ordinary chunk must not reclaim a fenced generation.
``_send_single_request``'s ``finally`` also tears a cancelled generation
down. If the stop and the abort both land while an earlier chunk is on the
wire, that teardown would unregister the token and the queued terminal
would then fail the identity check.
"""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
running = _req("req-sibling", RequestStatus.RUNNING, external_req_id="ext-sibling")
running.resumable = True
adapter.save_async(None, running, is_segment_finished=True)
ordinary_task = adapter._pending_save_reqs.popleft()
real_send = adapter._send_single_request_for_generation
def send_then_stop_and_abort(inner_task, sender_token=None):
result = real_send(inner_task, sender_token)
# Terminal stop and abort both land while this chunk is on the wire.
# The segment stop above already armed the next segment's watermark.
_queue_terminal(adapter, "ext-sibling", req_id="req-sibling", segment_generation=1)
adapter.cleanup_sender("ext-sibling")
return result
adapter._send_single_request_for_generation = send_then_stop_and_abort
adapter._send_single_request(ordinary_task)
adapter._send_single_request_for_generation = real_send
assert len(adapter._pending_save_reqs) == 1
connector.put.reset_mock()
adapter._send_single_request(adapter._pending_save_reqs.popleft())
_assert_terminal_put(connector)
assert "ext-sibling" not in adapter._sender_tokens
@pytest.mark.parametrize("finish_inside_put", [False, True])
def test_successful_put_keeps_terminal_chunk_sequence(build_adapter, finish_inside_put):
"""A cleanup during put must not reuse a key already consumed downstream."""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
receiver, receiver_connector = build_adapter(stage_id=2, model_mode="generation")
session = _req("req-sequence", RequestStatus.WAITING_FOR_STREAMING_REQ, external_req_id="ext-sequence")
session.resumable = True
session._omni_segment_generation = 0
scheduler = _parked_sender_scheduler(adapter, session)
final_update = _req(session.request_id, RequestStatus.WAITING)
final_update.resumable = False
downstream = _req("req-downstream", RequestStatus.WAITING, external_req_id=session.external_req_id)
receiver.request_ids_mapping[downstream.request_id] = session.external_req_id
payloads = {}
def receive_next():
entry = _dequeue_load_entry(receiver, downstream)
return receiver._poll_single_request(entry)
def put_then_finish(**kwargs):
payload = kwargs["data"]
payloads[kwargs["put_key"]] = {
"meta": {
"finished": bool(payload.meta.finished.item()),
"is_segment_finished": bool(payload.meta.is_segment_finished.item()),
}
}
if not payload.meta.finished.item():
# Consume the ordinary boundary while put is still in flight.
assert receive_next()
if finish_inside_put:
scheduler.add_request(final_update)
return True, 1, {}
def get_by_key(_from_stage, _to_stage, key, metadata=None):
assert metadata is None
payload = payloads.pop(key, None)
return (payload, 1) if payload is not None else None
connector.put.side_effect = put_then_finish
receiver_connector.get.side_effect = get_by_key
adapter.save_async(None, session, is_segment_finished=True)
adapter._send_single_request(adapter._pending_save_reqs.popleft())
if not finish_inside_put:
scheduler.add_request(final_update)
while adapter._pending_save_reqs:
adapter._send_single_request(adapter._pending_save_reqs.popleft())
assert [call.kwargs["put_key"] for call in connector.put.call_args_list] == ["ext-sequence_1_0", "ext-sequence_1_1"]
assert receive_next()
assert downstream.request_id in receiver.upstream_exhausted_requests
assert not payloads
assert session.external_req_id not in adapter._sender_tokens
def test_late_successful_put_cannot_advance_replacement_generation(build_adapter):
"""The post-put guard must retain identity checks while relaxing cancellation."""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
request = _req("req-old", RequestStatus.RUNNING, external_req_id="ext-replaced")
request.resumable = True
adapter.save_async(None, request, is_segment_finished=True)
task = adapter._pending_save_reqs.popleft()
old_token = task["sender_token"]
replacement_token = type(old_token)()
def replace_during_put(**kwargs):
# Model a completion that no longer owns its external id. The old
# terminal fence must never authorize writes to the successor's state.
with adapter._sender_state_lock:
old_token.cancelled = True
old_token.terminal_pending = True
adapter._sender_tokens[request.external_req_id] = replacement_token
adapter.put_req_chunk[request.external_req_id] = 4
adapter.ramp_chunk_count[request.external_req_id] = 2
return True, 1, {}
connector.put.side_effect = replace_during_put
adapter._send_single_request(task)
assert adapter._sender_tokens[request.external_req_id] is replacement_token
assert adapter.put_req_chunk[request.external_req_id] == 4
assert adapter.ramp_chunk_count[request.external_req_id] == 2
def test_new_request_cannot_join_a_generation_with_a_queued_terminal(build_adapter):
"""#6670: the fence must not readmit a request reusing the external id.
``save_async`` rejects chunks whose generation is ``cancelled``; that is
what stops a new request from inheriting a retiring generation's token and
chunk counters. Deferring the reclaim must not also defer that retirement.
"""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
_queue_terminal(adapter, "ext-shared", req_id="req-old")
adapter.cleanup_sender("ext-shared")
newcomer = _req("req-new", RequestStatus.RUNNING, external_req_id="ext-shared")
newcomer.resumable = True
adapter.save_async(None, newcomer, is_segment_finished=True)
# Rejected loudly, so the scheduler fails the newcomer now instead of its
# chunks being silently discarded at send time.
assert len(adapter._pending_save_reqs) == 1
assert adapter.collect_failed_send_request_ids() == {"req-new": "previous sender generation is still draining"}
adapter._send_single_request(adapter._pending_save_reqs.popleft())
_assert_terminal_put(connector)
assert "ext-shared" not in adapter._sender_tokens
def test_failed_terminal_put_runs_the_deferred_cleanup(build_adapter):
"""#6670: a silent connector drop must not strand the deferred cleanup.
``connector.put`` returning False skips the success path that would have
called ``cleanup``. The reclaim ``cleanup_sender`` handed over is owed on
that path too, or the whole per-request sender state leaks for good.
"""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
_queue_terminal(adapter, "ext-put-failed")
adapter.cleanup_sender("ext-put-failed")
connector.put.return_value = (False, 0, {})
adapter._send_single_request(adapter._pending_save_reqs.popleft())
assert "ext-put-failed" not in adapter._sender_tokens
assert "ext-put-failed" not in adapter.put_req_chunk
def test_raising_terminal_put_runs_the_deferred_cleanup(build_adapter):
"""#6670: the same reclaim is owed when the connector raises."""
adapter, connector = build_adapter(stage_id=1, model_mode="ar")
_queue_terminal(adapter, "ext-put-raised")
adapter.cleanup_sender("ext-put-raised")
connector.put.side_effect = RuntimeError("connector down")
# ``save_loop`` catches this and records a send failure.
with pytest.raises(RuntimeError):
adapter._send_single_request(adapter._pending_save_reqs.popleft())
assert "ext-put-raised" not in adapter._sender_tokens
assert "ext-put-raised" not in adapter.put_req_chunk
def test_wire_round_trip_struct_to_dict_contract():
"""Pin the wire contract: encoding ``OmniPayloadStruct`` and decoding it
yields a dict equivalent to ``to_dict(struct)``.
The chunk-adapter sender uses struct attribute access while the receiver
uses dict-key access. This works only because ``OmniMsgpackDecoder`` has
no target type and decodes structs back to plain dicts. If this test
breaks, the receiver's dict access will silently drop fields or KeyError.
"""
from vllm_omni.data_entry_keys import CodesStruct, to_dict
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackDecoder,
OmniMsgpackEncoder,
)
struct = OmniPayloadStruct(
meta=MetaStruct(
finished=torch.tensor(True, dtype=torch.bool),
left_context_size=12,
),
codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.int64)),
)
encoded = OmniMsgpackEncoder().encode(struct)
decoded = OmniMsgpackDecoder().decode(encoded)
assert isinstance(decoded, dict)
assert isinstance(decoded["meta"], dict)
assert isinstance(decoded["meta"]["finished"], torch.Tensor)
assert bool(decoded["meta"]["finished"].item()) is True
assert decoded["meta"]["left_context_size"] == 12
assert torch.equal(decoded["codes"]["audio"], torch.tensor([1, 2, 3], dtype=torch.int64))
expected = to_dict(struct)
assert set(decoded.keys()) == set(expected.keys())
assert set(decoded["meta"].keys()) == set(expected["meta"].keys())
assert set(decoded["codes"].keys()) == set(expected["codes"].keys())
# ---------------------------------------------------------------
# Deferred finish for upstream-completed requests
# ---------------------------------------------------------------
def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
"""Build a mock scheduler with requests queued for deferred finish."""
adapter_mock = mocker.MagicMock()
adapter_mock.upstream_exhausted_requests = {r.request_id for r in pending_finish_reqs}
cleanup_calls = []
adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
scheduler = mocker.MagicMock()
scheduler.chunk_transfer_adapter = adapter_mock
scheduler.connector = None
scheduler.ec_connector = None
scheduler.perf_metrics = None
scheduler.log_stats = False
scheduler.recompute_kv_load_failures = False
scheduler.structured_output_manager = mocker.MagicMock()
scheduler.structured_output_manager.accept_tokens.return_value = True
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager.take_events.return_value = None
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = mocker.MagicMock()
scheduler._pending_finish_reqs = list(pending_finish_reqs)
scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
scheduler._free_request = mocker.MagicMock(return_value=(None, None))
scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
scheduler.running = list(running)
scheduler.waiting = mocker.MagicMock()
scheduler.waiting.remove_requests = mocker.MagicMock()
scheduler.make_stats = mocker.MagicMock(return_value=None)
scheduler.requests = {r.request_id: r for r in running}
scheduler_output = SimpleNamespace(
num_scheduled_tokens={},
scheduled_spec_decode_tokens={},
num_invalid_spec_tokens=0,
)
model_runner_output = SimpleNamespace(
sampled_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
pooler_output=None,
num_nans_in_logits=None,
kv_connector_output=None,
cudagraph_stats=None,
req_id_to_index={},
)
return scheduler, scheduler_output, model_runner_output, cleanup_calls
def test_deferred_finish_emits_finished_output(mocker: MockerFixture):
"""A request whose upstream completed with no remaining tokens should
emit a FINISHED EngineCoreOutput, free resources, and clean up adapter state."""
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
request = _HashableRequest(
request_id="req-df1",
external_req_id="ext-df1",
status=RequestStatus.RUNNING,
is_finished=lambda: False,
num_computed_tokens=16,
num_prompt_tokens=16,
prompt_token_ids=list(range(1, 17)),
num_output_placeholders=0,
sampling_params=None,
pooling_params=None,
stop_reason=None,
client_index=0,
take_events=lambda: [],
trace_headers=None,
has_encoder_inputs=False,
take_prefill_stats=lambda: None,
num_nans_in_logits=0,
get_finished_reason=lambda: "stop",
)
scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
mocker,
running=[request],
pending_finish_reqs=[request],
)
scheduler._free_request.return_value = ({"mock": "kv_params"}, None)
result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
assert request.status == RequestStatus.FINISHED_STOPPED
scheduler._handle_stopped_request.assert_called_once_with(request)
scheduler._free_request.assert_called_once_with(request)
assert len(cleanup_calls) == 1
assert cleanup_calls[0][0] == ("req-df1", "ext-df1")
eco = result[0]
assert len(eco.outputs) == 1
assert eco.outputs[0].request_id == "req-df1"
assert eco.outputs[0].finish_reason == "stop"
assert eco.outputs[0].kv_transfer_params == {"mock": "kv_params"}
assert scheduler._pending_finish_reqs == []
def test_deferred_finish_empty_prompt(mocker: MockerFixture):
"""A request that never received any tokens (finished immediately upstream)
should still emit a FINISHED output and clean up."""
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
request = _HashableRequest(
request_id="req-df2",
external_req_id="ext-df2",
status=RequestStatus.WAITING,
is_finished=lambda: False,
num_computed_tokens=0,
num_prompt_tokens=0,
prompt_token_ids=[],
num_output_placeholders=0,
sampling_params=None,
pooling_params=None,
stop_reason=None,
client_index=0,
take_events=lambda: [],
trace_headers=None,
has_encoder_inputs=False,
take_prefill_stats=lambda: None,
num_nans_in_logits=0,
get_finished_reason=lambda: "stop",
)
scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
mocker,
running=[],
pending_finish_reqs=[request],
)
result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
assert request.status == RequestStatus.FINISHED_STOPPED
scheduler._free_request.assert_called_once_with(request)
assert len(cleanup_calls) == 1
eco = result[0]
assert len(eco.outputs) == 1
assert eco.outputs[0].finish_reason == "stop"
assert scheduler._pending_finish_reqs == []
def test_deferred_finish_skips_already_finished(mocker: MockerFixture):
"""A request that was aborted between schedule() and update_from_output()
should be skipped without emitting output or freeing resources twice."""
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
request = _HashableRequest(
request_id="req-df3",
external_req_id="ext-df3",
status=RequestStatus.FINISHED_ABORTED,
is_finished=lambda: True,
num_computed_tokens=0,
num_prompt_tokens=0,
prompt_token_ids=[],
num_output_placeholders=0,
sampling_params=None,
pooling_params=None,
stop_reason=None,
client_index=0,
take_events=lambda: [],
trace_headers=None,
has_encoder_inputs=False,
take_prefill_stats=lambda: None,
num_nans_in_logits=0,
get_finished_reason=lambda: "stop",
)
scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
mocker,
running=[],
pending_finish_reqs=[request],
)
result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
scheduler._handle_stopped_request.assert_not_called()
scheduler._free_request.assert_not_called()
assert len(cleanup_calls) == 0
assert 0 not in result
assert scheduler._pending_finish_reqs == []
def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture):
"""When _handle_stopped_request returns False (resumable request), the
output must still be emitted so the client stream doesn't hang."""
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
request = _HashableRequest(
request_id="req-df4",
external_req_id="ext-df4",
status=RequestStatus.RUNNING,
is_finished=lambda: False,
num_computed_tokens=16,
num_prompt_tokens=16,
prompt_token_ids=list(range(1, 17)),
num_output_placeholders=0,
sampling_params=None,
pooling_params=None,
stop_reason=None,
client_index=0,
take_events=lambda: [],
trace_headers=None,
has_encoder_inputs=False,
take_prefill_stats=lambda: None,
num_nans_in_logits=0,
get_finished_reason=lambda: "stop",
)
scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
mocker,
running=[request],
pending_finish_reqs=[request],
)
scheduler._handle_stopped_request.return_value = False
result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
scheduler._handle_stopped_request.assert_called_once_with(request)
scheduler._free_request.assert_not_called()
assert len(cleanup_calls) == 0
eco = result[0]
assert len(eco.outputs) == 1
assert eco.outputs[0].request_id == "req-df4"
assert eco.outputs[0].finish_reason == "stop"
assert eco.outputs[0].kv_transfer_params is None
assert scheduler._pending_finish_reqs == []
# ---------------------------------------------------------------
# Zombie purge tests (regression for vllm-project/vllm-omni#3736)
# ---------------------------------------------------------------
def test_process_pending_chunks_purges_zombies_in_running_deque(
build_adapter,
):
"""A request aborted while parked in ``waiting_for_chunk_running_requests``
is no longer tracked by ``scheduler.requests`` once
``Scheduler._free_request`` runs. ``process_pending_chunks`` must drop
that zombie *before* ``restore_queues`` would re-inject it onto the
running queue.
Regression for https://github.com/vllm-project/vllm-omni/issues/3736
(engine-core ``KeyError`` on aborts under chunk-transfer pressure).
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
live_req = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
zombie_req = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
adapter.waiting_for_chunk_running_requests.append(live_req)
adapter.waiting_for_chunk_running_requests.append(zombie_req)
# Mirror state the adapter would carry for an in-flight request so we can
# later assert ``cleanup_receiver`` actually fired against the zombie.
adapter.requests_origin_status[zombie_req.request_id] = RequestStatus.RUNNING
adapter.requests_origin_status[live_req.request_id] = RequestStatus.RUNNING
waiting_queue = DummyWaitingQueue()
running_queue: list = []
# Simulate the post-abort scheduler: only ``live-1`` is still tracked.
scheduler_requests = {live_req.request_id: live_req}
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
# 1. Zombie was popped from the deque.
assert zombie_req not in adapter.waiting_for_chunk_running_requests
# 2. Live request is still in the deque.
assert live_req in adapter.waiting_for_chunk_running_requests
# 3. ``cleanup_receiver`` ran for the zombie (drops origin-status mapping
# and registers the id as cancelled so a late load/poll is dropped too).
assert zombie_req.request_id not in adapter.requests_origin_status
assert zombie_req.request_id in adapter._cancelled_load_reqs
# 4. Live request's bookkeeping is untouched.
assert adapter.requests_origin_status[live_req.request_id] == RequestStatus.RUNNING
assert live_req.request_id not in adapter._cancelled_load_reqs
# 5. ``restore_queues`` (which the scheduler runs in its ``finally``
# clause) now only re-injects the live request -- the zombie is
# gone, so the worker's ``_update_states`` cannot crash on it.
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert running_queue == [live_req]
def test_process_pending_chunks_purges_zombies_in_waiting_deque(build_adapter):
"""Zombie purge applies symmetrically to the waiting-side deque.
Regression for https://github.com/vllm-project/vllm-omni/issues/3736 --
aborted requests that landed in ``waiting_for_chunk_waiting_requests``
must also be removed before ``restore_queues`` re-injects them, since
``restore_queues`` uses ``waiting_queue.add_request`` rather than
``running_queue.extend`` for that path.
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
live_req = _req("live-w", RequestStatus.WAITING_FOR_CHUNK)
zombie_req = _req("zombie-w", RequestStatus.WAITING_FOR_CHUNK)
adapter.waiting_for_chunk_waiting_requests.append(live_req)
adapter.waiting_for_chunk_waiting_requests.append(zombie_req)
waiting_queue = DummyWaitingQueue()
running_queue: list = []
scheduler_requests = {live_req.request_id: live_req}
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert live_req in adapter.waiting_for_chunk_waiting_requests
assert zombie_req not in adapter.waiting_for_chunk_waiting_requests
assert zombie_req.request_id in adapter._cancelled_load_reqs
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert waiting_queue == [live_req]
def test_purge_preserves_live_order_with_interleaved_zombies(build_adapter):
"""Interleaved live/zombie ordering -- pin the ``popleft`` + append-live
in-place filter at chunk_transfer_adapter._purge_untracked_chunk_requests:
survivor order must match insertion order, and ``cleanup_receiver`` must
run exactly once per zombie.
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
live1 = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
zombie1 = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
live2 = _req("live-2", RequestStatus.WAITING_FOR_CHUNK)
zombie2 = _req("zombie-2", RequestStatus.WAITING_FOR_CHUNK)
for req in (live1, zombie1, live2, zombie2):
adapter.waiting_for_chunk_running_requests.append(req)
waiting_queue = DummyWaitingQueue()
running_queue: list = []
scheduler_requests = {live1.request_id: live1, live2.request_id: live2}
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert list(adapter.waiting_for_chunk_running_requests) == [live1, live2]
assert zombie1.request_id in adapter._cancelled_load_reqs
assert zombie2.request_id in adapter._cancelled_load_reqs
def test_restore_queues_purges_late_aborts_after_process_pending_chunks(
build_adapter,
):
"""Race window guard: an abort can fire between ``process_pending_chunks``
and the scheduler's ``finally``-clause ``restore_queues`` call. The
second purge inside ``restore_queues`` must drop the now-untracked
request so it does not get re-injected onto ``running_queue``.
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
req = _req("late-abort", RequestStatus.WAITING_FOR_CHUNK)
adapter.waiting_for_chunk_running_requests.append(req)
waiting_queue = DummyWaitingQueue()
running_queue: list = []
# Tick N: process_pending_chunks runs while req is still tracked.
scheduler_requests: dict = {req.request_id: req}
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert req in adapter.waiting_for_chunk_running_requests
# Mid-tick: abort fires and the scheduler's free path deletes the entry.
del scheduler_requests[req.request_id]
# finally: restore_queues sees the now-untracked req and must drop it
# instead of blindly extending it onto running_queue.
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert running_queue == []
assert req.request_id in adapter._cancelled_load_reqs
def test_purge_is_noop_on_empty_deques(build_adapter):
"""Empty deques short-circuit -- guards against any accidental
``IndexError`` from ``popleft`` on an empty deque if a future caller
shadows the empty check.
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
assert len(adapter.waiting_for_chunk_waiting_requests) == 0
assert len(adapter.waiting_for_chunk_running_requests) == 0
waiting_queue = DummyWaitingQueue()
running_queue: list = []
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests={})
assert len(adapter.waiting_for_chunk_waiting_requests) == 0
assert len(adapter.waiting_for_chunk_running_requests) == 0
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
assert running_queue == []
assert waiting_queue == []
# --------------------------------------------------------------------------- #
# Chunk-wait deadline (RFC #4855 R1.1, issue #3833)
#
# Before this, the async-chunk path had no deadline of any kind, so a dropped
# terminal chunk or an upstream stage that died mid-stream parked the request
# in WAITING_FOR_CHUNK forever. None of these scenarios were covered.
# --------------------------------------------------------------------------- #
def _park_in_chunk_wait(adapter, request, *, waiting=True):
"""Drive a request through one process_pending_chunks round into the wait."""
queue = DummyWaitingQueue([request]) if waiting else [request]
if waiting:
adapter.process_pending_chunks(queue, [])
else:
adapter.process_pending_chunks(DummyWaitingQueue(), queue)
return queue
def test_chunk_wait_clock_starts_when_a_request_parks(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
request = _req("r1", RequestStatus.WAITING)
_park_in_chunk_wait(adapter, request)
assert request.status == RequestStatus.WAITING_FOR_CHUNK
assert "r1" in adapter._waiting_since
def test_chunk_wait_does_not_expire_before_the_deadline(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
assert "r1" in adapter._waiting_since
def test_dropped_terminal_chunk_expires_the_request(build_adapter):
"""The #3833 scenario: upstream stops sending and never marks the stream done."""
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
adapter._waiting_since["r1"] -= 601.0
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"r1"}
# Cleared, so a second sweep does not re-report the same request.
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
def test_arriving_chunk_resets_the_clock(build_adapter):
"""A slow but healthy stream must never expire: the deadline measures stall
time between chunks, not the lifetime of the stream."""
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
request = _req("r1", RequestStatus.WAITING)
queue = _park_in_chunk_wait(adapter, request)
adapter.restore_queues(queue, [], scheduler_requests={"r1": request})
# Age the wait to just under the deadline, then deliver a chunk.
adapter._waiting_since["r1"] -= 599.0
adapter._finished_load_reqs.add("r1")
adapter.process_pending_chunks(queue, [])
assert "r1" not in adapter._waiting_since
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
def test_a_disabled_deadline_never_expires(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
adapter._waiting_since["r1"] -= 10_000.0
assert adapter.collect_timed_out_request_ids(timeout_s=0.0) == set()
assert adapter.collect_timed_out_request_ids(timeout_s=-1.0) == set()
def test_finished_request_leaves_no_stale_timestamp(build_adapter):
"""Otherwise a completed request would be reported as timed out later."""
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
request = _req("r1", RequestStatus.WAITING)
_park_in_chunk_wait(adapter, request)
adapter.finish_requests(["r1"], RequestStatus.FINISHED_STOPPED, {"r1": request})
assert "r1" not in adapter._waiting_since
assert adapter.collect_timed_out_request_ids(timeout_s=0.001) == set()
def test_aborted_request_leaves_no_stale_timestamp(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
adapter.cleanup_receiver("r1")
assert "r1" not in adapter._waiting_since
def test_expiry_is_per_request(build_adapter):
"""One stalled stream must not take down its healthy neighbours."""
adapter, _ = build_adapter(stage_id=1, model_mode="ar", max_num_seqs=4)
stalled = _req("stalled", RequestStatus.WAITING)
healthy = _req("healthy", RequestStatus.WAITING)
adapter.process_pending_chunks(DummyWaitingQueue([stalled, healthy]), [])
adapter._waiting_since["stalled"] -= 601.0
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
assert "healthy" in adapter._waiting_since
def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter):
from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
talker2code2wav_async_chunk,
)
adapter, _ = build_adapter(stage_id=1, connector_extra={"codec_chunk_frames": 1})
external_req_id = "ext-native-abort"
first = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
first.resumable = True
first.model_intermediate_buffer = None
first.additional_information = {"meta": {"codec_streaming": True}, "nvc_logical_prompt_len": 1}
first_frame = torch.full((1, 31), 101, dtype=torch.long)
talker2code2wav_async_chunk(
adapter,
{"codes": {"audio": first_frame}, "meta": {"codec_streaming": True}},
first,
is_finished=False,
)
adapter.requests_num_chunks_sent[external_req_id] = 1
adapter.finish_requests([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first})
assert external_req_id not in adapter.request_payload
assert external_req_id not in adapter.requests_num_chunks_sent
replacement = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
replacement.resumable = True
replacement.model_intermediate_buffer = None
replacement.additional_information = first.additional_information
replacement_payload = talker2code2wav_async_chunk(
adapter,
{"codes": {"audio": torch.full((1, 31), 202, dtype=torch.long)}, "meta": {"codec_streaming": True}},
replacement,
is_finished=False,
)
assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long))
def test_save_async_boundary_holds_generation_without_request_counter(build_adapter):
"""A boundary must not advance past a generation nothing can catch up to.
When the producing stage maintains ``_omni_segment_generation`` on the
request, a boundary may start the next segment's watermark at
``generation + 1`` because the following frame will carry it. On paths
where the request has no such counter, that speculative advance makes every
subsequent frame look stale and starves the stream permanently
(vllm-project/vllm-omni#6816).
"""
adapter, _ = build_adapter(stage_id=1)
request = _req("req-nogen", RequestStatus.WAITING, external_req_id="ext-nogen")
request.resumable = True
request._omni_segment_generation = None
adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
assert adapter._segment_generation["ext-nogen"] == 0
follow_up = _req("req-nogen-next", RequestStatus.WAITING, external_req_id="ext-nogen")
follow_up.resumable = True
follow_up.num_computed_tokens = 3
follow_up._omni_segment_generation = None
queued_before = len(adapter._pending_save_reqs)
adapter.save_async(multimodal_output=None, request=follow_up, is_segment_finished=False)
assert len(adapter._pending_save_reqs) == queued_before + 1
assert adapter._pending_save_reqs[-1]["request"] is follow_up
Reference in new issue View Git Blame Copy Permalink