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(*args, **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._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 = [] 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()