vllm-omni/tests/entrypoints/test_async_omni.py

import asyncio
import re
from types import SimpleNamespace

import pytest
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.entrypoints.async_omni import AsyncOmni

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(**kw):
    pass


def get_fake_add_request(submitted_request_ids):
    async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
        del prompt, sampling_params_list, final_stage_id, kwargs
        submitted_request_ids.append(request_id)

    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._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 = []
        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
@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 = []
        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
        assert len(omni.request_states) == len(req_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_generate_accepts_request_after_repeated_cancellations():
    async def run_test():
        submitted_request_ids = []
        aborted_request_batches = []

        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
@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.
#
# NOTE: It seems like we currently need the shutdowns here, otherwise
# running the tests sequentially seems to leave a zombie process in diffusion
# that can OOM the Omni tests.
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.omni
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "req_id",
    ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"],
    ids=["plain", "prefixed-img", "prefixed-chat"],
)
async def test_diffusion_generate_preserves_request_id(req_id):
    """Ensure diffusion model requests don't leak internal UUID-suffixed req id back to user."""
    engine = AsyncOmni(model=DIFFUSION_MODEL)
    try:
        async for output in engine.generate("a white cat", request_id=req_id):
            assert output.request_id == req_id
    finally:
        engine.shutdown()


@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.omni
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "req_id",
    ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"],
    ids=["plain", "prefixed-img", "prefixed-chat"],
)
async def test_omni_generate_preserves_request_id(req_id):
    """Ensure omni model requests don't leak internal UUID-suffixed req id back to user."""
    engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
    try:
        async for output in engine.generate(
            "Say hello in one word.",
            request_id=req_id,
            output_modalities=["text"],
        ):
            assert output.request_id == req_id
    finally:
        engine.shutdown()


@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.omni
@pytest.mark.asyncio
async def test_diffusion_generate_empty_request_id():
    """Empty request_id should get a generated internal ID, not stay empty (diffusion)"""
    engine = AsyncOmni(model=DIFFUSION_MODEL)
    try:
        async for output in engine.generate("a white cat", request_id=""):
            assert output.request_id != ""
    finally:
        engine.shutdown()


@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.omni
@pytest.mark.asyncio
async def test_omni_generate_empty_request_id():
    """Empty request_id should get a generated internal ID, not stay empty (omni)"""
    engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
    try:
        async for output in engine.generate(
            "Say hello in one word.",
            request_id="",
            output_modalities=["text"],
        ):
            assert output.request_id != ""
    finally:
        engine.shutdown()