anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Entrypoint][Refactor] Make field type hint more concrete (#3139)

Browse Source 
Signed-off-by: wuhang <wuhang6@huawei.com>
This commit is contained in:
wuhang authored and GitHub committed 2026-05-14 11:25:55 +08:00
1 parent 51b4b1131e
commit 0d9d57acd9
20 files changed
+908 -543
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/engine/test_async_omni_engine_input.py
+5 -5
View File
@@ -4,7 +4,7 @@ from vllm.sampling_params import SamplingParams
from vllm.v1.engine import EngineCoreRequest
from vllm_omni.engine import OmniEngineCoreRequest
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -27,7 +27,7 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
engine = object.__new__(AsyncOmniEngine)
params = SamplingParams(max_tokens=8)
engine.default_sampling_params_list = [params]
engine.stage_metadata = [{"stage_type": "llm"}]
engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
engine.supported_tasks = ("speech",)
input_processor = mocker.Mock()
@@ -53,7 +53,7 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
arrival_time=0.0,
)
request = msg["prompt"]
request = msg.prompt
assert isinstance(request, OmniEngineCoreRequest)
assert request.external_req_id == "req-1"
assert request.additional_information is not None
@@ -66,7 +66,7 @@ def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixtur
engine = object.__new__(AsyncOmniEngine)
params = SamplingParams(max_tokens=8)
engine.default_sampling_params_list = [params]
engine.stage_metadata = [{"stage_type": "llm"}]
engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
engine.supported_tasks = ("generate",)
input_processor = mocker.Mock()
@@ -85,6 +85,6 @@ def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixtur
message_type="streaming_update",
)
assert msg["type"] == "streaming_update"
assert msg.type == "streaming_update"
input_processor.process_inputs.assert_called_once()
assert input_processor.process_inputs.call_args.kwargs["resumable"] is True
tests/engine/test_async_omni_engine_outputs.py
+23 -11
View File
@@ -10,6 +10,8 @@ import pytest
from pytest_mock import MockerFixture
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -30,14 +32,19 @@ def test_try_get_output_raises_after_orchestrator_dies(mocker: MockerFixture):
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"},
OutputMessage(
request_id="r1",
stage_id=0,
engine_outputs=OmniRequestOutput(request_id="r1"),
finished=False,
),
queue.Empty,
]
engine = _make_engine(mock_queue, mocker, thread_alive=True)
# Collect the one buffered result.
assert engine.try_get_output()["request_id"] == "r1"
assert engine.try_get_output().request_id == "r1"
# Orchestrator thread crashes between polls.
engine.orchestrator_thread.is_alive.return_value = False
@@ -51,13 +58,18 @@ async def test_try_get_output_async_raises_after_orchestrator_dies(mocker: Mocke
"""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"},
OutputMessage(
request_id="r1",
stage_id=0,
engine_outputs=OmniRequestOutput(request_id="r1"),
finished=False,
),
queue.Empty,
]
engine = _make_engine(mock_queue, mocker, thread_alive=True)
assert (await engine.try_get_output_async())["request_id"] == "r1"
assert (await engine.try_get_output_async()).request_id == "r1"
engine.orchestrator_thread.is_alive.return_value = False
@@ -71,7 +83,7 @@ def test_fatal_error_message_surfaces_through_try_get_output(mocker: MockerFixtu
``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}
fatal_msg = ErrorMessage(error="Orchestrator thread crashed", fatal=True)
mock_queue = mocker.MagicMock()
mock_queue.sync_q.get.return_value = fatal_msg
@@ -80,15 +92,15 @@ def test_fatal_error_message_surfaces_through_try_get_output(mocker: MockerFixtu
msg = engine.try_get_output()
assert msg is not None
assert msg["type"] == "error"
assert msg["fatal"] is True
assert "crashed" in msg["error"]
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}
fatal_msg = ErrorMessage(error="Orchestrator thread crashed", fatal=True)
mock_queue = mocker.MagicMock()
mock_queue.sync_q.get_nowait.return_value = fatal_msg
@@ -97,5 +109,5 @@ async def test_fatal_error_message_surfaces_through_try_get_output_async(mocker:
msg = await engine.try_get_output_async()
assert msg is not None
assert msg["type"] == "error"
assert msg["fatal"] is True
assert msg.type == "error"
assert msg.fatal is True
tests/engine/test_async_omni_engine_stage_init.py
+3 -3
View File
@@ -7,7 +7,7 @@ import types
import pytest
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
from vllm_omni.engine.stage_init_utils import (
LogicalStageInitPlan,
ReplicaInitPlan,
@@ -332,8 +332,8 @@ def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_inpu
stage1_client_r0.default_sampling_params,
]
assert engine.stage_metadata == [
{"final_output": False, "final_output_type": None, "stage_type": "llm"},
{"final_output": True, "final_output_type": None, "stage_type": "llm"},
StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm"),
StageRuntimeInfo(final_output=True, final_output_type=None, stage_type="llm"),
]
tests/engine/test_orchestrator.py
+116 -63
View File
@@ -17,6 +17,15 @@ from vllm.outputs import CompletionOutput, RequestOutput
from vllm.sampling_params import SamplingParams
from vllm.v1.engine.core_client import AsyncMPClient
from vllm_omni.engine.messages import (
AbortRequestMessage,
AddCompanionRequestMessage,
CollectiveRPCRequestMessage,
CollectiveRPCResultMessage,
OutputMessage,
ShutdownRequestMessage,
StageSubmissionMessage,
)
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
from vllm_omni.engine.stage_pool import StagePool
@@ -44,14 +53,27 @@ class FakeStageClient:
final_output: bool = False,
final_output_type: str = "text",
next_inputs: list[dict] | None = None,
engine_input_source: list[int] | None = None,
is_comprehension: bool = False,
model_stage: str | None = None,
kv_sender_info: dict[str, Any] | None = None,
) -> None:
self.stage_id = 0
self.replica_id = 0
self.stage_type = stage_type
self.final_output = final_output
self.final_output_type = final_output_type
self.default_sampling_params = SamplingParams(max_tokens=1)
self.requires_multimodal_data = False
self.engine_input_source = list(engine_input_source or [0])
self.is_comprehension = is_comprehension
self.model_stage = model_stage
self.next_inputs = list(next_inputs or [])
self.custom_process_input_func = None
self._kv_sender_info = dict(kv_sender_info) if kv_sender_info is not None else None
self.add_request_calls: list[tuple] = []
self.abort_calls: list[list[str]] = []
self.collective_rpc_calls: list[tuple[str, float | None, tuple[Any, ...], dict[str, Any]]] = []
self.shutdown_calls = 0
self._engine_core_outputs = queue.Queue()
self._diffusion_outputs = queue.Queue()
@@ -81,6 +103,30 @@ class FakeStageClient:
async def abort_requests_async(self, request_ids: list[str]) -> None:
self.abort_calls.append(list(request_ids))
async def collective_rpc_async(
self,
*,
method: str,
timeout: float | None = None,
args: tuple[Any, ...] = (),
kwargs: dict[str, Any] | None = None,
) -> Any:
normalized_kwargs = dict(kwargs or {})
self.collective_rpc_calls.append((method, timeout, args, normalized_kwargs))
return {
"supported": False,
"todo": True,
"reason": f"{self.__class__.__name__}.collective_rpc_async is not implemented yet",
}
def get_kv_sender_info(self) -> dict[str, Any] | None:
if self._kv_sender_info is None:
return None
return dict(self._kv_sender_info)
def check_health(self) -> None:
return None
def shutdown(self) -> None:
self.shutdown_calls += 1
@@ -96,7 +142,6 @@ class FakeCollectiveRpcStageClient(FakeStageClient):
def __init__(self, *args, rpc_result: Any = None, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.rpc_result = rpc_result
self.collective_rpc_calls: list[tuple[str, float | None, tuple[Any, ...], dict[str, Any]]] = []
async def collective_rpc_async(
self,
@@ -318,7 +363,7 @@ def _build_harness(
async def _shutdown_orchestrator(orchestrator_fixture: OrchestratorFixture) -> None:
orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
await asyncio.to_thread(orchestrator_fixture.thread.join, 5)
if orchestrator_fixture.thread.is_alive():
raise AssertionError("Timed out waiting for orchestrator thread shutdown")
@@ -333,7 +378,7 @@ async def _wait_for(predicate, *, timeout: float = 2.0) -> None:
await asyncio.sleep(0.01)
async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> OutputMessage:
deadline = time.monotonic() + timeout
while True:
if time.monotonic() >= deadline:
@@ -343,11 +388,15 @@ async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, time
except queue.Empty:
await asyncio.sleep(0.01)
continue
if msg.get("type") == "output":
if isinstance(msg, OutputMessage):
return msg
async def _get_rpc_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
async def _get_rpc_message(
orchestrator_fixture: OrchestratorFixture,
*,
timeout: float = 2.0,
) -> CollectiveRPCResultMessage:
deadline = time.monotonic() + timeout
rpc_sync_q = orchestrator_fixture.queues[2].sync_q
while True:
@@ -369,24 +418,22 @@ async def _enqueue_add_request(
final_stage_id: int,
) -> None:
orchestrator_fixture.request_sync_q.put_nowait(
{
"type": "add_request",
"request_id": request_id,
"prompt": prompt,
"original_prompt": original_prompt,
"sampling_params_list": sampling_params_list,
"final_stage_id": final_stage_id,
}
StageSubmissionMessage(
type="add_request",
request_id=request_id,
prompt=prompt,
original_prompt=original_prompt,
output_prompt_text=None,
sampling_params_list=sampling_params_list,
final_stage_id=final_stage_id,
preprocess_ms=0.0,
enqueue_ts=time.perf_counter(),
)
)
async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, request_ids: list[str]) -> None:
orchestrator_fixture.request_sync_q.put_nowait(
{
"type": "abort",
"request_ids": request_ids,
}
)
orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=request_ids))
def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fails(monkeypatch):
@@ -457,7 +504,7 @@ def orchestrator_factory():
for fixture in fixtures:
if fixture.thread.is_alive():
fixture.request_sync_q.put_nowait({"type": "shutdown"})
fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
fixture.thread.join(timeout=5)
for q in fixture.queues:
q.close()
@@ -505,10 +552,10 @@ async def test_run_two_stage_llm(orchestrator_factory) -> None:
output_msg = await _get_output_message(orchestrator_fixture)
assert output_msg["request_id"] == "req-llm"
assert output_msg["stage_id"] == 1
assert output_msg["finished"] is True
assert output_msg["engine_outputs"].request_id == "req-llm"
assert output_msg.request_id == "req-llm"
assert output_msg.stage_id == 1
assert output_msg.finished is True
assert output_msg.engine_outputs.request_id == "req-llm"
assert "req-llm" not in orchestrator_fixture.orchestrator.request_states
finally:
await _shutdown_orchestrator(orchestrator_fixture)
@@ -541,10 +588,10 @@ async def test_run_single_stage_diffusion(orchestrator_factory) -> None:
output_msg = await _get_output_message(orchestrator_fixture)
assert output_msg["request_id"] == "req-diff"
assert output_msg["stage_id"] == 0
assert output_msg["finished"] is True
assert output_msg["engine_outputs"].request_id == "req-diff"
assert output_msg.request_id == "req-diff"
assert output_msg.stage_id == 0
assert output_msg.finished is True
assert output_msg.engine_outputs.request_id == "req-diff"
assert "req-diff" not in orchestrator_fixture.orchestrator.request_states
finally:
await _shutdown_orchestrator(orchestrator_fixture)
@@ -589,10 +636,10 @@ async def test_run_llm_to_diffusion(orchestrator_factory) -> None:
output_msg = await _get_output_message(orchestrator_fixture)
assert output_msg["request_id"] == "req-img"
assert output_msg["stage_id"] == 1
assert output_msg["finished"] is True
assert output_msg["engine_outputs"].request_id == "req-img"
assert output_msg.request_id == "req-img"
assert output_msg.stage_id == 1
assert output_msg.finished is True
assert output_msg.engine_outputs.request_id == "req-img"
assert "req-img" not in orchestrator_fixture.orchestrator.request_states
finally:
await _shutdown_orchestrator(orchestrator_fixture)
@@ -633,9 +680,9 @@ async def test_run_async_chunk(orchestrator_factory) -> None:
output_msg = await _get_output_message(orchestrator_fixture)
assert output_msg["request_id"] == "req-async"
assert output_msg["stage_id"] == 1
assert output_msg["finished"] is True
assert output_msg.request_id == "req-async"
assert output_msg.stage_id == 1
assert output_msg.finished is True
assert "req-async" not in orchestrator_fixture.orchestrator.request_states
finally:
await _shutdown_orchestrator(orchestrator_fixture)
@@ -757,9 +804,9 @@ async def test_multi_replica_round_robin_distribution(orchestrator_factory) -> N
stage1.push_engine_core_outputs(_engine_core_outputs("s1-raw", 2.0))
output_msg = await _get_output_message(orchestrator_fixture)
assert output_msg["request_id"] == "req-0"
assert output_msg["stage_id"] == 1
assert output_msg["finished"] is True
assert output_msg.request_id == "req-0"
assert output_msg.stage_id == 1
assert output_msg.finished is True
assert "req-0" not in orchestrator_fixture.orchestrator.request_states
finally:
await _shutdown_orchestrator(orchestrator_fixture)
@@ -927,12 +974,17 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non
orchestrator.stage_pools = [pool]
await orchestrator._handle_streaming_update(
{
"request_id": "req-stream",
"prompt": SimpleNamespace(request_id="req-stream", prompt_token_ids=[1], resumable=True),
"sampling_params_list": [_sampling_params()],
"output_prompt_text": "segment-2",
}
StageSubmissionMessage(
type="streaming_update",
request_id="req-stream",
prompt=SimpleNamespace(request_id="req-stream", prompt_token_ids=[1], resumable=True),
original_prompt={"prompt": "segment-2"},
output_prompt_text="segment-2",
sampling_params_list=[_sampling_params()],
final_stage_id=0,
preprocess_ms=0.0,
enqueue_ts=time.perf_counter(),
)
)
assert pool.calls == [("req-stream", "segment-2")]
@@ -1024,12 +1076,14 @@ async def test_collective_rpc_ignores_invalid_stage_ids(orchestrator_factory, ca
target_logger.setLevel(logging.WARNING)
try:
orchestrator_fixture.request_sync_q.put_nowait(
{
"type": "collective_rpc",
"rpc_id": "rpc-1",
"method": "list_loras",
"stage_ids": [99, 1],
}
CollectiveRPCRequestMessage(
rpc_id="rpc-1",
method="list_loras",
timeout=None,
args=(),
kwargs={},
stage_ids=[99, 1],
)
)
msg = await _get_rpc_message(orchestrator_fixture)
@@ -1037,10 +1091,10 @@ async def test_collective_rpc_ignores_invalid_stage_ids(orchestrator_factory, ca
target_logger.removeHandler(caplog.handler)
target_logger.setLevel(prev_level)
assert msg["type"] == "collective_rpc_result"
assert msg["rpc_id"] == "rpc-1"
assert msg["stage_ids"] == [1]
assert msg["results"] == [{"stage": 1}]
assert msg.type == "collective_rpc_result"
assert msg.rpc_id == "rpc-1"
assert msg.stage_ids == [1]
assert msg.results == [{"stage": 1}]
assert not stage0.collective_rpc_calls
assert len(stage1.collective_rpc_calls) == 1
assert "collective_rpc: ignoring invalid stage_id 99" in caplog.text
@@ -1084,15 +1138,14 @@ async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator
await _wait_for(lambda: len(stage0_r1.add_request_calls) == 1)
orchestrator_fixture.request_sync_q.put_nowait(
{
"type": "add_companion_request",
"companion_id": "parent-neg",
"parent_id": "parent",
"role": "negative",
"prompt": SimpleNamespace(request_id="parent-neg", prompt_token_ids=[9]),
"companion_prompt_text": {"prompt": "negative"},
"sampling_params_list": [_sampling_params()],
}
AddCompanionRequestMessage(
companion_id="parent-neg",
parent_id="parent",
role="negative",
prompt=SimpleNamespace(request_id="parent-neg", prompt_token_ids=[9]),
companion_prompt_text={"prompt": "negative"},
sampling_params_list=[_sampling_params()],
)
)
await _wait_for(lambda: len(stage0_r1.add_request_calls) == 2)
tests/engine/test_orchestrator_error_handling.py
+15 -12
View File
@@ -17,6 +17,7 @@ from types import SimpleNamespace
import pytest
from vllm.v1.engine.exceptions import EngineDeadError
from vllm_omni.engine.messages import EngineQueueMessage, ErrorMessage, ShutdownRequestMessage
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -37,7 +38,7 @@ def _sampling_params(max_tokens: int = 4):
return SamplingParams(max_tokens=max_tokens)
async def _get_any_output_message(fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
async def _get_any_output_message(fixture: OrchestratorFixture, *, timeout: float = 2.0) -> EngineQueueMessage:
"""Like _get_output_message but returns any message type (including errors)."""
deadline = time.monotonic() + timeout
while True:
@@ -62,7 +63,7 @@ def orchestrator_factory():
for fixture in fixtures:
if fixture.thread.is_alive():
fixture.request_sync_q.put_nowait({"type": "shutdown"})
fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
fixture.thread.join(timeout=5)
for q in fixture.queues:
q.close()
@@ -101,10 +102,11 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa
# Collect the fatal error message.
msg = await _get_any_output_message(orchestrator_fixture)
assert msg["type"] == "error"
assert msg["fatal"] is True
assert msg["request_id"] == "req-dead"
assert "Stage-0 engine core is dead" in msg["error"]
assert isinstance(msg, ErrorMessage)
assert msg.type == "error"
assert msg.fatal is True
assert msg.request_id == "req-dead"
assert "Stage-0 engine core is dead" in msg.error
# The orchestrator thread should exit after the fatal error.
orchestrator_fixture.thread.join(timeout=5)
@@ -114,7 +116,7 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa
assert "req-dead" not in orchestrator_fixture.orchestrator.request_states
finally:
if orchestrator_fixture.thread.is_alive():
orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
orchestrator_fixture.thread.join(timeout=5)
@@ -148,13 +150,14 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -
msg = await _get_any_output_message(orchestrator_fixture)
assert msg["type"] == "error"
assert msg["request_id"] == "req-err"
assert msg["stage_id"] == 0
assert msg["error"] == "gpu fault"
assert isinstance(msg, ErrorMessage)
assert msg.type == "error"
assert msg.request_id == "req-err"
assert msg.stage_id == 0
assert msg.error == "gpu fault"
# Request state should be cleaned up.
await _wait_for(lambda: "req-err" not in orchestrator_fixture.orchestrator.request_states)
finally:
orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
orchestrator_fixture.thread.join(timeout=5)
tests/engine/test_orchestrator_kv_sender_info.py
+7 -5
View File
@@ -5,6 +5,7 @@ import pytest
from vllm import SamplingParams
from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
from vllm_omni.engine.messages import OutputMessage
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
from vllm_omni.engine.stage_pool import StagePool
@@ -225,11 +226,12 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
assert diffusion_stage.calls == []
assert len(orchestrator.output_async_queue.items) == 1
terminal_msg = orchestrator.output_async_queue.items[0]
assert terminal_msg["type"] == "output"
assert terminal_msg["request_id"] == "req-empty"
assert terminal_msg["stage_id"] == 1
assert terminal_msg["finished"] is True
assert "produced no valid inputs" in terminal_msg["engine_outputs"].error
assert isinstance(terminal_msg, OutputMessage)
assert terminal_msg.type == "output"
assert terminal_msg.request_id == "req-empty"
assert terminal_msg.stage_id == 1
assert terminal_msg.finished is True
assert "produced no valid inputs" in terminal_msg.engine_outputs.error
assert "req-empty" not in orchestrator.request_states
tests/entrypoints/test_omni_base_profiler.py
+6 -4
View File
@@ -1,5 +1,7 @@
"""Unit tests for OmniBase and AsyncOmni profiler methods."""
from types import SimpleNamespace
import pytest
from pytest_mock import MockerFixture
@@ -16,10 +18,10 @@ class TestOmniBaseProfiler:
engine.num_stages = 3
engine.is_alive.return_value = True
engine.default_sampling_params_list = [mocker.MagicMock() for _ in range(3)]
engine.get_stage_metadata.side_effect = lambda i: {
"final_output_type": "text" if i == 0 else "audio",
"final_output": True,
}
engine.get_stage_metadata.side_effect = lambda i: SimpleNamespace(
final_output_type="text" if i == 0 else "audio",
final_output=True,
)
engine.collective_rpc.return_value = [None, None, None]
return engine
tests/entrypoints/test_omni_entrypoints.py
+60 -67
View File
@@ -13,44 +13,40 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _stage_meta(*, stage_type: str, final_output: bool, final_output_type: str | None) -> StageRuntimeInfo:
return StageRuntimeInfo(
stage_type=stage_type,
final_output=final_output,
final_output_type=final_output_type,
)
THREE_STAGE_META = [
{"stage_type": "llm", "final_output": True, "final_output_type": "text"},
{"stage_type": "llm", "final_output": False, "final_output_type": None},
{"stage_type": "diffusion", "final_output": True, "final_output_type": "image"},
_stage_meta(stage_type="llm", final_output=True, final_output_type="text"),
_stage_meta(stage_type="llm", final_output=False, final_output_type=None),
_stage_meta(stage_type="diffusion", final_output=True, final_output_type="image"),
]
DIFFUSION_ONLY_META = [
{"stage_type": "diffusion", "final_output": True, "final_output_type": "image"},
_stage_meta(stage_type="diffusion", final_output=True, final_output_type="image"),
]
LLM_DIFFUSION_META = [
{"stage_type": "llm", "final_output": True, "final_output_type": "text"},
{"stage_type": "diffusion", "final_output": True, "final_output_type": "image"},
_stage_meta(stage_type="llm", final_output=True, final_output_type="text"),
_stage_meta(stage_type="diffusion", final_output=True, final_output_type="image"),
]
class FakeEngineOutput:
def __init__(
self,
*,
payload: str,
finished: bool,
images: list[str] | None = None,
stage_durations: dict[str, float] | None = None,
) -> None:
self.payload = payload
self.finished = finished
self.images = images or []
self.stage_durations = stage_durations or {}
def make_output_msg(
request_id: str,
stage_id: int,
@@ -60,21 +56,25 @@ def make_output_msg(
finished: bool | None = None,
images: list[str] | None = None,
metrics: Any = None,
) -> dict[str, Any]:
) -> OutputMessage:
if finished is None:
finished = output_finished
return {
"type": "output",
"request_id": request_id,
"stage_id": stage_id,
"engine_outputs": FakeEngineOutput(
payload=payload,
finished=output_finished,
images=images,
),
"finished": finished,
"metrics": metrics,
}
final_output_type = "image" if images else "text"
engine_output = OmniRequestOutput(
request_id=request_id,
finished=output_finished,
final_output_type=final_output_type,
images=images or [],
stage_durations={},
)
engine_output.payload = payload
return OutputMessage(
request_id=request_id,
stage_id=stage_id,
engine_outputs=engine_output,
finished=finished,
metrics=metrics,
)
class FakeAsyncOmniEngine:
@@ -82,7 +82,7 @@ class FakeAsyncOmniEngine:
self,
model: str = "dummy-model",
*,
stage_metadata: list[dict[str, Any]] | None = None,
stage_metadata: list[StageRuntimeInfo] | None = None,
default_sampling_params_list: list[Any] | None = None,
on_add_request: Callable[[FakeAsyncOmniEngine, dict[str, Any]], None] | None = None,
rpc_results: list[Any] | None = None,
@@ -102,7 +102,7 @@ class FakeAsyncOmniEngine:
self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)]
self.input_processor = None
self.output_q: queue.Queue[dict[str, Any]] = queue.Queue()
self.output_q: queue.Queue[Any] = queue.Queue()
self.submitted: list[dict[str, Any]] = []
self.aborted: list[list[str]] = []
self.rpc_results = rpc_results or []
@@ -133,16 +133,16 @@ class FakeAsyncOmniEngine:
async def add_request_async(self, *args, **kwargs) -> None:
self.add_request(*args, **kwargs)
def try_get_output(self, timeout: float = 0.001) -> dict[str, Any] | None:
def try_get_output(self, timeout: float = 0.001) -> Any | None:
try:
return self.output_q.get_nowait()
except queue.Empty:
return None
async def try_get_output_async(self) -> dict[str, Any] | None:
async def try_get_output_async(self) -> Any | None:
return self.try_get_output()
def get_stage_metadata(self, stage_id: int) -> dict[str, Any]:
def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
return self.stage_metadata[stage_id]
def abort(self, request_ids: list[str]) -> None:
@@ -405,24 +405,22 @@ def _enqueue_async_llm_diffusion_outputs(engine: FakeAsyncOmniEngine, msg: dict[
def _enqueue_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
engine.output_q.put_nowait(
{
"type": "error",
"request_id": msg["request_id"],
"stage_id": 0,
"error": "engine boom",
}
ErrorMessage(
request_id=msg["request_id"],
stage_id=0,
error="engine boom",
)
)
def _enqueue_fatal_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
engine.output_q.put_nowait(
{
"type": "error",
"fatal": True,
"request_id": msg["request_id"],
"stage_id": 2,
"error": "engine dead",
}
ErrorMessage(
fatal=True,
request_id=msg["request_id"],
stage_id=2,
error="engine dead",
)
)
@@ -798,7 +796,7 @@ def test_omni_forces_final_only_on_llm_stages(monkeypatch: pytest.MonkeyPatch):
def test_fatal_error_raises_engine_dead():
base = _make_base()
msg = {"type": "error", "error": "orchestrator crashed", "fatal": True}
msg = ErrorMessage(error="orchestrator crashed", fatal=True)
with pytest.raises(EngineDeadError, match="orchestrator crashed"):
base._handle_output_message(msg)
@@ -806,7 +804,7 @@ def test_fatal_error_raises_engine_dead():
def test_non_fatal_error_raises_runtime():
base = _make_base()
msg = {"type": "error", "error": "something wrong"}
msg = ErrorMessage(error="something wrong")
with pytest.raises(RuntimeError, match="something wrong"):
base._handle_output_message(msg)
@@ -853,20 +851,15 @@ def _enqueue_stage_error(
"""Enqueue a stage error output, optionally killing the engine."""
if kill_engine:
engine._alive = False
engine_output = OmniRequestOutput.from_error(msg["request_id"], error_text)
engine_output.payload = ""
engine.output_q.put_nowait(
{
"type": "output",
"request_id": msg["request_id"],
"stage_id": 0,
"engine_outputs": SimpleNamespace(
payload="",
finished=True,
images=[],
stage_durations={},
error=error_text,
),
"finished": False,
}
OutputMessage(
request_id=msg["request_id"],
stage_id=0,
engine_outputs=engine_output,
finished=False,
)
)
vllm_omni/diffusion/inline_stage_diffusion_client.py
+3 -1
View File
@@ -20,6 +20,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
from vllm_omni.diffusion.data import DiffusionRequestAbortedError
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.engine.stage_client import StageClientBase
from vllm_omni.engine.stage_init_utils import StageMetadata
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -31,11 +32,12 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
class InlineStageDiffusionClient:
class InlineStageDiffusionClient(StageClientBase):
"""Runs DiffusionEngine in a thread executor inside the Orchestrator."""
stage_type: str = "diffusion"
replica_id: int = 0
is_comprehension: bool = False
def __init__(
self,
vllm_omni/diffusion/stage_diffusion_client.py
+3 -1
View File
@@ -29,6 +29,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackDecoder,
OmniMsgpackEncoder,
)
from vllm_omni.engine.stage_client import StageClientBase
from vllm_omni.engine.stage_init_utils import StageMetadata, terminate_alive_proc
from vllm_omni.outputs import OmniRequestOutput
@@ -57,7 +58,7 @@ def create_diffusion_client(
)
class StageDiffusionClient:
class StageDiffusionClient(StageClientBase):
"""Communicates with StageDiffusionProc via ZMQ for use inside the Orchestrator.
Exposes the same attributes and async methods the Orchestrator
@@ -68,6 +69,7 @@ class StageDiffusionClient:
stage_type: str = "diffusion"
replica_id: int = 0
is_comprehension: bool = False
def __init__(
self,
vllm_omni/engine/async_omni_engine.py
+107 -90
View File
@@ -20,8 +20,8 @@ import uuid
import weakref
from collections.abc import Mapping, Sequence
from contextlib import ExitStack
from dataclasses import asdict
from typing import TYPE_CHECKING, Any
from dataclasses import asdict, dataclass
from typing import TYPE_CHECKING, Any, Literal, cast
import janus
import torch
@@ -36,6 +36,7 @@ from vllm.v1.engine.input_processor import InputProcessor
from vllm_omni.config.stage_config import strip_parent_engine_args
from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
from vllm_omni.diffusion.diffusion_engine import supports_audio_output
from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
from vllm_omni.diffusion.stage_diffusion_proc import (
complete_diffusion_handshake,
@@ -45,11 +46,23 @@ from vllm_omni.distributed.omni_connectors.utils.initialization import (
resolve_omni_kv_config_for_stage,
)
from vllm_omni.engine import OmniEngineCoreRequest
from vllm_omni.engine.messages import (
AbortRequestMessage,
AddCompanionRequestMessage,
CollectiveRPCRequestMessage,
CollectiveRPCResultMessage,
EngineQueueMessage,
ErrorMessage,
ShutdownRequestMessage,
StageSubmissionMessage,
)
from vllm_omni.engine.orchestrator import Orchestrator
from vllm_omni.engine.output_modality import FinalOutputModalityType
from vllm_omni.engine.serialization import (
deserialize_additional_information,
serialize_additional_information,
)
from vllm_omni.engine.stage_client import StageClient
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
from vllm_omni.engine.stage_engine_core_proc import (
complete_stage_handshake,
@@ -83,12 +96,13 @@ from vllm_omni.engine.stage_init_utils import (
setup_stage_devices,
terminate_alive_proc,
)
from vllm_omni.engine.stage_pool import StagePool
from vllm_omni.engine.stage_pool import StagePool, StagePoolClient
from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
from vllm_omni.entrypoints.utils import (
inject_omni_kv_config,
load_and_resolve_stage_configs,
)
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.platforms import current_omni_platform
if TYPE_CHECKING:
@@ -99,6 +113,13 @@ logger = init_logger(__name__)
_STARTUP_POLL_INTERVAL_S = 1.0
@dataclass(frozen=True, slots=True)
class StageRuntimeInfo:
final_output: bool
final_output_type: FinalOutputModalityType | None
stage_type: str
# ============================================================================
# Parent-EngineArgs field-routing contracts (consumed by
# AsyncOmniEngine._strip_parent_engine_args when ``stage_configs_path`` is set).
@@ -190,14 +211,14 @@ def _apply_omni_final_stage_metadata(
def _weak_shutdown_async_omni_engine(
orchestrator_thread: threading.Thread | None,
request_queue: janus.Queue[dict[str, Any]] | None,
output_queue: janus.Queue[dict[str, Any]] | None,
rpc_output_queue: janus.Queue[dict[str, Any]] | None,
request_queue: janus.Queue[EngineQueueMessage] | None,
output_queue: janus.Queue[EngineQueueMessage] | None,
rpc_output_queue: janus.Queue[EngineQueueMessage] | None,
) -> None:
"""Best-effort orchestrator cleanup for GC finalization."""
try:
if request_queue is not None:
request_queue.sync_q.put_nowait({"type": "shutdown"})
request_queue.sync_q.put_nowait(ShutdownRequestMessage())
except Exception:
pass
@@ -295,14 +316,14 @@ class AsyncOmniEngine:
stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None
self.async_chunk = bool(getattr(stage0_args, "async_chunk", False))
self.stage_pools: list[StagePool] = []
self.stage_clients: list[Any] = [] # logical-stage view for external readers
self.stage_clients: list[StageClient] = [] # logical-stage view for external readers
self.input_processor: InputProcessor | None = None
self.supported_tasks: tuple[str, ...] = ("generate",)
self.default_sampling_params_list: list[Any] = []
self.stage_metadata: list[dict[str, Any]] = []
self.request_queue: janus.Queue[dict[str, Any]] | None = None
self.output_queue: janus.Queue[dict[str, Any]] | None = None
self.rpc_output_queue: janus.Queue[dict[str, Any]] | None = None
self.default_sampling_params_list: list[OmniSamplingParams] = []
self.stage_metadata: list[StageRuntimeInfo] = []
self.request_queue: janus.Queue[EngineQueueMessage] | None = None
self.output_queue: janus.Queue[EngineQueueMessage] | None = None
self.rpc_output_queue: janus.Queue[EngineQueueMessage] | None = None
self._shutdown_called = False
self._weak_finalizer: weakref.finalize | None = None
self._rpc_lock = threading.Lock()
@@ -379,15 +400,15 @@ class AsyncOmniEngine:
@staticmethod
def _collect_initialized_clients_for_cleanup(
stage_pools: Sequence[Any],
initialized_clients_by_stage: Mapping[int, Sequence[Any | None]],
) -> list[Any]:
stage_pools: Sequence[StagePool],
initialized_clients_by_stage: Mapping[int, Sequence[StagePoolClient | None]],
) -> list[StageClient]:
"""Collect initialized clients exactly once for failure cleanup."""
collected: list[Any] = []
collected: list[StageClient] = []
seen: set[int] = set()
def _add_client(client: Any) -> None:
def _add_client(client: StageClient | None) -> None:
if client is None:
return
client_id = id(client)
@@ -407,7 +428,7 @@ class AsyncOmniEngine:
return collected
@staticmethod
def _shutdown_initialized_clients(clients: Sequence[Any]) -> None:
def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> None:
"""Best-effort shutdown for attached clients after init failure."""
for client in reversed(list(clients)):
@@ -585,7 +606,7 @@ class AsyncOmniEngine:
plan: ReplicaInitPlan,
stage_init_timeout: int,
llm_stage_launch_lock: threading.Lock,
) -> Any:
) -> StageEngineCoreClientBase:
"""Initialize one LLM replica end-to-end."""
proc = None
@@ -748,7 +769,7 @@ class AsyncOmniEngine:
plan: ReplicaInitPlan,
stage_init_timeout: int,
stage_launch_lock: threading.Lock,
) -> Any:
) -> StageDiffusionClient | InlineStageDiffusionClient:
"""Initialize one diffusion replica end-to-end."""
client = None
@@ -865,7 +886,7 @@ class AsyncOmniEngine:
plan: ReplicaInitPlan,
stage_init_timeout: int,
stage_launch_lock: threading.Lock,
) -> Any:
) -> StagePoolClient:
"""Initialize one replica, regardless of backend type."""
if plan.metadata.stage_type == "diffusion":
@@ -876,7 +897,7 @@ class AsyncOmniEngine:
self,
stage_plans: Sequence[LogicalStageInitPlan],
stage_init_timeout: int,
) -> dict[int, list[Any | None]]:
) -> dict[int, list[StagePoolClient | None]]:
"""Initialize all stage replicas.
Diffusion replicas are launched **inline on the orchestrator thread**
@@ -892,7 +913,7 @@ class AsyncOmniEngine:
"""
stage_launch_lock = threading.Lock()
initialized_clients_by_stage: dict[int, list[Any | None]] = {
initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = {
plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans
}
primary_exc: Exception | None = None
@@ -922,7 +943,7 @@ class AsyncOmniEngine:
# --- 2) LLM replicas: parallel init via a scoped ThreadPoolExecutor. ---
if primary_exc is None and llm_replicas:
future_to_replica: dict[concurrent.futures.Future[Any], tuple[int, int]] = {}
future_to_replica: dict[concurrent.futures.Future[StagePoolClient], tuple[int, int]] = {}
with concurrent.futures.ThreadPoolExecutor(
max_workers=max(1, len(llm_replicas)),
thread_name_prefix="stage-init",
@@ -959,13 +980,13 @@ class AsyncOmniEngine:
def _assemble_stage_pools(
self,
stage_plans: Sequence[LogicalStageInitPlan],
initialized_clients_by_stage: Mapping[int, Sequence[Any | None]],
initialized_clients_by_stage: Mapping[int, Sequence[StagePoolClient | None]],
) -> list[StagePool]:
"""Assemble logical stage pools and update top-level stage metadata."""
stage_pools: list[StagePool] = []
default_sampling_params_list: list[Any] = []
stage_metadata_list: list[dict[str, Any]] = []
default_sampling_params_list: list[OmniSamplingParams] = []
stage_metadata_list: list[StageRuntimeInfo] = []
for plan in stage_plans:
replica_clients = initialized_clients_by_stage[plan.stage_idx]
@@ -973,7 +994,7 @@ class AsyncOmniEngine:
if first_client is None:
raise RuntimeError(f"Stage {plan.stage_idx} initialization completed with a missing client")
clients = [client for client in replica_clients if client is not None]
clients: list[StagePoolClient] = [client for client in replica_clients if client is not None]
stage_vllm_config = None
output_processor = None
if plan.replicas[0].metadata.stage_type != "diffusion":
@@ -991,11 +1012,11 @@ class AsyncOmniEngine:
)
default_sampling_params_list.append(first_client.default_sampling_params)
stage_metadata_list.append(
{
"final_output": first_client.final_output,
"final_output_type": first_client.final_output_type,
"stage_type": first_client.stage_type,
}
StageRuntimeInfo(
final_output=first_client.final_output,
final_output_type=first_client.final_output_type,
stage_type=first_client.stage_type,
)
)
self.default_sampling_params_list = list(default_sampling_params_list)
@@ -1033,7 +1054,7 @@ class AsyncOmniEngine:
stage_pools: list[StagePool] = []
input_processor: InputProcessor | None = None
initialized_clients_by_stage: dict[int, list[Any | None]] = {
initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = {
plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans
}
@@ -1071,7 +1092,7 @@ class AsyncOmniEngine:
self.prompt_expand_func = prompt_expand_func
# Derive logical-stage views for external readers (entrypoints/async_omni.py).
self.stage_clients = [pool.stage_client for pool in self.stage_pools]
self.stage_clients = [cast(StageClient, pool.stage_client) for pool in self.stage_pools]
self.stage_vllm_configs = [pool.stage_vllm_config for pool in self.stage_pools]
self.output_processors = [pool.output_processor for pool in self.stage_pools]
@@ -1079,7 +1100,7 @@ class AsyncOmniEngine:
supported_tasks: set[str] = set()
if any(getattr(pool.stage_client, "is_comprehension", False) for pool in self.stage_pools):
supported_tasks.add("generate")
if any(m.get("final_output_type") == "audio" for m in self.stage_metadata):
if any(m.final_output_type == "audio" for m in self.stage_metadata):
supported_tasks.add("speech")
self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",)
@@ -1127,7 +1148,7 @@ class AsyncOmniEngine:
logger.exception("[AsyncOmniEngine] Orchestrator thread crashed")
error_text = str(e) or "Orchestrator thread crashed"
try:
error_msg = {"type": "error", "error": error_text, "fatal": True}
error_msg = ErrorMessage(error=error_text, fatal=True)
if self.output_queue is not None:
self.output_queue.sync_q.put_nowait(error_msg)
if self.rpc_output_queue is not None:
@@ -1194,11 +1215,13 @@ class AsyncOmniEngine:
reasoning_ended: bool | None = None,
*,
resumable: bool = False,
message_type: str = "add_request",
) -> dict[str, Any]:
message_type: Literal["add_request", "streaming_update"] = "add_request",
) -> StageSubmissionMessage:
"""Build an add_request message after stage-0 preprocessing."""
effective_sampling_params_list = (
list(sampling_params_list) if sampling_params_list is not None else list(self.default_sampling_params_list)
effective_sampling_params_list: list[OmniSamplingParams] = (
list(cast(Sequence[OmniSamplingParams], sampling_params_list))
if sampling_params_list is not None
else list(self.default_sampling_params_list)
)
if not effective_sampling_params_list:
raise ValueError(
@@ -1210,7 +1233,7 @@ class AsyncOmniEngine:
# dict, e.g. for multi_modal_data).
original_prompt = prompt
stage_type = self.stage_metadata[0].get("stage_type")
stage_type = self.stage_metadata[0].stage_type
output_prompt_text: Any = None
_preprocess_ms = 0.0
if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest):
@@ -1261,17 +1284,17 @@ class AsyncOmniEngine:
output_prompt_text = original_prompt.get("prompt")
prompt = request
return {
"type": message_type,
"request_id": request_id,
"prompt": prompt,
"original_prompt": original_prompt,
"output_prompt_text": output_prompt_text,
"sampling_params_list": effective_sampling_params_list,
"final_stage_id": final_stage_id,
"preprocess_ms": _preprocess_ms,
"enqueue_ts": time.perf_counter(),
}
return StageSubmissionMessage(
type=message_type,
request_id=request_id,
prompt=prompt,
original_prompt=original_prompt,
output_prompt_text=output_prompt_text,
sampling_params_list=effective_sampling_params_list,
final_stage_id=final_stage_id,
preprocess_ms=_preprocess_ms,
enqueue_ts=time.perf_counter(),
)
def _enqueue_cfg_companions(
self,
@@ -1311,15 +1334,14 @@ class AsyncOmniEngine:
# deferred to Orchestrator._handle_add_companion, which routes
# admission through StagePool.submit_initial(..., affinity_request_id=...).
self.request_queue.sync_q.put_nowait(
{
"type": "add_companion_request",
"companion_id": cid,
"parent_id": parent_id,
"role": ep.role,
"prompt": request,
"companion_prompt_text": companion_prompt,
"sampling_params_list": companion_spl,
}
AddCompanionRequestMessage(
companion_id=cid,
parent_id=parent_id,
role=ep.role,
prompt=request,
companion_prompt_text=companion_prompt,
sampling_params_list=companion_spl,
)
)
logger.info(
@@ -1766,8 +1788,8 @@ class AsyncOmniEngine:
# CFG companion expansion: create and enqueue companion requests
# so the AR stage also generates their KV caches.
if self.prompt_expand_func is not None and final_stage_id > 0:
original_prompt = msg.get("original_prompt", prompt)
effective_spl = msg.get("sampling_params_list", [])
original_prompt = msg.original_prompt
effective_spl = msg.sampling_params_list
stage0_params = effective_spl[0] if effective_spl else None
if stage0_params is not None:
self._enqueue_cfg_companions(request_id, original_prompt, stage0_params, effective_spl)
@@ -1854,7 +1876,7 @@ class AsyncOmniEngine:
resumable=resumable,
)
def try_get_output(self, timeout: float = 0.001) -> dict[str, Any] | None:
def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None:
"""Read one output message from the Orchestrator output queue."""
if self.output_queue is None:
return None
@@ -1865,7 +1887,7 @@ class AsyncOmniEngine:
raise RuntimeError("Orchestrator died unexpectedly. See logs above.")
return None
async def try_get_output_async(self) -> dict[str, Any] | None:
async def try_get_output_async(self) -> EngineQueueMessage | None:
"""Async read from the Orchestrator output queue."""
if self.output_queue is None:
return None
@@ -1876,7 +1898,7 @@ class AsyncOmniEngine:
raise RuntimeError("Orchestrator died unexpectedly. See logs above.")
return None
def get_stage_metadata(self, stage_id: int) -> dict[str, Any]:
def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
"""Get cached metadata for a stage."""
return self.stage_metadata[stage_id]
@@ -1884,12 +1906,7 @@ class AsyncOmniEngine:
"""Send abort message to the Orchestrator."""
if self.request_queue is None:
raise RuntimeError("request_queue is not initialized")
self.request_queue.sync_q.put_nowait(
{
"type": "abort",
"request_ids": request_ids,
}
)
self.request_queue.sync_q.put_nowait(AbortRequestMessage(request_ids=request_ids))
async def abort_async(self, request_ids: list[str]) -> None:
"""Async abort API."""
@@ -1914,14 +1931,14 @@ class AsyncOmniEngine:
raise RuntimeError("rpc_output_queue is not initialized")
rpc_id = uuid.uuid4().hex
msg = {
"type": "collective_rpc",
"rpc_id": rpc_id,
"method": method,
"args": tuple(args),
"kwargs": kwargs or {},
"stage_ids": stage_ids,
}
msg = CollectiveRPCRequestMessage(
rpc_id=rpc_id,
method=method,
timeout=timeout,
args=tuple(args),
kwargs=kwargs or {},
stage_ids=stage_ids,
)
with self._rpc_lock:
self.request_queue.sync_q.put_nowait(msg)
@@ -1934,25 +1951,25 @@ class AsyncOmniEngine:
except queue.Empty as exc:
raise TimeoutError(f"collective_rpc timed out after {timeout} seconds") from exc
if result_msg.get("type") == "error":
raise RuntimeError(result_msg.get("error", "Orchestrator returned an error message"))
if isinstance(result_msg, ErrorMessage):
raise RuntimeError(result_msg.error)
if result_msg.get("type") != "collective_rpc_result":
if not isinstance(result_msg, CollectiveRPCResultMessage):
logger.warning(
"[AsyncOmniEngine] Dropping unexpected rpc queue message type=%s",
result_msg.get("type"),
getattr(result_msg, "type", type(result_msg).__name__),
)
continue
if result_msg.get("rpc_id") != rpc_id:
if result_msg.rpc_id != rpc_id:
logger.warning(
"[AsyncOmniEngine] Dropping mismatched rpc result rpc_id=%s expected=%s",
result_msg.get("rpc_id"),
result_msg.rpc_id,
rpc_id,
)
continue
return list(result_msg.get("results", []))
return list(result_msg.results)
async def collective_rpc_async(
self,
@@ -1991,7 +2008,7 @@ class AsyncOmniEngine:
logger.info("[AsyncOmniEngine] Shutting down Orchestrator")
try:
if self.request_queue is not None:
self.request_queue.sync_q.put_nowait({"type": "shutdown"})
self.request_queue.sync_q.put_nowait(ShutdownRequestMessage())
except Exception:
pass
if self.is_alive():
vllm_omni/engine/messages.py
+90
View File
@@ -0,0 +1,90 @@
from __future__ import annotations
from typing import Literal
import msgspec
from vllm.inputs import PromptType
from vllm.v1.engine import EngineCoreRequest
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
from vllm_omni.outputs import OmniRequestOutput
class EngineQueueMessage(msgspec.Struct, forbid_unknown_fields=True):
pass
class StageSubmissionMessage(EngineQueueMessage, kw_only=True):
type: Literal["add_request", "streaming_update"]
request_id: str
prompt: EngineCoreRequest | PromptType
original_prompt: EngineCoreRequest | PromptType
output_prompt_text: object | None
sampling_params_list: list[OmniSamplingParams]
final_stage_id: int
preprocess_ms: float
enqueue_ts: float
class AddCompanionRequestMessage(EngineQueueMessage, kw_only=True):
type: Literal["add_companion_request"] = "add_companion_request"
companion_id: str
parent_id: str
role: str
prompt: EngineCoreRequest
companion_prompt_text: object | None
sampling_params_list: list[OmniSamplingParams]
class AbortRequestMessage(EngineQueueMessage, kw_only=True):
type: Literal["abort"] = "abort"
request_ids: list[str]
class CollectiveRPCRequestMessage(EngineQueueMessage, kw_only=True):
type: Literal["collective_rpc"] = "collective_rpc"
rpc_id: str
method: str
timeout: float | None = None
args: tuple[object, ...]
kwargs: dict[str, object]
stage_ids: list[int] | None
class ShutdownRequestMessage(EngineQueueMessage, kw_only=True):
type: Literal["shutdown"] = "shutdown"
class ErrorMessage(EngineQueueMessage, kw_only=True):
type: Literal["error"] = "error"
error: str
fatal: bool = False
request_id: str | None = None
stage_id: int | None = None
class OutputMessage(EngineQueueMessage, kw_only=True):
type: Literal["output"] = "output"
request_id: str
stage_id: int
engine_outputs: OmniRequestOutput
metrics: StageRequestMetrics | None = None
finished: bool
stage_submit_ts: float | None = None
class StageMetricsMessage(EngineQueueMessage, kw_only=True):
type: Literal["stage_metrics"] = "stage_metrics"
request_id: str
stage_id: int
metrics: StageRequestMetrics
stage_submit_ts: float | None = None
class CollectiveRPCResultMessage(EngineQueueMessage, kw_only=True):
type: Literal["collective_rpc_result"] = "collective_rpc_result"
rpc_id: str
method: str
stage_ids: list[int]
results: list[object]
vllm_omni/engine/orchestrator.py
+233 -175
View File
@@ -24,6 +24,18 @@ from vllm.v1.engine.exceptions import EngineDeadError
from vllm_omni.engine import OmniEngineCoreRequest
from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
from vllm_omni.engine.messages import (
AbortRequestMessage,
AddCompanionRequestMessage,
CollectiveRPCRequestMessage,
CollectiveRPCResultMessage,
EngineQueueMessage,
ErrorMessage,
OutputMessage,
ShutdownRequestMessage,
StageMetricsMessage,
StageSubmissionMessage,
)
from vllm_omni.engine.serialization import serialize_additional_information
from vllm_omni.engine.stage_pool import StagePool
from vllm_omni.outputs import OmniRequestOutput
@@ -115,9 +127,9 @@ class Orchestrator:
def __init__(
self,
request_async_queue: janus.AsyncQueue[dict[str, Any]],
output_async_queue: janus.AsyncQueue[dict[str, Any]],
rpc_async_queue: janus.AsyncQueue[dict[str, Any]],
request_async_queue: janus.AsyncQueue[EngineQueueMessage],
output_async_queue: janus.AsyncQueue[EngineQueueMessage],
rpc_async_queue: janus.AsyncQueue[EngineQueueMessage],
stage_pools: list[StagePool],
*,
async_chunk: bool = False,
@@ -196,36 +208,38 @@ class Orchestrator:
"""Read messages from the main thread via request_async_queue."""
while True:
msg = await self.request_async_queue.get()
msg_type = msg.get("type")
if msg_type == "add_request":
await self._handle_add_request(msg)
elif msg_type == "streaming_update":
await self._handle_streaming_update(msg)
elif msg_type == "add_companion_request":
if isinstance(msg, StageSubmissionMessage):
if msg.type == "add_request":
await self._handle_add_request(msg)
elif msg.type == "streaming_update":
await self._handle_streaming_update(msg)
else:
logger.warning("[Orchestrator] Unknown stage submission type: %s", msg.type)
elif isinstance(msg, AddCompanionRequestMessage):
await self._handle_add_companion(msg)
elif msg_type == "abort":
elif isinstance(msg, AbortRequestMessage):
await self._handle_abort(msg)
elif msg_type == "collective_rpc":
elif isinstance(msg, CollectiveRPCRequestMessage):
await self._handle_collective_rpc(msg)
elif msg_type == "shutdown":
elif isinstance(msg, ShutdownRequestMessage):
logger.info("[Orchestrator] Received shutdown signal")
self._shutdown_event.set()
self._shutdown_stages()
break
else:
logger.warning("[Orchestrator] Unknown message type: %s", msg_type)
logger.warning("[Orchestrator] Unknown message type: %s", type(msg).__name__)
async def _handle_add_request(self, msg: dict[str, Any]) -> None:
async def _handle_add_request(self, msg: StageSubmissionMessage) -> None:
"""Handle an add_request message from the main thread."""
stage_id = 0
request_id = msg["request_id"]
prompt = msg["prompt"]
original_prompt = msg.get("original_prompt", prompt)
sampling_params_list = msg["sampling_params_list"]
request_id = msg.request_id
prompt = msg.prompt
original_prompt = msg.original_prompt
sampling_params_list = msg.sampling_params_list
if not sampling_params_list:
raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.")
final_stage_id = msg["final_stage_id"]
final_stage_id = msg.final_stage_id
logger.debug(
"[Orchestrator] _handle_add_request: stage=%s req=%s "
@@ -249,27 +263,27 @@ class Orchestrator:
self.request_states[request_id] = req_state
req_state.streaming.enabled = bool(getattr(prompt, "resumable", False))
req_state.stage_submit_ts[stage_id] = _time.time()
enqueue_ts = msg.get("enqueue_ts", 0.0)
enqueue_ts = msg.enqueue_ts
if enqueue_ts > 0:
req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - enqueue_ts) * 1000.0
preprocess_ms = msg.get("preprocess_ms", 0.0)
preprocess_ms = msg.preprocess_ms
if preprocess_ms > 0:
req_state.pipeline_timings["preprocess_ms"] = preprocess_ms
await self.stage_pools[stage_id].submit_initial(
request_id,
req_state,
prompt,
prompt_text=msg.get("output_prompt_text"),
prompt_text=msg.output_prompt_text,
)
if self.async_chunk and stage_id == 0 and final_stage_id > 0:
await self._prewarm_async_chunk_stages(request_id, prompt, req_state)
async def _handle_streaming_update(self, msg: dict[str, Any]) -> None:
async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None:
"""Handle a streaming_update message for an existing request."""
stage_id = 0
request_id = msg["request_id"]
request = msg["prompt"]
request_id = msg.request_id
request = msg.prompt
req_state = self.request_states.get(request_id)
if req_state is None:
@@ -277,13 +291,22 @@ class Orchestrator:
"[Orchestrator] streaming_update for unknown req=%s, falling back to add_request",
request_id,
)
fallback_msg = dict(msg)
fallback_msg["type"] = "add_request"
fallback_msg = StageSubmissionMessage(
type="add_request",
request_id=msg.request_id,
prompt=msg.prompt,
original_prompt=msg.original_prompt,
output_prompt_text=msg.output_prompt_text,
sampling_params_list=msg.sampling_params_list,
final_stage_id=msg.final_stage_id,
preprocess_ms=msg.preprocess_ms,
enqueue_ts=msg.enqueue_ts,
)
await self._handle_add_request(fallback_msg)
return
if "sampling_params_list" in msg and msg["sampling_params_list"]:
req_state.sampling_params_list = msg["sampling_params_list"]
if msg.sampling_params_list:
req_state.sampling_params_list = msg.sampling_params_list
req_state.streaming.enabled = True
req_state.stage_submit_ts[stage_id] = _time.time()
@@ -291,16 +314,16 @@ class Orchestrator:
request_id,
req_state,
request,
prompt_text=msg.get("output_prompt_text"),
prompt_text=msg.output_prompt_text,
)
async def _handle_add_companion(self, msg: dict[str, Any]) -> None:
async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None:
"""Handle an add_companion_request message: submit companion to stage 0."""
companion_id = msg["companion_id"]
parent_id = msg["parent_id"]
role = msg["role"]
companion_prompt = msg["prompt"]
sampling_params_list = msg["sampling_params_list"]
companion_id = msg.companion_id
parent_id = msg.parent_id
role = msg.role
companion_prompt = msg.prompt
sampling_params_list = msg.sampling_params_list
parent_state = self.request_states.get(parent_id)
if parent_state is None:
@@ -326,7 +349,7 @@ class Orchestrator:
companion_id,
companion_state,
companion_prompt,
prompt_text=msg.get("companion_prompt_text"),
prompt_text=msg.companion_prompt_text,
affinity_request_id=parent_id,
)
@@ -338,9 +361,9 @@ class Orchestrator:
companion_replica_id,
)
async def _handle_abort(self, msg: dict[str, Any]) -> None:
async def _handle_abort(self, msg: AbortRequestMessage) -> None:
"""Handle an abort message from the main thread."""
request_ids = msg["request_ids"]
request_ids = msg.request_ids
await self._cleanup_request_ids(
self._cfg_tracker.abort_parents(request_ids),
abort=True,
@@ -359,14 +382,14 @@ class Orchestrator:
for pool in self.stage_pools:
pool.release_bindings(request_ids)
async def _handle_collective_rpc(self, msg: dict[str, Any]) -> None:
async def _handle_collective_rpc(self, msg: CollectiveRPCRequestMessage) -> None:
"""Handle a control-plane RPC request from the main thread."""
rpc_id = msg["rpc_id"]
method = msg["method"]
timeout = msg.get("timeout")
args = tuple(msg.get("args", ()))
kwargs = dict(msg.get("kwargs") or {})
requested_stage_ids = msg.get("stage_ids")
rpc_id = msg.rpc_id
method = msg.method
timeout = msg.timeout
args = msg.args
kwargs = dict(msg.kwargs)
requested_stage_ids = msg.stage_ids
target_pools: list[StagePool] = []
if requested_stage_ids is None:
@@ -393,13 +416,12 @@ class Orchestrator:
results.append(stage_result)
await self.rpc_async_queue.put(
{
"type": "collective_rpc_result",
"rpc_id": rpc_id,
"method": method,
"stage_ids": stage_ids,
"results": results,
}
CollectiveRPCResultMessage(
rpc_id=rpc_id,
method=method,
stage_ids=stage_ids,
results=results,
)
)
# ---- Orchestration loop ----
@@ -469,13 +491,12 @@ class Orchestrator:
for req_id, req_state in list(self.request_states.items()):
if stage_id in req_state.stage_submit_ts:
await self.output_async_queue.put(
{
"type": "error",
"error": str(e),
"fatal": True,
"request_id": req_id,
"stage_id": stage_id,
}
ErrorMessage(
error=str(e),
fatal=True,
request_id=req_id,
stage_id=stage_id,
)
)
self.request_states.pop(req_id, None)
self._shutdown_event.set()
@@ -534,12 +555,11 @@ class Orchestrator:
else:
parent_id = output.request_id
await self.output_async_queue.put(
{
"type": "error",
"request_id": parent_id,
"stage_id": stage_id,
"error": output.error,
}
ErrorMessage(
request_id=parent_id,
stage_id=stage_id,
error=output.error,
)
)
await self._cleanup_request_ids(
[parent_id, *self._cfg_tracker.cleanup_parent(parent_id)],
@@ -596,25 +616,23 @@ class Orchestrator:
if self.stage_pools[stage_id].final_output:
await self.output_async_queue.put(
{
"type": "output",
"request_id": req_id,
"stage_id": stage_id,
"engine_outputs": output,
"metrics": stage_metrics,
"finished": finished and stage_id == req_state.final_stage_id,
"stage_submit_ts": submit_ts,
}
OutputMessage(
request_id=req_id,
stage_id=stage_id,
engine_outputs=output,
metrics=stage_metrics,
finished=finished and stage_id == req_state.final_stage_id,
stage_submit_ts=submit_ts,
)
)
elif stage_metrics is not None:
await self.output_async_queue.put(
{
"type": "stage_metrics",
"request_id": req_id,
"stage_id": stage_id,
"metrics": stage_metrics,
"stage_submit_ts": submit_ts,
}
StageMetricsMessage(
request_id=req_id,
stage_id=stage_id,
metrics=stage_metrics,
stage_submit_ts=submit_ts,
)
)
if self._pd_pair is not None and finished and stage_id == self._pd_pair[0]:
@@ -775,79 +793,120 @@ class Orchestrator:
"""Forward output from the current logical stage to the next one."""
next_logical = src_stage_id + 1
next_pool = self.stage_pools[next_logical]
next_client = next_pool.stage_client
if next_pool.stage_type == "diffusion":
await self._forward_to_diffusion_stage(
req_id,
src_stage_id,
output,
req_state,
)
else:
await self._forward_to_llm_stage(
req_id,
src_stage_id,
output,
req_state,
is_streaming_session=is_streaming_session,
is_final_update=is_final_update,
)
async def _forward_to_diffusion_stage(
self,
req_id: str,
src_stage_id: int,
output: Any,
req_state: OrchestratorRequestState,
) -> None:
"""Forward output from an LLM stage to the next diffusion stage."""
next_logical = src_stage_id + 1
next_pool = self.stage_pools[next_logical]
params = req_state.sampling_params_list[next_logical]
source_outputs = [output]
already_submitted = self._next_stage_already_submitted(src_stage_id, req_state)
diffusion_client = next_pool.stage_client
if diffusion_client.custom_process_input_func is not None:
_t_ar2d = _time.perf_counter()
diffusion_prompt = diffusion_client.custom_process_input_func(
source_outputs,
req_state.prompt,
getattr(diffusion_client, "requires_multimodal_data", False),
)
_dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d
logger.info(
"[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d",
req_id,
_dt_ar2d,
src_stage_id,
next_logical,
)
if isinstance(diffusion_prompt, list):
if not diffusion_prompt:
error_output = OmniRequestOutput.from_error(
req_id,
f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}",
)
logger.warning(
"[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; "
"routing terminal error output",
req_id,
src_stage_id,
next_logical,
)
await self.output_async_queue.put(
OutputMessage(
request_id=req_id,
stage_id=next_logical,
engine_outputs=error_output,
metrics=None,
finished=True,
)
)
await self._cleanup_request_ids(
[req_id, *self._cfg_tracker.cleanup_parent(req_id)],
)
return
if already_submitted and len(diffusion_prompt) == 1:
diffusion_prompt = diffusion_prompt[0]
else:
diffusion_prompt = req_state.prompt
if already_submitted:
await next_pool.submit_update(req_id, req_state, diffusion_prompt)
else:
await next_pool.submit_initial(
req_id,
req_state,
diffusion_prompt,
submit_kwargs={
"kv_sender_info": self._build_kv_sender_info(
list(getattr(diffusion_client, "engine_input_source", None) or [src_stage_id]),
request_id=req_id,
)
},
params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
)
req_state.stage_submit_ts[next_logical] = _time.time()
async def _forward_to_llm_stage(
self,
req_id: str,
src_stage_id: int,
output: Any,
req_state: OrchestratorRequestState,
*,
is_streaming_session: bool = False,
is_final_update: bool = False,
) -> None:
"""Forward output from the current stage to the next LLM stage."""
next_logical = src_stage_id + 1
next_pool = self.stage_pools[next_logical]
params = req_state.sampling_params_list[next_logical]
source_outputs = [output]
next_stage_resumable = is_streaming_session and not is_final_update
already_submitted = self._next_stage_already_submitted(src_stage_id, req_state)
requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False)
if next_pool.stage_type == "diffusion":
if next_client.custom_process_input_func is not None:
_t_ar2d = _time.perf_counter()
diffusion_prompt = next_client.custom_process_input_func(
source_outputs,
req_state.prompt,
requires_multimodal_data,
)
_dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d
logger.info(
"[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d",
req_id,
_dt_ar2d,
src_stage_id,
next_logical,
)
if isinstance(diffusion_prompt, list):
if not diffusion_prompt:
error_output = OmniRequestOutput.from_error(
req_id,
f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}",
)
logger.warning(
"[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; "
"routing terminal error output",
req_id,
src_stage_id,
next_logical,
)
await self.output_async_queue.put(
{
"type": "output",
"request_id": req_id,
"stage_id": next_logical,
"engine_outputs": error_output,
"metrics": None,
"finished": True,
}
)
await self._cleanup_request_ids(
[req_id, *self._cfg_tracker.cleanup_parent(req_id)],
)
return
if already_submitted and len(diffusion_prompt) == 1:
diffusion_prompt = diffusion_prompt[0]
else:
diffusion_prompt = req_state.prompt
if already_submitted:
await next_pool.submit_update(req_id, req_state, diffusion_prompt)
else:
await next_pool.submit_initial(
req_id,
req_state,
diffusion_prompt,
submit_kwargs={
"kv_sender_info": self._build_kv_sender_info(
list(getattr(next_client, "engine_input_source", None) or [src_stage_id]),
request_id=req_id,
)
},
params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
)
req_state.stage_submit_ts[next_logical] = _time.time()
return
# PD disaggregation: prefill → decode routing uses original prompt + KV transfer params
if self._pd_pair is not None and (src_stage_id, next_logical) == self._pd_pair:
@@ -894,6 +953,7 @@ class Orchestrator:
{},
)[req_id] = req_state.pd_prefill_multimodal_output
next_client = next_pool.llm_stage_client
try:
next_inputs = next_client.process_engine_inputs(
source_outputs,
@@ -912,7 +972,7 @@ class Orchestrator:
for next_input in next_inputs:
# Only AR thinker stages consume encoder mm_features; downstream
# (talker/code2wav/…) must not see them (avoids encoder-cache misses).
model_stage = getattr(next_client, "model_stage", None)
model_stage = next_client.model_stage
mm_features = req_state.mm_features if model_stage == "thinker" else None
request = build_engine_core_request_from_tokens(
request_id=req_id,
@@ -962,7 +1022,7 @@ class Orchestrator:
req_state.prompt,
submit_kwargs={
"kv_sender_info": self._build_kv_sender_info(
list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]),
list(next_pool.stage_client.engine_input_source or [next_stage_id - 1]),
request_id=request_id,
)
},
@@ -1009,14 +1069,14 @@ class Orchestrator:
continue
sender_pool = self.stage_pools[sender_stage_id]
sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None
if sender_stage is None:
sender_stage = sender_pool.stage_client
get_sender_info = getattr(sender_stage, "get_kv_sender_info", None)
if not callable(get_sender_info):
if sender_pool.stage_type == "diffusion":
continue
sender_info = get_sender_info()
sender_stage = sender_pool.get_bound_llm_client(request_id) if request_id is not None else None
if sender_stage is None:
sender_stage = sender_pool.llm_stage_client
sender_info = sender_stage.get_kv_sender_info()
if not sender_info:
logger.warning(
"[Orchestrator] Stage-%s has no KV sender info available",
@@ -1050,16 +1110,15 @@ class Orchestrator:
msg = self.request_async_queue.get_nowait()
except Exception:
break
if msg.get("type") == "add_request":
req_id = msg["request_id"]
if isinstance(msg, StageSubmissionMessage) and msg.type == "add_request":
req_id = msg.request_id
await self.output_async_queue.put(
{
"type": "error",
"error": self._fatal_error,
"fatal": True,
"request_id": req_id,
"stage_id": self._fatal_error_stage_id,
}
ErrorMessage(
error=self._fatal_error,
fatal=True,
request_id=req_id,
stage_id=self._fatal_error_stage_id,
)
)
notified.add(req_id)
@@ -1069,13 +1128,12 @@ class Orchestrator:
for req_id in list(self.request_states):
if req_id not in notified:
await self.output_async_queue.put(
{
"type": "error",
"error": self._fatal_error,
"fatal": True,
"request_id": req_id,
"stage_id": self._fatal_error_stage_id,
}
ErrorMessage(
error=self._fatal_error,
fatal=True,
request_id=req_id,
stage_id=self._fatal_error_stage_id,
)
)
self.request_states.pop(req_id, None)
vllm_omni/engine/output_modality.py
+3
View File
@@ -9,6 +9,9 @@ from __future__ import annotations
import re
from enum import Enum, Flag, StrEnum, auto
from typing import Literal, TypeAlias
FinalOutputModalityType: TypeAlias = Literal["text", "image", "audio", "video"]
_MODALITY_ALIASES: dict[str, str] = {
"speech": "audio",
vllm_omni/engine/stage_client.py
+111
View File
@@ -0,0 +1,111 @@
"""Shared stage-client typing for vLLM-Omni runtime surfaces."""
from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Protocol
if TYPE_CHECKING:
from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniTokensPrompt
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.engine.output_modality import FinalOutputModalityType
from vllm_omni.inputs.data import OmniSamplingParams
class StageClient(Protocol):
"""Shared metadata and lifecycle surface for all stage clients.
This protocol intentionally covers the common attributes consumed by
orchestration and entrypoint code. Backend-specific request APIs stay
typed with the concrete client union.
"""
stage_id: int
replica_id: int
stage_type: str
final_output: bool
final_output_type: FinalOutputModalityType | None
default_sampling_params: OmniSamplingParams
requires_multimodal_data: bool
custom_process_input_func: Any | None
engine_input_source: Sequence[int]
is_comprehension: bool
def check_health(self) -> None: ...
def shutdown(self) -> None: ...
class StageClientBase:
"""Runtime nominal base for stage clients.
Keeping the runtime base separate from the protocol avoids Protocol-related
MRO interference with concrete client super() calls.
"""
pass
class StagePoolClient(StageClient, Protocol):
"""Common pool-facing client surface shared by every stage backend."""
async def abort_requests_async(self, request_ids: list[str]) -> None: ...
async def collective_rpc_async(
self,
method: str,
timeout: float | None = None,
args: tuple[Any, ...] = (),
kwargs: dict[str, Any] | None = None,
) -> Any: ...
class StagePoolLLMClient(StagePoolClient, Protocol):
"""Pool-facing API for LLM-style stages."""
model_stage: str | None
async def add_request_async(self, request: EngineCoreRequest) -> None: ...
async def get_output_async(self) -> EngineCoreOutputs: ...
def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None: ...
def process_engine_inputs(
self,
source_outputs: list[Any],
prompt: Any = None,
streaming_context: Any | None = None,
) -> list[OmniTokensPrompt]: ...
def get_kv_sender_info(
self,
*,
base_port: int = ...,
kv_transfer_port_offset: int = ...,
) -> dict[str, Any] | None: ...
class StagePoolDiffusionClient(StagePoolClient, Protocol):
"""Pool-facing API for diffusion stages."""
async def add_request_async(
self,
request_id: str,
prompt: OmniPromptType,
sampling_params: OmniDiffusionSamplingParams,
kv_sender_info: dict[int, dict[str, Any]] | None = None,
) -> None: ...
async def add_batch_request_async(
self,
request_id: str,
prompts: list[OmniPromptType],
sampling_params: OmniDiffusionSamplingParams,
kv_sender_info: dict[int, dict[str, Any]] | None = None,
) -> None: ...
def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: ...
vllm_omni/engine/stage_engine_core_client.py
+2 -1
View File
@@ -23,6 +23,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
from vllm_omni.distributed.omni_connectors.utils.initialization import (
KV_TRANSFER_PORT_OFFSET,
)
from vllm_omni.engine.stage_client import StageClientBase
from vllm_omni.engine.stage_init_utils import StageMetadata
if TYPE_CHECKING:
@@ -56,7 +57,7 @@ def _default_process_engine_inputs(
]
class StageEngineCoreClientBase:
class StageEngineCoreClientBase(StageClientBase):
"""Shared stage-aware behavior for async EngineCore clients.
The concrete transport/load-balancing behavior is supplied by the
vllm_omni/engine/stage_pool.py
+44 -28
View File
@@ -5,11 +5,16 @@ from __future__ import annotations
import asyncio
import time as _time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, cast
from vllm.logger import init_logger
from vllm.v1.engine import EngineCoreOutputs
from vllm_omni.engine.stage_client import (
StagePoolClient,
StagePoolDiffusionClient,
StagePoolLLMClient,
)
from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
from vllm_omni.metrics.stats import StageStats
from vllm_omni.metrics.utils import count_tokens_from_outputs
@@ -35,20 +40,20 @@ class StagePool:
def __init__(
self,
stage_id: int,
clients: Any | list[Any],
clients: StagePoolClient | list[StagePoolClient],
*,
output_processor: Any = None,
stage_vllm_config: Any = None,
) -> None:
if isinstance(clients, list):
normalized_clients = list(clients)
normalized_clients: list[StagePoolClient] = list(clients)
else:
normalized_clients = [clients]
if not normalized_clients:
raise ValueError(f"StagePool for stage {stage_id} has no replicas")
self.stage_id = stage_id
self.clients: list[Any] = normalized_clients
self.clients: list[StagePoolClient] = normalized_clients
self._output_processor = output_processor
self._stage_vllm_config = stage_vllm_config
self._next_replica_id = 0
@@ -63,16 +68,20 @@ class StagePool:
@property
def stage_type(self) -> str | None:
return getattr(self.stage_client, "stage_type", None)
return self.stage_client.stage_type
@property
def final_output(self) -> bool:
return bool(getattr(self.clients[0], "final_output", False))
return self.clients[0].final_output
@property
def stage_client(self) -> Any:
def stage_client(self) -> StagePoolClient:
return self.clients[0]
@property
def llm_stage_client(self) -> StagePoolLLMClient:
return cast(StagePoolLLMClient, self.stage_client)
@property
def stage_vllm_config(self) -> Any:
return self._stage_vllm_config
@@ -87,13 +96,20 @@ class StagePool:
"""Return the currently bound replica id for *request_id* if present."""
return self._request_bindings.get(request_id)
def get_bound_client(self, request_id: str) -> Any | None:
def get_bound_client(self, request_id: str) -> StagePoolClient | None:
"""Return the currently bound client for *request_id* if present."""
replica_id = self.get_bound_replica_id(request_id)
if replica_id is None:
return None
return self.clients[replica_id]
def get_bound_llm_client(self, request_id: str) -> StagePoolLLMClient | None:
"""Return the currently bound LLM client for *request_id* if present."""
client = self.get_bound_client(request_id)
if client is None:
return None
return cast(StagePoolLLMClient, client)
def release_binding(self, request_id: str) -> None:
"""Drop the route binding for *request_id* in this stage."""
self._request_bindings.pop(request_id, None)
@@ -125,6 +141,12 @@ class StagePool:
self._request_bindings[request_id] = chosen
return chosen
def _llm_client(self, replica_id: int) -> StagePoolLLMClient:
return cast(StagePoolLLMClient, self.clients[replica_id])
def _diffusion_client(self, replica_id: int) -> StagePoolDiffusionClient:
return cast(StagePoolDiffusionClient, self.clients[replica_id])
# ---- Metrics ----
def build_stage_metrics(
@@ -188,7 +210,7 @@ class StagePool:
request_id,
affinity_request_id=affinity_request_id,
)
client = self.clients[replica_id]
client = self._diffusion_client(replica_id)
if isinstance(request, list):
await client.add_batch_request_async(request_id, request, params, **submit_kwargs)
else:
@@ -212,7 +234,7 @@ class StagePool:
raise
try:
await self.clients[replica_id].add_request_async(request, **submit_kwargs)
await self._llm_client(replica_id).add_request_async(request, **submit_kwargs)
except Exception:
self.release_binding(request_id)
rollback = getattr(self.output_processor, "remove_request", None)
@@ -244,7 +266,7 @@ class StagePool:
replica_id = self.select_replica_id(request_id)
if self.stage_type == "diffusion":
await self.clients[replica_id].add_request_async(request_id, request, params)
await self._diffusion_client(replica_id).add_request_async(request_id, request, params)
else:
# Refresh the shared output-processor state before yielding to the
# stage client so streaming segments are merged against the latest
@@ -256,12 +278,12 @@ class StagePool:
request_index=0,
queue=None,
)
await self.clients[replica_id].add_request_async(request)
await self._llm_client(replica_id).add_request_async(request)
return replica_id
# ---- Stage-local polling ----
async def _poll_stage_raw(self, client: Any) -> EngineCoreOutputs | None:
async def _poll_stage_raw(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None:
"""Pull raw EngineCoreOutputs from a stage replica without processing."""
outputs = await client.get_output_async()
if not outputs.outputs:
@@ -274,7 +296,7 @@ class StagePool:
raw_outputs: EngineCoreOutputs,
) -> list[Any]:
"""Run the shared LLM output processor on one raw poll result."""
client = self.clients[replica_id]
client = self._llm_client(replica_id)
processor = self.output_processor
processed = processor.process_outputs(
raw_outputs.outputs,
@@ -297,7 +319,7 @@ class StagePool:
timeout_s: float = 0.001,
) -> EngineCoreOutputs | None:
"""Poll raw EngineCore outputs from one LLM replica once."""
client = self.clients[replica_id]
client = self._llm_client(replica_id)
try:
return await asyncio.wait_for(
self._poll_stage_raw(client),
@@ -317,7 +339,7 @@ class StagePool:
def poll_diffusion_output(self, replica_id: int) -> Any | None:
"""Drain one ready diffusion output from the given replica if present."""
return self.clients[replica_id].get_diffusion_output_nowait()
return self._diffusion_client(replica_id).get_diffusion_output_nowait()
# ---- Stage-local control plane ----
@@ -360,18 +382,12 @@ class StagePool:
kwargs = dict(kwargs or {})
client = self.clients[replica_id]
try:
if hasattr(client, "collective_rpc_async"):
return await client.collective_rpc_async(
method=method,
timeout=timeout,
args=args,
kwargs=kwargs,
)
return {
"supported": False,
"todo": True,
"reason": f"{client.__class__.__name__}.collective_rpc_async is not implemented yet",
}
return await client.collective_rpc_async(
method=method,
timeout=timeout,
args=args,
kwargs=kwargs,
)
except Exception as exc:
logger.exception(
"[StagePool] collective_rpc failed: stage=%s replica=%s method=%s",
vllm_omni/entrypoints/async_omni.py
+38 -30
View File
@@ -26,6 +26,7 @@ from vllm.tasks import SupportedTask
from vllm.v1.engine.exceptions import EngineDeadError
from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
from vllm_omni.entrypoints.client_request_state import ClientRequestState
from vllm_omni.entrypoints.omni_base import (
OmniBase,
@@ -416,7 +417,12 @@ class AsyncOmni(EngineClient, OmniBase):
except (asyncio.CancelledError, GeneratorExit):
cancelled = True
except Exception as error:
await req_state.queue.put({"request_id": request_id, "error": error})
await req_state.queue.put(
ErrorMessage(
request_id=request_id,
error=str(error),
)
)
finally:
if not cancelled:
# Send empty final request to indicate that inputs have
@@ -501,23 +507,25 @@ class AsyncOmni(EngineClient, OmniBase):
while True:
result = await req_state.queue.get()
stage_id = result.get("stage_id", 0)
if result.get("type") == "error" and result.get("fatal"):
raise OmniEngineDeadError(
result.get("error", ""),
error_stage_id=result.get("stage_id"),
)
# Check for errors
if "error" in result:
if isinstance(result, ErrorMessage):
logger.error(
"[AsyncOmni] Orchestrator error for req=%s stage-%s: %s",
request_id,
stage_id,
result["error"],
result.stage_id,
result.error,
)
raise RuntimeError(result)
if result.fatal:
raise OmniEngineDeadError(
result.error,
error_stage_id=result.stage_id,
)
raise RuntimeError(result.error)
if not isinstance(result, OutputMessage):
logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result)
continue
stage_id = result.stage_id
self._check_engine_output_error(result, request_id, stage_id)
@@ -541,7 +549,7 @@ class AsyncOmni(EngineClient, OmniBase):
yield output_to_yield
# The Orchestrator sets "finished" when the final stage is done
if result.get("finished"):
if result.finished:
break
# ==================== Output Handler ====================
@@ -596,29 +604,29 @@ class AsyncOmni(EngineClient, OmniBase):
except OmniEngineDeadError as e:
logger.error("[AsyncOmni] Engine dead: %s", e)
for req_state in list(self.request_states.values()):
error_msg = {
"type": "error",
"error": str(e),
"fatal": True,
"request_id": req_state.request_id,
}
if e.error_stage_id is not None:
error_msg["stage_id"] = e.error_stage_id
error_msg = ErrorMessage(
error=str(e),
fatal=True,
request_id=req_state.request_id,
stage_id=e.error_stage_id,
)
await req_state.queue.put(error_msg)
except EngineDeadError as e:
logger.error("[AsyncOmni] Engine dead: %s", e)
for req_state in list(self.request_states.values()):
error_msg = {
"type": "error",
"error": str(e),
"fatal": True,
"request_id": req_state.request_id,
}
error_msg = ErrorMessage(
error=str(e),
fatal=True,
request_id=req_state.request_id,
)
await req_state.queue.put(error_msg)
except Exception as e:
logger.exception("[AsyncOmni] final_output_loop failed.")
for req_state in list(self.request_states.values()):
error_msg = {"request_id": req_state.request_id, "error": str(e)}
error_msg = ErrorMessage(
request_id=req_state.request_id,
error=str(e),
)
await req_state.queue.put(error_msg)
self.final_output_task = None
vllm_omni/entrypoints/omni.py
+3 -2
View File
@@ -10,6 +10,7 @@ from tqdm.auto import tqdm
from vllm.logger import init_logger
from vllm.sampling_params import RequestOutputKind
from vllm_omni.engine.messages import OutputMessage
from vllm_omni.entrypoints.client_request_state import ClientRequestState
from vllm_omni.entrypoints.omni_base import OmniBase
from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
@@ -33,7 +34,7 @@ class Omni(OmniBase):
for stage_id, params in enumerate(sampling_params_list):
sp = copy.deepcopy(params)
stage_meta = self.engine.get_stage_metadata(stage_id)
if stage_meta.get("stage_type") != "diffusion" and hasattr(sp, "output_kind"):
if stage_meta.stage_type != "diffusion" and hasattr(sp, "output_kind"):
sp.output_kind = RequestOutputKind.FINAL_ONLY
effective_params.append(sp)
return effective_params
@@ -181,7 +182,7 @@ class Omni(OmniBase):
if output_to_yield is not None:
yield output_to_yield
if msg.get("finished"):
if isinstance(msg, OutputMessage) and msg.finished:
active_reqs.discard(req_id)
if pbar is not None:
pbar.update(1)
vllm_omni/entrypoints/omni_base.py
+36 -45
View File
@@ -4,7 +4,6 @@ import argparse
import os
import sys
import time
import types
import warnings
import weakref
from collections.abc import Sequence
@@ -15,6 +14,12 @@ from vllm.logger import init_logger
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.messages import (
EngineQueueMessage,
ErrorMessage,
OutputMessage,
StageMetricsMessage,
)
from vllm_omni.entrypoints.client_request_state import ClientRequestState
from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e
@@ -189,12 +194,10 @@ class OmniBase(PDDisaggregationMixin):
self.default_sampling_params_list = self.engine.default_sampling_params_list
if not self.output_modalities:
self.output_modalities = [
self.engine.get_stage_metadata(i).get("final_output_type") for i in range(self.engine.num_stages)
self.engine.get_stage_metadata(i).final_output_type for i in range(self.engine.num_stages)
]
self._stage_meta_list = [
types.SimpleNamespace(**self.engine.get_stage_metadata(i)) for i in range(self.engine.num_stages)
]
self._stage_meta_list = [self.engine.get_stage_metadata(i) for i in range(self.engine.num_stages)]
logger.info(
"[%s] Initialized with %s stages for model %s",
@@ -289,17 +292,15 @@ class OmniBase(PDDisaggregationMixin):
self._stage_meta_list,
)
def _process_stage_metrics_message(self, msg: dict[str, Any]) -> None:
req_id = msg.get("request_id")
def _process_stage_metrics_message(self, msg: StageMetricsMessage) -> None:
req_id = msg.request_id
req_state = self.request_states.get(req_id)
if req_state is None or req_state.metrics is None:
return
_m = msg.get("metrics")
if _m is None:
return
stage_id = msg.get("stage_id", 0)
_m = msg.metrics
stage_id = msg.stage_id
req_state.metrics.on_stage_metrics(stage_id, req_id, _m)
submit_ts = msg.get("stage_submit_ts")
submit_ts = msg.stage_submit_ts
now = time.time()
if req_state.metrics.stage_first_ts[stage_id] is None:
req_state.metrics.stage_first_ts[stage_id] = submit_ts if submit_ts is not None else now
@@ -307,40 +308,30 @@ class OmniBase(PDDisaggregationMixin):
def _handle_output_message(
self,
msg: dict[str, Any] | None,
msg: EngineQueueMessage | None,
) -> OutputMessageHandleResult:
"""Handle one Orchestrator output-queue message."""
if msg is None:
return True, None, None, None
msg_type = msg.get("type")
if msg_type == "stage_metrics":
if isinstance(msg, StageMetricsMessage):
self._process_stage_metrics_message(msg)
return True, None, None, None
if msg_type == "error":
error_text = msg.get("error", "Orchestrator returned an error message")
stage_id = msg.get("stage_id")
if msg.get("fatal"):
if isinstance(msg, ErrorMessage):
if msg.fatal:
raise OmniEngineDeadError(
error_text,
error_stage_id=stage_id,
msg.error,
error_stage_id=msg.stage_id,
)
raise RuntimeError(error_text)
raise RuntimeError(msg.error)
if msg_type != "output":
logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg_type)
if not isinstance(msg, OutputMessage):
logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg.type)
return True, None, None, None
req_id = msg.get("request_id")
if req_id is None:
logger.warning("[%s] got output message without request_id", self.__class__.__name__)
return True, None, None, None
stage_id = msg.get("stage_id")
if stage_id is None:
logger.warning("[%s] got output message without stage_id for req=%s", self.__class__.__name__, req_id)
return True, None, None, None
req_id = msg.request_id
stage_id = msg.stage_id
req_state = self.request_states.get(req_id)
if req_state is None:
@@ -357,7 +348,7 @@ class OmniBase(PDDisaggregationMixin):
def _check_engine_output_error(
self,
result: dict[str, Any],
result: OutputMessage,
request_id: str,
stage_id: int,
) -> None:
@@ -367,7 +358,7 @@ class OmniBase(PDDisaggregationMixin):
engine is unrecoverable, otherwise raises :class:`EngineGenerateError`
(recoverable, single-request failure).
"""
engine_outputs = result.get("engine_outputs")
engine_outputs = result.engine_outputs
error_text = getattr(engine_outputs, "error", None)
if error_text is None:
return
@@ -388,17 +379,17 @@ class OmniBase(PDDisaggregationMixin):
def _process_single_result(
self,
result: dict[str, Any],
result: OutputMessage,
stage_id: int,
metrics: OrchestratorMetrics,
req_start_ts: dict[str, float],
wall_start_ts: float,
final_stage_id_for_e2e: int,
) -> OmniRequestOutput | None:
req_id = result.get("request_id")
engine_outputs = result.get("engine_outputs")
stage_durations = getattr(result["engine_outputs"], "stage_durations", {})
peak_memory_mb = getattr(result["engine_outputs"], "peak_memory_mb", 0.0)
req_id = result.request_id
engine_outputs = result.engine_outputs
stage_durations = getattr(engine_outputs, "stage_durations", {})
peak_memory_mb = getattr(engine_outputs, "peak_memory_mb", 0.0)
# Merge AR stage timing from OrchestratorAggregator.stage_events
if self._enable_ar_profiler:
@@ -408,7 +399,7 @@ class OmniBase(PDDisaggregationMixin):
stage_durations[f"ar_stage_{evt.stage_id}"] = evt.stage_gen_time_ms / 1000.0
# Merge pipeline timings from Orchestrator into stage_durations
_m = result.get("metrics")
_m = result.metrics
if _m is not None and hasattr(_m, "pipeline_timings") and _m.pipeline_timings:
for key, value in _m.pipeline_timings.items():
if key not in stage_durations:
@@ -425,18 +416,18 @@ class OmniBase(PDDisaggregationMixin):
finished = engine_outputs.finished
submit_ts = result.get("stage_submit_ts")
submit_ts = result.stage_submit_ts
now = time.time()
if metrics.stage_first_ts[stage_id] is None:
metrics.stage_first_ts[stage_id] = submit_ts if submit_ts is not None else now
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, now)
_m = result.get("metrics")
_m = result.metrics
if finished and _m is not None:
metrics.on_stage_metrics(stage_id, req_id, _m)
stage_meta = self.engine.get_stage_metadata(stage_id)
if not stage_meta["final_output"]:
if not stage_meta.final_output:
return None
try:
@@ -450,7 +441,7 @@ class OmniBase(PDDisaggregationMixin):
except Exception:
logger.exception("[%s] Finalize request handling error", self.__class__.__name__)
output_type = getattr(engine_outputs, "final_output_type", stage_meta["final_output_type"])
output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type)
images = getattr(engine_outputs, "images", []) if output_type == "image" else []
return OmniRequestOutput(
request_id=req_id or "",