141 lines
5.3 KiB
Python
141 lines
5.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Entrypoint contract tests for ``AsyncOmni.submit_interaction_async``."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
|
from vllm_omni.entrypoints.client_request_state import ClientRequestState
|
|
from vllm_omni.inputs.data import OmniInteractionPrompt
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _make_async_omni(*, num_stages: int = 1, stage_type: str = "diffusion") -> AsyncOmni:
|
|
omni = object.__new__(AsyncOmni)
|
|
omni.log_stats = False
|
|
omni.request_states = {
|
|
"external-abc-uuid-1": ClientRequestState(
|
|
request_id="external-abc-uuid-1",
|
|
external_request_id="external-abc",
|
|
queue=AsyncMock(),
|
|
),
|
|
"external-abc-uuid-2": ClientRequestState(
|
|
request_id="external-abc-uuid-2",
|
|
external_request_id="external-abc",
|
|
queue=AsyncMock(),
|
|
),
|
|
"other-req-uuid": ClientRequestState(
|
|
request_id="other-req-uuid",
|
|
external_request_id="other-req",
|
|
queue=AsyncMock(),
|
|
),
|
|
}
|
|
omni.engine = SimpleNamespace( # pyright: ignore[reportAttributeAccessIssue]
|
|
num_stages=num_stages,
|
|
get_stage_metadata=lambda stage_id: SimpleNamespace(stage_type=stage_type),
|
|
submit_interaction_async=AsyncMock(),
|
|
)
|
|
return omni
|
|
|
|
|
|
def _prompt_interaction(prompt: str = "new prompt", transition_chunks: int | None = None) -> OmniInteractionPrompt:
|
|
interaction: OmniInteractionPrompt = {"event_id": "ui-update-1", "event": {"prompt": prompt}}
|
|
if transition_chunks is not None:
|
|
interaction["transition_chunks"] = transition_chunks
|
|
return interaction
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_interaction_async_maps_external_to_internal_id() -> None:
|
|
omni = _make_async_omni()
|
|
with pytest.raises(ValueError, match="exactly one active request"):
|
|
await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction(transition_chunks=2))
|
|
|
|
omni.request_states.pop("external-abc-uuid-2")
|
|
await omni.submit_interaction_async(
|
|
"external-abc",
|
|
interaction=_prompt_interaction(transition_chunks=2),
|
|
)
|
|
omni.engine.submit_interaction_async.assert_awaited_once_with( # pyright: ignore[reportAttributeAccessIssue]
|
|
"external-abc-uuid-1",
|
|
interaction=_prompt_interaction(transition_chunks=2),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_interaction_async_passes_missing_transition_chunks() -> None:
|
|
omni = _make_async_omni()
|
|
omni.request_states.pop("external-abc-uuid-2")
|
|
await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction())
|
|
omni.engine.submit_interaction_async.assert_awaited_once_with( # pyright: ignore[reportAttributeAccessIssue]
|
|
"external-abc-uuid-1",
|
|
interaction=_prompt_interaction(),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_interaction_async_rejects_empty_prompt() -> None:
|
|
omni = _make_async_omni()
|
|
omni.request_states.pop("external-abc-uuid-2")
|
|
with pytest.raises(ValueError, match="prompt must be non-empty"):
|
|
await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction(""))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_interaction_async_rejects_negative_transition_chunks() -> None:
|
|
omni = _make_async_omni()
|
|
omni.request_states.pop("external-abc-uuid-2")
|
|
with pytest.raises(ValueError, match="transition_chunks must be >= 0"):
|
|
await omni.submit_interaction_async(
|
|
"external-abc",
|
|
interaction=_prompt_interaction(transition_chunks=-1),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_interaction_async_rejects_inactive_request() -> None:
|
|
omni = _make_async_omni()
|
|
omni.request_states.clear()
|
|
with pytest.raises(ValueError, match="No active request"):
|
|
await omni.submit_interaction_async("missing", interaction=_prompt_interaction(transition_chunks=2))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_interaction_async_rejects_non_diffusion() -> None:
|
|
omni = _make_async_omni(stage_type="llm")
|
|
omni.request_states = {
|
|
"req-uuid": ClientRequestState(
|
|
request_id="req-uuid",
|
|
external_request_id="req",
|
|
queue=AsyncMock(),
|
|
),
|
|
}
|
|
with pytest.raises(ValueError, match="requires a diffusion stage"):
|
|
await omni.submit_interaction_async("req", interaction=_prompt_interaction(transition_chunks=2))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_interaction_async_forwards_unsupported_dict_shape_to_engine() -> None:
|
|
omni = _make_async_omni()
|
|
omni.request_states.pop("external-abc-uuid-2")
|
|
payload: OmniInteractionPrompt = { # pyright: ignore[reportAssignmentType] # deliberately wrong
|
|
"event": {"multi_modal_data": {"unsupported_modality": {"type": "pose"}}},
|
|
"transition_chunks": 2,
|
|
}
|
|
|
|
await omni.submit_interaction_async(
|
|
"external-abc",
|
|
interaction=payload,
|
|
)
|
|
|
|
omni.engine.submit_interaction_async.assert_awaited_once_with( # pyright: ignore[reportAttributeAccessIssue]
|
|
"external-abc-uuid-1",
|
|
interaction=payload,
|
|
)
|