52b53367ef
Signed-off-by: Daniel Huang <daniel1.huang@intel.com> Signed-off-by: Daniel Huang <pilotflyer824@gmail.com> Signed-off-by: bjf-frz <frz123db@gmail.com> Co-authored-by: Chenguang Zheng <645327136@qq.com> Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: bjf-frz <frz123db@gmail.com>
102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
"""Tests for AsyncOmniEngine.try_get_output and try_get_output_async.
|
|
|
|
Focuses on the critical behavior: when the orchestrator thread dies,
|
|
subsequent attempts to collect output raise RuntimeError.
|
|
"""
|
|
|
|
import queue
|
|
|
|
import pytest
|
|
from pytest_mock import MockerFixture
|
|
|
|
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _make_engine(output_queue, mocker: MockerFixture, *, thread_alive: bool = True) -> AsyncOmniEngine:
|
|
"""Create an AsyncOmniEngine bypassing __init__."""
|
|
engine = object.__new__(AsyncOmniEngine)
|
|
engine.output_queue = output_queue
|
|
engine.orchestrator_thread = mocker.MagicMock(
|
|
is_alive=mocker.MagicMock(return_value=thread_alive),
|
|
)
|
|
return engine
|
|
|
|
|
|
def test_try_get_output_raises_after_orchestrator_dies(mocker: MockerFixture):
|
|
"""Draining remaining results then hitting an empty queue with a dead
|
|
orchestrator must raise RuntimeError so callers know the pipeline is gone."""
|
|
mock_queue = mocker.MagicMock()
|
|
# First call succeeds; second call finds the queue empty.
|
|
mock_queue.sync_q.get.side_effect = [
|
|
{"type": "output", "request_id": "r1"},
|
|
queue.Empty,
|
|
]
|
|
|
|
engine = _make_engine(mock_queue, mocker, thread_alive=True)
|
|
|
|
# Collect the one buffered result.
|
|
assert engine.try_get_output()["request_id"] == "r1"
|
|
|
|
# Orchestrator thread crashes between polls.
|
|
engine.orchestrator_thread.is_alive.return_value = False
|
|
|
|
with pytest.raises(RuntimeError, match="Orchestrator died unexpectedly"):
|
|
engine.try_get_output()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_try_get_output_async_raises_after_orchestrator_dies(mocker: MockerFixture):
|
|
"""Same scenario as above but for the async variant."""
|
|
mock_queue = mocker.MagicMock()
|
|
mock_queue.sync_q.get_nowait.side_effect = [
|
|
{"type": "output", "request_id": "r1"},
|
|
queue.Empty,
|
|
]
|
|
|
|
engine = _make_engine(mock_queue, mocker, thread_alive=True)
|
|
|
|
assert (await engine.try_get_output_async())["request_id"] == "r1"
|
|
|
|
engine.orchestrator_thread.is_alive.return_value = False
|
|
|
|
with pytest.raises(RuntimeError, match="Orchestrator died unexpectedly"):
|
|
await engine.try_get_output_async()
|
|
|
|
|
|
def test_fatal_error_message_surfaces_through_try_get_output(mocker: MockerFixture):
|
|
"""When the orchestrator thread crashes, it enqueues a fatal error message.
|
|
|
|
``try_get_output`` must return this message so the caller
|
|
(``OmniBase._handle_output_message``) can detect the fatal flag.
|
|
"""
|
|
fatal_msg = {"type": "error", "error": "Orchestrator thread crashed", "fatal": True}
|
|
|
|
mock_queue = mocker.MagicMock()
|
|
mock_queue.sync_q.get.return_value = fatal_msg
|
|
|
|
engine = _make_engine(mock_queue, mocker, thread_alive=False)
|
|
|
|
msg = engine.try_get_output()
|
|
assert msg is not None
|
|
assert msg["type"] == "error"
|
|
assert msg["fatal"] is True
|
|
assert "crashed" in msg["error"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fatal_error_message_surfaces_through_try_get_output_async(mocker: MockerFixture):
|
|
"""Async variant of the fatal error message test."""
|
|
fatal_msg = {"type": "error", "error": "Orchestrator thread crashed", "fatal": True}
|
|
|
|
mock_queue = mocker.MagicMock()
|
|
mock_queue.sync_q.get_nowait.return_value = fatal_msg
|
|
|
|
engine = _make_engine(mock_queue, mocker, thread_alive=False)
|
|
|
|
msg = await engine.try_get_output_async()
|
|
assert msg is not None
|
|
assert msg["type"] == "error"
|
|
assert msg["fatal"] is True
|