303 lines
10 KiB
Python
303 lines
10 KiB
Python
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()
|