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

[large-scale-serving] Integrate OmniCoordinator into stage engine pipeline (#3569)

Browse Source 
Signed-off-by: chickeyton <ngton2014@gmail.com>
Signed-off-by: herotai214 <herotai214@gmail.com>
Co-authored-by: herotai214 <herotai214@gmail.com>
This commit is contained in:
chickeytonandherotai214 authored and GitHub committed 2026-05-19 12:08:20 +08:00
1 parent ca9fd0b71c
commit e277feacaf
33 files changed
+3033 -782
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/engine/test_async_omni_engine_stage_init.py
+2 -12
View File
@@ -171,18 +171,6 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
assert cleanup_clients == [shared, extra]
def test_initialize_stages_rejects_non_diffusion_replicas_in_single_stage_mode():
engine = object.__new__(AsyncOmniEngine)
engine.single_stage_mode = True
engine.stage_configs = [types.SimpleNamespace(stage_id=0, runtime={"num_replicas": 2})]
with pytest.raises(
ValueError,
match="single_stage_mode only supports num_replicas > 1 for diffusion stages",
):
engine._validate_single_stage_mode_replica_constraints()
def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
import vllm_omni.engine.async_omni_engine as engine_mod
from vllm_omni.platforms import current_omni_platform
@@ -427,6 +415,7 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep
engine.single_stage_mode = False
engine._single_stage_id_filter = None
engine._omni_master_server = None
engine._coordinator_runtime = None
engine.stage_configs = [types.SimpleNamespace()]
cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
@@ -471,6 +460,7 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_
engine.single_stage_mode = False
engine._single_stage_id_filter = None
engine._omni_master_server = None
engine._coordinator_runtime = None
engine.stage_configs = [types.SimpleNamespace()]
cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
tests/engine/test_orchestrator.py
+1 -1
View File
@@ -1052,7 +1052,7 @@ async def test_stage_pool_abort_requests_logs_when_binding_is_missing(caplog) ->
target_logger.setLevel(prev_level)
assert not stage0.abort_calls
assert "abort: no binding for req=missing-req in stage-0" in caplog.text
assert "abort: no live binding for req=missing-req in stage-0" in caplog.text
@pytest.mark.asyncio
tests/engine/test_single_stage_mode.py
+54 -46
View File
@@ -203,6 +203,32 @@ class TestOmniMasterServerAllocation:
server = OmniMasterServer(master_address="127.0.0.1", master_port=15004, stage_ids=[3])
assert server.get_allocation(3) is server._stage_routes[(3, 0)]
def test_next_free_replica_id_skips_head_local_slot_until_filled(self):
# Head pre-allocates slot (0, 0) for its own register_stage_with_omni_master
# call. A same-host headless that registers with auto-assign BEFORE the
# head's own registration must NOT be handed slot 0 — it should land on
# slot 1 instead. Without the head_local_replicas reservation,
# _next_free_replica_id would see (0, 0) absent from _stage_configs and
# return 0, colliding with the head's own bound sockets.
server = OmniMasterServer(
master_address="127.0.0.1",
master_port=15010,
stage_ids=[0],
head_local_replicas={0: [0]},
)
assert server._next_free_replica_id(0) == 1
def test_next_free_replica_id_uses_remote_slot_when_unowned(self):
# When the head pre-allocates a remote-only slot (the head's _initialize_*
# path waits on get_stage_config), auto-assign SHOULD fill it so the
# head's wait unblocks. This is the original behavior, preserved.
server = OmniMasterServer(
master_address="127.0.0.1",
master_port=15011,
stage_ids=[0],
)
assert server._next_free_replica_id(0) == 0
# ---------------------------------------------------------------------------
# OmniMasterServer registration flow
@@ -233,36 +259,6 @@ class TestOmniMasterServerRegistration:
ctx.term()
server.stop()
def test_server_handles_unknown_stage_id_gracefully(self):
import msgspec
import zmq
from vllm.utils.network_utils import get_open_port
master_port = get_open_port()
server = OmniMasterServer(master_address="127.0.0.1", master_port=master_port, stage_ids=[0])
server.start()
ctx = zmq.Context()
bad_sock = None
good_sock = None
try:
bad_sock = ctx.socket(zmq.DEALER)
bad_sock.connect(f"tcp://127.0.0.1:{master_port}")
bad_sock.send(msgspec.msgpack.encode({"stage_id": 99}))
assert not bad_sock.poll(timeout=500)
good_sock = ctx.socket(zmq.DEALER)
good_sock.connect(f"tcp://127.0.0.1:{master_port}")
good_sock.send(msgspec.msgpack.encode({"stage_id": 0}))
assert good_sock.poll(timeout=2_000)
good_sock.recv()
finally:
for sock in (bad_sock, good_sock):
if sock is not None:
sock.close(linger=0)
ctx.term()
server.stop()
def test_registration_stores_stage_config(self):
import msgspec
import zmq
@@ -474,6 +470,8 @@ class TestSingleStageInitialization:
engine._omni_master_address = "127.0.0.1"
engine._omni_master_port = 26000
engine._omni_master_server = None
engine._omni_heartbeat_timeout = 30.0
engine._coordinator_runtime = None
engine.async_chunk = False
engine.diffusion_batch_size = 2
return engine
@@ -508,10 +506,16 @@ class TestSingleStageInitialization:
def test_start_omni_master_server_uses_configured_stage_ids(self, mocker: MockerFixture):
import vllm_omni.engine.async_omni_engine as engine_mod
from vllm_omni.distributed import omni_coordinator as omni_coord_mod
engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7)
mock_oms = mocker.Mock(spec=OmniMasterServer)
mocker.patch.object(engine_mod, "OmniMasterServer", return_value=mock_oms)
mocker.patch.object(
omni_coord_mod,
"OmniCoordinatorRuntime",
return_value=mocker.Mock(router_address="tcp://127.0.0.1:9999"),
)
stage_plans = [
_make_llm_plan(0, configured_stage_id=7, launch_mode="local"),
@@ -520,12 +524,16 @@ class TestSingleStageInitialization:
engine._start_omni_master_server(stage_plans)
engine_mod.OmniMasterServer.assert_called_once_with(
master_address="127.0.0.1",
master_port=26000,
stage_ids=[7, 11],
stage_replica_counts={7: 1, 11: 1},
)
call_kwargs = engine_mod.OmniMasterServer.call_args.kwargs
assert call_kwargs["master_address"] == "127.0.0.1"
assert call_kwargs["master_port"] == 26000
assert call_kwargs["stage_ids"] == [7, 11]
assert call_kwargs["stage_replica_counts"] == {7: 1, 11: 1}
# head_local_replicas reserves slots that the head will register
# itself (launch_mode == "local"). Stage 11 is remote, so it must
# NOT appear in the head-owned set — that slot is for the headless
# to fill via auto-assign.
assert call_kwargs["head_local_replicas"] == {7: [0]}
mock_oms.start.assert_called_once()
def test_start_omni_master_server_duplicate_stage_ids_raise(self):
@@ -545,7 +553,9 @@ class TestSingleStageInitialization:
with pytest.raises(ValueError, match="requires both"):
engine._start_omni_master_server([_make_llm_plan(0, configured_stage_id=7, launch_mode="local")])
def test_build_logical_stage_init_plans_clears_runtime_cfg_in_single_stage_mode(self, mocker: MockerFixture):
def test_build_logical_stage_init_plans_preserves_runtime_cfg_for_local_llm_in_single_stage_mode(
self, mocker: MockerFixture
):
import vllm_omni.engine.async_omni_engine as engine_mod
engine = self._build_engine([_make_stage_cfg(7)], single_stage_mode=True, stage_id_filter=7)
@@ -574,7 +584,7 @@ class TestSingleStageInitialization:
finally:
monkeypatch.undo()
assert stage_plans[0].replicas[0].metadata.runtime_cfg is None
assert stage_plans[0].replicas[0].metadata.runtime_cfg == {"devices": "0"}
def test_validate_single_stage_mode_allows_diffusion_replicas(self):
stage_cfg = _make_stage_cfg(0, stage_type="diffusion")
@@ -583,14 +593,6 @@ class TestSingleStageInitialization:
engine._validate_single_stage_mode_replica_constraints()
def test_validate_single_stage_mode_rejects_llm_replicas(self):
stage_cfg = _make_stage_cfg(0, stage_type="llm")
stage_cfg.runtime.num_replicas = 2
engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=0)
with pytest.raises(ValueError, match="only supports num_replicas > 1 for diffusion"):
engine._validate_single_stage_mode_replica_constraints()
def test_build_logical_stage_init_plans_preserves_diffusion_runtime_cfg_in_single_stage_mode(
self, mocker: MockerFixture
):
@@ -809,6 +811,7 @@ class TestSingleStageReplicaInitialization:
engine.model = "fake-model"
engine.single_stage_mode = True
engine._omni_master_server = mocker.Mock(spec=OmniMasterServer)
engine._coordinator_runtime = None
engine.stage_configs = []
fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace())
@@ -896,6 +899,7 @@ class TestSingleStageReplicaInitialization:
engine._omni_master_server = mocker.Mock(spec=OmniMasterServer)
engine._omni_master_server.address = "127.0.0.1"
engine._omni_master_server.port = 25000
engine._coordinator_runtime = None
plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0]
sentinel_client = SimpleNamespace()
@@ -950,6 +954,9 @@ class TestSingleStageReplicaInitialization:
handshake_address="tcp://hs",
request_address="tcp://req",
response_address="tcp://resp",
omni_coordinator_address=None,
omni_stage_id=5,
omni_replica_id=0,
)
mock_handshake.assert_called_once_with(proc, "tcp://hs", 60)
mock_from_addresses.assert_called_once_with(
@@ -972,6 +979,7 @@ class TestSingleStageReplicaInitialization:
engine._omni_master_server = mocker.Mock(spec=OmniMasterServer)
engine._omni_master_server.address = "127.0.0.1"
engine._omni_master_server.port = 25000
engine._coordinator_runtime = None
plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0]
proc = mocker.Mock()