amy-why-3459andHongsheng Liu
01f500a5bf
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
44 lines
1.6 KiB
Python
44 lines
1.6 KiB
Python
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _make_runner(engine_output_type: str | None, downstream_req_ids: set[str]) -> GPUARModelRunner:
|
|
runner = object.__new__(GPUARModelRunner)
|
|
runner.vllm_config = SimpleNamespace(
|
|
model_config=SimpleNamespace(engine_output_type=engine_output_type),
|
|
)
|
|
runner._request_needs_downstream_stage_payload = lambda rid: rid in downstream_req_ids
|
|
return runner
|
|
|
|
|
|
def test_resolve_pooler_payload_req_ids_audio_terminal_stage_keeps_payload():
|
|
runner = _make_runner(engine_output_type="audio", downstream_req_ids=set())
|
|
|
|
engine_output_type, payload_req_ids = GPUARModelRunner._resolve_pooler_payload_req_ids(runner, ["r1", "r2"])
|
|
|
|
assert engine_output_type == "audio"
|
|
assert payload_req_ids == ["r1", "r2"]
|
|
|
|
|
|
def test_resolve_pooler_payload_req_ids_text_terminal_stage_drops_payload():
|
|
runner = _make_runner(engine_output_type="text", downstream_req_ids=set())
|
|
|
|
engine_output_type, payload_req_ids = GPUARModelRunner._resolve_pooler_payload_req_ids(runner, ["r1", "r2"])
|
|
|
|
assert engine_output_type == "text"
|
|
assert payload_req_ids == []
|
|
|
|
|
|
def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests():
|
|
runner = _make_runner(engine_output_type="latent", downstream_req_ids={"r2"})
|
|
|
|
engine_output_type, payload_req_ids = GPUARModelRunner._resolve_pooler_payload_req_ids(runner, ["r1", "r2", "r3"])
|
|
|
|
assert engine_output_type == "latent"
|
|
assert payload_req_ids == ["r2"]
|