# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import threading from collections import deque from types import SimpleNamespace from unittest.mock import patch import pytest import torch from pytest_mock import MockerFixture from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.request import RequestStatus from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct 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 _req(req_id: str, status: RequestStatus, external_req_id: str | None = None): return SimpleNamespace( request_id=req_id, external_req_id=external_req_id or req_id, status=status, prompt_token_ids=[], num_computed_tokens=0, num_output_placeholders=0, additional_information=None, is_finished=lambda: status == RequestStatus.FINISHED_STOPPED, ) @pytest.fixture def build_adapter(monkeypatch, mocker: MockerFixture): def _build( *, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2, 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) 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.finished_requests assert "req-1" not in adapter._pending_load_reqs def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter): adapter, connector = build_adapter(stage_id=1, model_mode="generation") request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor") codes = torch.tensor([[1, 2], [3, 4]], dtype=torch.long) payload: OmniPayload = { "codes": {"audio": codes}, "meta": { "left_context_size": 1, "finished": torch.tensor(False, dtype=torch.bool), }, } connector.get.return_value = (payload, 16) adapter._poll_single_request(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(pooling_output=None, request=request) adapter.custom_process_next_stage_input_func = lambda **kwargs: {} adapter.save_async(pooling_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(pooling_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_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({"pooling_output": None, "request": request, "is_finished": False}) assert cleanup_calls == [] # no terminal cleanup; meta.finished is unobservable 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({"pooling_output": None, "request": request, "is_finished": False}) assert connector.put.called sent_payload = connector.put.call_args.kwargs["data"] assert isinstance(sent_payload, OmniPayloadStruct) assert sent_payload.meta is None # confirms it's the empty struct on the wire 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({"pooling_output": None, "request": request, "is_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.finished_requests 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_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] adapter.process_pending_chunks(waiting_queue, running_queue) 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) 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_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"} # --------------------------------------------------------------- # 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.finished_requests.add(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.finished_requests 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.finished_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.finished_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.finished_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.finished_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 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 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.finished_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.finished_requests 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).""" 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.finished_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) 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_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) 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_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 = {} 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.finished_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) 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"} 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 == []