vllm-omni/tests/entrypoints/test_async_omni.py

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

import asyncio
import re
from types import SimpleNamespace

import pytest
from vllm.lora.request import LoRARequest
from vllm.sampling_params import RequestOutputKind, SamplingParams

from tests.helpers.mark import hardware_test
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.outputs import OmniRequestOutput

pytestmark = [pytest.mark.core_model]

DIFFUSION_MODEL = "riverclouds/qwen_image_random"
OMNI_MODEL = "Qwen/Qwen2.5-Omni-7B"
OMNI_STAGE_CONFIG = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml")


async def _noop(*args, **kw):
    pass


def get_fake_add_request(submitted_request_ids, submitted_lora_requests=None):
    async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
        lora_request = kwargs.get("lora_request")
        del prompt, sampling_params_list, final_stage_id, kwargs
        submitted_request_ids.append(request_id)
        if submitted_lora_requests is not None:
            submitted_lora_requests.append(lora_request)

    return fake_add_request_async


def get_fake_abort(aborted_request_batches):
    async def fake_abort_async(request_ids):
        aborted_request_batches.append(list(request_ids))

    return fake_abort_async


async def fake_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts):
    del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts
    if request_id.startswith("cancel-"):
        await asyncio.Future()
        return
    yield SimpleNamespace(
        stage_id=0,
        request_output=SimpleNamespace(outputs=[]),
        finished=True,
    )


def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) -> AsyncOmni:
    omni = object.__new__(AsyncOmni)
    omni._pause_cond = asyncio.Condition()
    omni._paused = False
    omni.engine = SimpleNamespace(
        num_stages=1,
        add_request_async=fake_add_request,
        abort_async=fake_abort_request,
    )
    omni.log_stats = False
    omni.request_states = {}
    omni._final_output_handler = lambda: None
    omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params
    omni._compute_final_stage_id = lambda output_modalities: 0
    omni._compute_final_output_stage_ids = lambda output_modalities: [0]
    omni._process_orchestrator_results = fake_process_results
    omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
    return omni


@pytest.mark.cpu
def test_generate_submits_randomized_id_to_engine():
    """Ensure the engine receives a UUID-suffixed ID, not the raw request ID"""

    async def run():
        submitted_ids: list[str] = []
        omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids))

        req_id = "my-req-1"
        async for _ in omni.generate(
            prompt={"prompt": "test"},
            request_id=req_id,
            sampling_params_list=[SimpleNamespace()],
            output_modalities=["text"],
        ):
            pass

        assert len(submitted_ids) == 1
        assert submitted_ids[0] != req_id
        assert submitted_ids[0].startswith(f"{req_id}-")

    asyncio.run(run())


@pytest.mark.cpu
def test_generate_forwards_lora_request_to_engine():
    """Ensure the lora_request passed to generate() reaches add_request_async.

    Regression test for https://github.com/vllm-project/vllm-omni/issues/5369:
    AsyncOmni.generate() accepted lora_request but silently dropped it when
    submitting the request, so generation always used the base model.
    """

    async def run():
        submitted_ids: list[str] = []
        submitted_loras: list[LoRARequest | None] = []
        omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids, submitted_loras))

        lora = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake")
        async for _ in omni.generate(
            prompt={"prompt": "test"},
            request_id="lora-req",
            sampling_params_list=[SimpleNamespace()],
            output_modalities=["text"],
            lora_request=lora,
        ):
            pass

        assert len(submitted_ids) == 1
        assert len(submitted_loras) == 1
        assert submitted_loras[0] is lora

    asyncio.run(run())


@pytest.mark.cpu
@pytest.mark.parametrize(
    "req_ids,cancel_prefix,expected_cancel_count",
    [
        (["cancel-me"], "cancel-me", 1),
        (["cancel-me", "cancel-me"], "cancel-me", 2),
        (["cancel-hello", "cancel-hello-world"], "cancel-hello", 1),
    ],
)
def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefix: str, expected_cancel_count: int):
    """Ensure that abort() with the user-visible ID resolves correctly.

    NOTE: In the case of concurrent / colliding request(s), all requests matching the
    user provided request ID will be aborted."""

    async def run():
        aborted_batches: list[list[str]] = []
        omni = get_async_omni_instance(
            fake_abort_request=get_fake_abort(aborted_batches),
        )

        async def exhaust(agen):
            async for _ in agen:
                pass

        tasks = []
        for user_request_id in req_ids:
            t = asyncio.create_task(
                exhaust(
                    omni.generate(
                        prompt={"prompt": "test"},
                        request_id=user_request_id,
                        sampling_params_list=[SimpleNamespace()],
                    )
                )
            )
            await asyncio.sleep(0)
            tasks.append(t)

        assert len(omni.request_states) == len(req_ids)
        await omni.abort(cancel_prefix)

        assert len(aborted_batches) == 1
        aborted_ids = aborted_batches[0]
        # Aborted requests will have fmt {ext_id}-{UUID} to avoid collisions
        for rid in aborted_ids:
            assert re.fullmatch(rf"{re.escape(cancel_prefix)}-[0-9a-f]+", rid)
        assert len(aborted_ids) == expected_cancel_count
        assert len(set(aborted_ids)) == expected_cancel_count

        for t in tasks:
            t.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)

    asyncio.run(run())


@pytest.mark.cpu
def test_abort_keeps_request_states_until_generate_cleanup():
    """Frontend abort must not drop request_states; generate() owns cleanup."""

    async def run():
        release = asyncio.Event()
        seen_during_wait: list[int] = []

        async def slow_abort_async(request_ids):
            del request_ids
            seen_during_wait.append(len(omni.request_states))
            await release.wait()

        omni = get_async_omni_instance(fake_abort_request=slow_abort_async)
        omni.request_states["req-1-aaaa"] = SimpleNamespace(
            request_id="req-1-aaaa",
            external_request_id="req-1",
            input_stream_task=None,
        )

        task = asyncio.create_task(omni.abort("req-1"))
        await asyncio.sleep(0)
        assert not task.done()
        assert "req-1-aaaa" in omni.request_states
        assert seen_during_wait == [1]

        release.set()
        await task
        assert "req-1-aaaa" in omni.request_states

    asyncio.run(run())


@pytest.mark.cpu
def test_abort_propagates_engine_errors_without_popping_state():
    async def run():
        async def failing_abort_async(request_ids):
            del request_ids
            raise RuntimeError("orchestrator abort failed")

        omni = get_async_omni_instance(fake_abort_request=failing_abort_async)
        omni.request_states["req-1-bbbb"] = SimpleNamespace(
            request_id="req-1-bbbb",
            external_request_id="req-1",
            input_stream_task=None,
        )

        with pytest.raises(RuntimeError, match="orchestrator abort failed"):
            await omni.abort("req-1")
        assert "req-1-bbbb" in omni.request_states

    asyncio.run(run())


@pytest.mark.cpu
def test_abort_enqueues_prefix_tokens_from_engine():
    """AR abort must deliver the generated prefix on the request queue (CPU)."""

    async def run():
        from vllm.outputs import CompletionOutput

        from vllm_omni.engine.messages import OutputMessage

        queue: asyncio.Queue = asyncio.Queue()

        async def abort_with_prefix(request_ids):
            rid = request_ids[0]
            engine_output = OmniRequestOutput(
                request_id=rid,
                finished=True,
                stage_id=0,
                final_output_type="text",
                outputs=[
                    CompletionOutput(
                        index=0,
                        text="partial",
                        token_ids=[7, 8, 9],
                        cumulative_logprob=None,
                        logprobs=None,
                        finish_reason="abort",
                        stop_reason=None,
                    )
                ],
            )
            return [
                OutputMessage(
                    request_id=rid,
                    stage_id=0,
                    replica_id=None,
                    engine_outputs=engine_output,
                    metrics=None,
                    finished=True,
                    stage_submit_ts=None,
                )
            ]

        omni = get_async_omni_instance(fake_abort_request=abort_with_prefix)
        omni.request_states["req-1-cccc"] = SimpleNamespace(
            request_id="req-1-cccc",
            external_request_id="req-1",
            input_stream_task=None,
            queue=queue,
        )
        await omni._abort(["req-1-cccc"])
        msg = queue.get_nowait()
        assert msg.finished is True
        assert list(msg.engine_outputs.outputs[0].token_ids) == [7, 8, 9]
        assert msg.engine_outputs.outputs[0].finish_reason == "abort"
        assert "req-1-cccc" in omni.request_states

    asyncio.run(run())


@pytest.mark.cpu
def test_abort_enqueues_synthetic_finished_when_engine_returns_empty():
    """Empty abort_async must still unblock generate() with finish_reason=abort."""

    async def run():
        queue: asyncio.Queue = asyncio.Queue()

        async def empty_abort_async(request_ids):
            del request_ids
            return []

        omni = get_async_omni_instance(fake_abort_request=empty_abort_async)
        omni.request_states["req-1-dddd"] = SimpleNamespace(
            request_id="req-1-dddd",
            external_request_id="req-1",
            input_stream_task=None,
            queue=queue,
        )
        await omni._abort(["req-1-dddd"])
        msg = queue.get_nowait()
        assert msg.finished is True
        assert msg.request_id == "req-1-dddd"
        assert msg.engine_outputs.outputs[0].finish_reason == "abort"
        assert "req-1-dddd" in omni.request_states

    asyncio.run(run())


@pytest.mark.cpu
def test_generate_accepts_request_after_repeated_cancellations():
    async def run_test():
        submitted_request_ids: list[str] = []
        aborted_request_batches: list[list[str]] = []

        async def collect_outputs(request_id):
            outputs = []
            async for output in AsyncOmni.generate(
                omni,
                prompt={"prompt": "prompt"},
                request_id=request_id,
                sampling_params_list=[SimpleNamespace()],
                output_modalities=["image"],
            ):
                outputs.append(output)
            return outputs

        omni = get_async_omni_instance(
            fake_add_request=get_fake_add_request(submitted_request_ids),
            fake_abort_request=get_fake_abort(aborted_request_batches),
        )

        assert len(await collect_outputs("baseline")) == 1

        for idx in range(3):
            task = asyncio.create_task(collect_outputs(f"cancel-{idx}"))
            await asyncio.sleep(0)
            task.cancel()
            with pytest.raises(asyncio.CancelledError):
                await task

        assert len(await collect_outputs("after-cancel")) == 1

        # Check prefixes instead of equality since generate will add
        # a random numeric suffix to ensure the request ID is unique
        expected_prefixes = ["baseline-", "cancel-0-", "cancel-1-", "cancel-2-", "after-cancel-"]
        assert len(submitted_request_ids) == len(expected_prefixes)
        for submitted, prefix in zip(submitted_request_ids, expected_prefixes):
            assert submitted.startswith(prefix)

        assert len(aborted_request_batches) == 3
        for batch, prefix in zip(aborted_request_batches, ["cancel-0-", "cancel-1-", "cancel-2-"]):
            assert len(batch) == 1
            assert batch[0].startswith(prefix)

    asyncio.run(run_test())


@pytest.mark.cpu
def test_generate_cancellation_converges_after_engine_shutdown_starts(mocker):
    async def run_test():
        submitted_request_ids: list[str] = []
        engine = object.__new__(AsyncOmniEngine)
        engine._shutdown_called = True
        engine.request_queue = mocker.Mock()
        engine.num_stages = 1
        engine.add_request_async = get_fake_add_request(submitted_request_ids)

        omni = get_async_omni_instance()
        omni.engine = engine

        async def collect_outputs():
            async for _ in omni.generate(
                prompt={"prompt": "prompt"},
                request_id="cancel-shutdown",
                sampling_params_list=[SimpleNamespace()],
                output_modalities=["image"],
            ):
                pass

        task = asyncio.create_task(collect_outputs())
        await asyncio.sleep(0)
        task.cancel()

        with pytest.raises(asyncio.CancelledError):
            await task

        assert len(submitted_request_ids) == 1
        assert submitted_request_ids[0].startswith("cancel-shutdown-")
        assert omni.request_states == {}
        engine.request_queue.sync_q.put.assert_not_called()

    asyncio.run(run_test())


@pytest.mark.cpu
def test_generate_yields_streaming_diffusion_chunks_before_final():
    """AsyncOmni.generate yields every intermediate diffusion chunk before the final one."""

    async def streaming_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts):
        del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts
        yield OmniRequestOutput.from_diffusion(
            request_id=request_id,
            images=[],
            final_output_type="image",
            custom_output={"chunk": 0},
            finished=False,
        )
        yield OmniRequestOutput.from_diffusion(
            request_id=request_id,
            images=[],
            final_output_type="image",
            custom_output={"chunk": 1},
            finished=True,
        )

    async def run_test():
        omni = get_async_omni_instance()
        omni._process_orchestrator_results = streaming_process_results

        outputs = []
        async for output in AsyncOmni.generate(
            omni,
            prompt={"prompt": "a cat"},
            request_id="req-stream",
            sampling_params_list=[SimpleNamespace()],
            output_modalities=["image"],
        ):
            outputs.append(output)

        assert len(outputs) == 2
        assert [output.finished for output in outputs] == [False, True]
        assert [output.custom_output["chunk"] for output in outputs] == [0, 1]

    asyncio.run(run_test())


@pytest.mark.cpu
@pytest.mark.parametrize(
    "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE]
)
def test_output_kind_is_preserved_with_explicit_sampling_params(output_kind):
    """Ensure we don't change the output kind in async generate if params are provided directly."""

    captured_params = []

    async def capturing_add_request(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
        del prompt, final_stage_id, kwargs
        captured_params.extend(sampling_params_list)

    async def run():
        omni = get_async_omni_instance(fake_add_request=capturing_add_request)
        sp = SamplingParams(output_kind=output_kind)
        async for _ in omni.generate(
            prompt={"prompt": "test"},
            request_id="test-req",
            sampling_params_list=[sp],
            output_modalities=["text"],
        ):
            pass

    asyncio.run(run())
    assert captured_params[0].output_kind == output_kind


# End to end tests for ensuring internal manipulation of request ID
# in diffusion / Omni models don't leak back to the user.
#
# One AsyncOmni per test function (all cases in a single asyncio loop) to avoid
# repeated cold starts. Do not use class/module-scoped engine fixtures here:
# pytest-asyncio uses a function-scoped event loop by default, so reusing an
# engine across tests can hang on the second generate() call.


# Covers:
#   * plain client ids (``my-req-1``)
#   * OpenAI-style prefixed ids (``img_gen-*``, ``chatcmpl-*``) that AsyncOmni
#     suffixes internally for engine routing — streamed outputs must still echo
#     the caller-visible id, not the internal UUID-suffixed id
#   * empty ``request_id`` — server assigns a non-empty id for the caller
_DIFFUSION_REQ_IDS = ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"]
_OMNI_REQ_IDS = ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"]


@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.omni
@pytest.mark.asyncio
async def test_diffusion_generate_request_id():
    """Diffusion E2E request-id contract (``riverclouds/qwen_image_random``).

    Scenarios (one engine, sequential ``generate`` calls):
    - plain id ``my-req-1``
    - image-style prefix ``img_gen-abc123``
    - chat-style prefix ``chatcmpl-xyz``
    - empty ``request_id`` → output id is non-empty (auto-assigned)

    Each streaming output must expose the user-supplied id unchanged; internal
    UUID suffixing must not leak into ``output.request_id``.
    """
    engine = AsyncOmni(model=DIFFUSION_MODEL)
    try:
        for req_id in _DIFFUSION_REQ_IDS:
            async for output in engine.generate("a white cat", request_id=req_id):
                assert output.request_id == req_id
        async for output in engine.generate("a white cat", request_id=""):
            assert output.request_id != ""
    finally:
        engine.shutdown()


@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.omni
@pytest.mark.asyncio
async def test_omni_generate_request_id():
    """Omni E2E request-id contract (``Qwen/Qwen2.5-Omni-7B``, thinker-only stage).

    Same scenarios as ``test_diffusion_generate_request_id``:
    - plain id ``my-req-1``
    - image-style prefix ``img_gen-abc123``
    - chat-style prefix ``chatcmpl-xyz``
    - empty ``request_id`` → output id is non-empty (auto-assigned)

    Text modality only; asserts caller-visible ids are preserved across the
    multi-stage orchestrator path on H100.
    """
    engine = AsyncOmni(model=OMNI_MODEL, deploy_config=OMNI_STAGE_CONFIG)
    try:
        for req_id in _OMNI_REQ_IDS:
            async for output in engine.generate(
                "Say hello in one word.",
                request_id=req_id,
                output_modalities=["text"],
            ):
                assert output.request_id == req_id
        async for output in engine.generate(
            "Say hello in one word.",
            request_id="",
            output_modalities=["text"],
        ):
            assert output.request_id != ""
    finally:
        engine.shutdown()