vllm-omni/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

# 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,
        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 = mocker.MagicMock()
        connector.stage_id = stage_id
        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_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_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


# ---------------------------------------------------------------
# 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 == []