# 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, 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 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_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.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, ) assert len(adapter._pending_save_reqs) == 1 boundary_task = adapter._pending_save_reqs.popleft() assert "ext-stream" not in adapter.requests_num_chunks_sent request.num_computed_tokens = 3 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]) 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_until_dedup_is_reset(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_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