# 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_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