anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
6a0a578eb25243f9240f15c9586074029830080f
vllm-omni/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
T

2148 lines
85 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
from collections import deque
from types import MethodType, SimpleNamespace
from unittest.mock import Mock, patch
import pytest
import torch
from pytest_mock import MockerFixture
from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
from vllm.v1.request import 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 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_streaming_payload_can_append_exact_prompt_length(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": {"next_stage_prompt_len": 3},
}
construct_next_stage_streaming_input_prompt(payload, request)
assert request.prompt_token_ids == [0, 0, 7, 8, 0, 0, 0]
assert request._all_token_ids == [0, 0, 7, 8, 0, 0, 0]
assert request._output_token_ids == []
assert request.num_computed_tokens == 4
assert request.num_prompt_tokens == 7
request.update_block_hashes.assert_called_once_with()
@pytest.fixture
def build_adapter(monkeypatch, mocker: MockerFixture):
def _build(
*,
stage_id: int = 1,
model_mode: str = "ar",
max_num_seqs: int = 2,
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.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),
)
model_config = SimpleNamespace(
worker_type=model_mode,
max_num_seqs=max_num_seqs,
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
@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
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(request)
assert request.additional_information == payload
assert adapter.get_req_chunk["req-1"] == 1
assert "req-1" in adapter._finished_load_reqs
assert "req-1" in adapter.upstream_exhausted_requests
assert "req-1" not in adapter._pending_load_reqs
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(request) is True
assert request.num_computed_tokens == 0
assert request.prompt_token_ids == [0] * 10
assert request.request_id in adapter.replaced_streaming_prompt_ids
request.status = RequestStatus.WAITING_FOR_CHUNK
running_queue = [request]
waiting_queue = DummyWaitingQueue()
adapter.process_pending_chunks(waiting_queue, running_queue)
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
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(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)]
assert adapter._poll_single_request(request) 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(request) 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_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_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(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(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(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(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(request)
# AR mode now forwards the latest payload directly.
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(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(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 list(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
adapter._pending_load_reqs.append(SimpleNamespace(request_id=req_id))
adapter._finished_load_reqs.add(req_id)
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 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_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_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(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.should_advance.return_value = False
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager.take_events.return_value = None
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.should_advance.return_value = False
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager = mocker.MagicMock()
scheduler.kv_cache_manager.take_events.return_value = None
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 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.should_advance.return_value = False
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager.take_events.return_value = None
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
Reference in new issue View Git Blame Copy Permalink