vllm-omni/tests/worker_v2/test_omni_generation_model_runner.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""OmniGenerationModelRunner contracts: empty-step lifecycle, output partition,
CPU-sync vs CUDA-async dispatch, and async-chunk slot recycling."""

from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import MagicMock

import numpy as np
import pytest
import torch
from vllm.distributed.aux_output_connector.connector import AuxRequestOutput

from vllm_omni.model_executor.models.output_templates import OmniOutput
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner, check_exact_input_shape

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


@pytest.mark.parametrize("randomize_inputs", [False, True])
def test_profile_run_accepts_release_flag_without_running_codec(monkeypatch, randomize_inputs):
    runner = object.__new__(OmniGenerationModelRunner)
    runner.model = MagicMock(side_effect=AssertionError("profiling ran codec model"))
    synchronize = MagicMock()
    monkeypatch.setattr(torch.accelerator, "synchronize", synchronize)
    runner.profile_run(randomize_inputs=randomize_inputs)
    synchronize.assert_called_once_with()
    runner.model.assert_not_called()


class _FakeStagedField:
    def __init__(self, data):
        self.np = data
        self._staged = []

    def stage_write_elem(self, idx, value):
        self._staged.append((idx, value))

    def apply_write(self):
        for idx, value in self._staged:
            self.np[idx] = value
        self._staged.clear()


def _make_runner(model_output, num_reqs=1, prompt_len=10):
    runner = object.__new__(OmniGenerationModelRunner)
    runner.device = torch.device("cpu")
    runner.model_config = MagicMock(hf_text_config=None)
    del runner.model_config.eos_token_id
    runner.postprocess = lambda *a, **kw: None
    runner._gen_model_output = model_output
    runner._gen_input_batch = SimpleNamespace(
        num_reqs=num_reqs, req_ids=[f"req-{i}" for i in range(num_reqs)], idx_mapping_np=np.arange(num_reqs)
    )
    runner.execute_model_state = SimpleNamespace(finished_req_ids={"finished"}, ec_connector_output=None)
    runner.kv_connector = SimpleNamespace(post_forward=MagicMock(return_value=None))
    runner.check_ep_fault = False
    runner.aux_output_connector = None
    runner.req_states = MagicMock(
        prompt_len=SimpleNamespace(np=np.full(num_reqs, prompt_len, dtype=np.int32)),
        num_computed_tokens=_FakeStagedField(np.zeros(num_reqs, dtype=np.int32)),
    )
    return runner


@pytest.mark.parametrize("aux_enabled", [False, True])
def test_control_only_step_keeps_lifecycle_and_skips_input_construction(aux_enabled):
    runner = object.__new__(OmniGenerationModelRunner)
    order = []
    for name in (
        "_prepare_native_data_plane",
        "finish_requests",
        "free_states",
        "_apply_block_table_staged_writes_if_available",
    ):
        setattr(runner, name, lambda *args, name=name: order.append(name))
    for name in ("_handle_async_chunk_updates", "add_requests", "update_requests", "_sync_native_data_plane_payloads"):
        setattr(runner, name, MagicMock(side_effect=AssertionError("control step built model inputs")))
    output = object()
    runner.kv_connector = SimpleNamespace(no_forward=lambda _s: output)
    runner._merge_ec_connector_no_forward = lambda _s, value: value
    runner._attach_native_data_plane_signals = lambda value: value
    metadata = object()
    runner.aux_output_connector = SimpleNamespace(begin_step=MagicMock()) if aux_enabled else None
    scheduler_output = SimpleNamespace(
        total_num_scheduled_tokens=0,
        scheduled_new_reqs=[],
        scheduled_cached_reqs=SimpleNamespace(req_ids=[]),
        aux_output_connector_metadata=metadata,
    )
    assert runner.execute_model(scheduler_output) is output
    assert order == [
        "_prepare_native_data_plane",
        "finish_requests",
        "free_states",
        "_apply_block_table_staged_writes_if_available",
    ]
    if aux_enabled:
        assert runner.aux_output_connector is not None
        runner.aux_output_connector.begin_step.assert_called_once_with(metadata)


def test_released_chunk_reuses_scheduler_output_and_slot_recycle_clears_state():
    from vllm.v1.core.sched.output import SchedulerOutput

    from vllm_omni.core.sched.output import OmniCachedRequestData

    runner = object.__new__(OmniGenerationModelRunner)
    runner.model = SimpleNamespace(logits_processor=None)
    runner._remove_request = MagicMock()
    req_id_to_index = MagicMock()
    req_id_to_index.get.return_value = None
    runner.req_states = SimpleNamespace(
        add_request=MagicMock(), req_id_to_index=req_id_to_index, apply_staged_writes=MagicMock()
    )
    runner.pooling_runner = None
    runner.encoder_cache = None
    runner.model_state = SimpleNamespace(
        add_request=MagicMock(),
        apply_staged_writes=MagicMock(),
        remove_request=MagicMock(),
        intermediate_buffer=SimpleNamespace(remove_request=MagicMock()),
    )
    runner.block_tables = SimpleNamespace(append_block_ids=MagicMock())
    runner.lora_state = SimpleNamespace(add_request=MagicMock())
    runner.is_last_pp_rank = False
    runner.adaptive_verification = None
    runner.sampler = None
    cached = OmniCachedRequestData(
        req_ids=["req"],
        resumed_req_ids=set(),
        new_token_ids=[[]],
        all_token_ids={"req": [1]},
        new_block_ids=[()],
        num_computed_tokens=[0],
        num_output_tokens=[0],
        prompt_token_ids={"req": [1]},
        additional_information={"req": None},
    )
    scheduler_output = SchedulerOutput(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=cached,
        num_scheduled_tokens={"req": 1},
        total_num_scheduled_tokens=1,
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        num_common_prefix_blocks=[0],
        finished_req_ids=set(),
        free_encoder_mm_hashes=[],
    )
    runner._handle_async_chunk_updates(scheduler_output)
    added = runner.model_state.add_request.call_args.args[1]
    assert (added.req_id, added.prompt_token_ids) == ("req", [1])


def test_sample_tokens_cpu_sync_owns_outputs_and_runs_connector_last():
    waveform = torch.arange(4, dtype=torch.float32)
    output = OmniOutput(text_hidden_states=torch.empty(0), multimodal_outputs={"codes": {"audio": [waveform]}})
    runner = _make_runner(output, num_reqs=1)
    runner._reserve_native_data_plane_outputs = MagicMock()
    runner._finalize_native_data_plane_output = MagicMock(side_effect=lambda value: value)

    def post_forward(_finished):
        assert runner.req_states.num_computed_tokens.np.tolist() == [10]

    runner.kv_connector.post_forward.side_effect = post_forward
    result = OmniGenerationModelRunner.sample_tokens(runner)
    assert isinstance(result, OmniModelRunnerOutput)
    # Reserve precedes finalize; connector runs after token state is applied.
    runner._reserve_native_data_plane_outputs.assert_called_once_with(["req-0"])
    runner.kv_connector.post_forward.assert_called_once_with({"finished"})
    waveform.fill_(-1)
    assert torch.equal(result.multimodal_outputs[0]["codes.audio"], torch.arange(4, dtype=torch.float32))


def test_sample_tokens_uses_async_output_for_cuda_and_snapshots_req_ids(monkeypatch):
    from vllm_omni.worker_v2 import omni_generation_model_runner as generation_runner

    output = OmniOutput(text_hidden_states=torch.zeros(1), multimodal_outputs={"model_outputs": [torch.randn(4)]})
    runner = _make_runner(output, num_reqs=1)
    runner.device = SimpleNamespace(type="cuda")
    runner.main_stream = object()
    runner.output_copy_stream = object()
    runner.model_config.async_chunk = True
    runner._release_generation_slots = MagicMock()
    runner._finalize_native_data_plane_output = MagicMock()
    monkeypatch.setattr(generation_runner, "_contains_cuda_tensor", lambda _: True)
    captured = {}

    class _FakeAsyncOutput:
        def __init__(self, **kwargs):
            captured.update(kwargs)

    monkeypatch.setattr(generation_runner, "OmniGenerationAsyncOutput", _FakeAsyncOutput)
    input_batch = runner._gen_input_batch
    result = generation_runner.OmniGenerationModelRunner.sample_tokens(runner)

    assert isinstance(result, _FakeAsyncOutput)
    assert captured["finalize_output"] is runner._finalize_native_data_plane_output
    assert captured["model_runner_output"].sampled_token_ids == [[]]
    runner._release_generation_slots.assert_called_once()
    output_req_ids = captured["model_runner_output"].req_ids
    input_batch.req_ids[0] = "reused"  # snapshot must insulate the published output
    assert output_req_ids == ["req-0"]


def test_generation_aux_output_copies_and_materializes_without_sampling(monkeypatch):
    from vllm_omni.worker_v2 import omni_generation_model_runner as module

    events = []

    class PendingOutput:
        def enqueue_cpu_copy(self, *, num_sampled, num_rejected):
            np.testing.assert_array_equal(num_sampled, [0])
            np.testing.assert_array_equal(num_rejected, [0])
            events.append("copy")

        def process_output(self):
            assert events == ["copy", "record", "synchronize"]
            return {"req-0": AuxRequestOutput(0, np.array([[2, 3]], dtype=np.uint8))}

    class CopyEvent:
        def record(self, stream):
            events.append("record")

        def synchronize(self):
            events.append("synchronize")

    monkeypatch.setattr(torch.cuda, "Event", lambda **kw: CopyEvent())
    monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
    monkeypatch.setattr(module, "_async_copy_mm", lambda *a, **kw: {})
    output = OmniModelRunnerOutput(["req-0"], {"req-0": 0}, [[]])
    result = module.OmniGenerationAsyncOutput(
        model_runner_output=output,
        multimodal_outputs={},
        num_reqs=1,
        main_stream=object(),
        copy_stream=SimpleNamespace(wait_stream=lambda stream: None),
        pending_aux_output=PendingOutput(),
    ).get_output()

    assert result.sampled_token_ids == [[]]
    np.testing.assert_array_equal(result.aux_output_connector_output["req-0"].rows, [[2, 3]])


def test_pending_aux_output_forces_async_delivery_without_cuda_payload(monkeypatch):
    from vllm_omni.worker_v2 import omni_generation_model_runner as module

    runner = _make_runner(OmniOutput(text_hidden_states=torch.empty(0), multimodal_outputs={}))
    runner.device = SimpleNamespace(type="cuda")
    runner.main_stream = object()
    runner.output_copy_stream = object()
    pending = object()
    runner.aux_output_connector = SimpleNamespace(prepare_output=MagicMock(return_value=pending))
    runner._reserve_native_data_plane_outputs = MagicMock()
    runner._release_generation_slots = MagicMock()
    runner._finalize_native_data_plane_output = MagicMock()
    captured = {}

    class AsyncOutput:
        def __init__(self, **kwargs):
            captured.update(kwargs)

    monkeypatch.setattr(module, "OmniGenerationAsyncOutput", AsyncOutput)
    input_batch = runner._gen_input_batch
    result = runner.sample_tokens()

    assert isinstance(result, AsyncOutput)
    runner.aux_output_connector.prepare_output.assert_called_once_with(input_batch)
    assert captured["pending_aux_output"] is pending


@pytest.mark.parametrize("exact", [True, False])
def test_exact_input_shape_models_reject_padded_batches(exact):
    # Qwen3-Omni Code2Wav splits input_ids by seq_token_counts (#6712).
    model = SimpleNamespace(requires_exact_input_shape=exact)
    check_exact_input_shape(model, torch.zeros(32, dtype=torch.long), 32)
    check_exact_input_shape(model, None, 16)
    if exact:
        with pytest.raises(RuntimeError, match="exact input shape"):
            check_exact_input_shape(model, torch.zeros(48, dtype=torch.long), 32)
    else:
        check_exact_input_shape(model, torch.zeros(48, dtype=torch.long), 32)