diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index e5add7621..70df01238 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -625,11 +625,11 @@ class TestStageDiffusionClientErrorPropagation: client._tasks = {} client._shutting_down = False client._engine_dead = engine_dead - client._owns_process = True - client._proc = MagicMock( + proc = MagicMock( is_alive=MagicMock(return_value=proc_alive), exitcode=1, ) + client._proc_manager = SimpleNamespace(proc=proc) client._request_socket = MagicMock() client._response_socket = MagicMock() client._encoder = MagicMock() @@ -672,7 +672,7 @@ class TestStageDiffusionClientErrorPropagation: assert client.get_diffusion_output_nowait() is None def test_check_health_raises_when_proc_dead(self): - """``check_health`` detects a dead subprocess via ``_proc.is_alive()`` + """``check_health`` detects a dead subprocess via the manager's proc and raises ``EngineDeadError``, setting ``_engine_dead`` as a side effect.""" client = self._make_client(proc_alive=False) @@ -699,7 +699,7 @@ class TestStageDiffusionClientErrorPropagation: ``get_diffusion_output_nowait`` returns ``None`` and sets ``_shutting_down`` instead of raising.""" client = self._make_client(proc_alive=False) - client._proc.exitcode = 137 # SIGKILL (128 + 9) + client._proc_manager.proc.exitcode = 137 # SIGKILL (128 + 9) client._response_socket.recv.side_effect = zmq.Again result = client.get_diffusion_output_nowait() @@ -727,7 +727,6 @@ class TestStageDiffusionClientErrorPropagation: metadata, "tcp://req", "tcp://resp", - proc=None, batch_size=1, ) @@ -836,7 +835,7 @@ class TestStageDiffusionClientProcMonitor: client._engine_dead = False proc = _make_short_lived_process() - client._proc = proc + client._proc_manager = SimpleNamespace(proc=proc) client._start_proc_monitor() proc.join(5) diff --git a/tests/distributed/omni_coordinator/test_omni_coordinator.py b/tests/distributed/omni_coordinator/test_omni_coordinator.py index c7b1b597e..dfb291ccf 100644 --- a/tests/distributed/omni_coordinator/test_omni_coordinator.py +++ b/tests/distributed/omni_coordinator/test_omni_coordinator.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +import threading import time import pytest @@ -48,6 +49,33 @@ def _drain_sub_messages(sub: zmq.Socket, max_seconds: float = 0.4) -> None: _recv_replica_list(sub, timeout_ms=50) +def test_omni_coordinator_wait_for_shutdown_unblocks_on_close(): + router_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + pub_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + coordinator = OmniCoordinator( + router_zmq_addr=router_addr, + pub_zmq_addr=pub_addr, + heartbeat_timeout=1000.0, + ) + + waiter = threading.Thread(target=coordinator.wait_for_shutdown) + waiter.start() + time.sleep(0.05) + assert waiter.is_alive() + + coordinator.close() + waiter.join(timeout=1.0) + assert not waiter.is_alive() + + def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates(): """Rapid updates should be coalesced into fewer PUB messages.""" router_addr = get_engine_client_zmq_addr( diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index bc163873b..69440ac3b 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1,21 +1,20 @@ -import concurrent.futures +import contextlib import importlib import os -import threading import time import types -from contextlib import contextmanager import pytest from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec -from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_init_utils import ( LogicalStageInitPlan, ReplicaInitPlan, build_stage0_input_processor, compute_replica_layout, ) +from vllm_omni.engine.stage_runtime import StageRuntime pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -63,7 +62,7 @@ def _make_diffusion_metadata(stage_id: int, *, replica_id: int = 0, final_output def _make_llm_plan( stage_idx: int, *, - configured_stage_id: int, + stage_id: int, vllm_config: object, num_replicas: int = 1, final_output: bool = False, @@ -73,7 +72,7 @@ def _make_llm_plan( replicas: list[ReplicaInitPlan] = [] for replica_id in range(num_replicas): stage_cfg = types.SimpleNamespace( - stage_id=configured_stage_id, + stage_id=stage_id, stage_type="llm", runtime=types.SimpleNamespace(devices=str(replica_id)), engine_args={}, @@ -85,7 +84,7 @@ def _make_llm_plan( launch_mode="local", stage_cfg=stage_cfg, metadata=_make_llm_metadata( - configured_stage_id, + stage_id, replica_id=replica_id, final_output=final_output, final_output_type=final_output_type, @@ -99,7 +98,7 @@ def _make_llm_plan( ) return LogicalStageInitPlan( stage_idx=stage_idx, - configured_stage_id=configured_stage_id, + stage_id=stage_id, replicas=replicas, ) @@ -107,13 +106,13 @@ def _make_llm_plan( def _make_diffusion_plan( stage_idx: int, *, - configured_stage_id: int, + stage_id: int, num_replicas: int = 1, ): replicas: list[ReplicaInitPlan] = [] for replica_id in range(num_replicas): stage_cfg = types.SimpleNamespace( - stage_id=configured_stage_id, + stage_id=stage_id, stage_type="diffusion", runtime=types.SimpleNamespace(devices=str(replica_id)), engine_args={}, @@ -124,14 +123,14 @@ def _make_diffusion_plan( num_replicas=num_replicas, launch_mode="local", stage_cfg=stage_cfg, - metadata=_make_diffusion_metadata(configured_stage_id, replica_id=replica_id), + metadata=_make_diffusion_metadata(stage_id, replica_id=replica_id), stage_connector_spec={}, omni_kv_connector=(None, None, None), ) ) return LogicalStageInitPlan( stage_idx=stage_idx, - configured_stage_id=configured_stage_id, + stage_id=stage_id, replicas=replicas, ) @@ -165,7 +164,7 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): shared = types.SimpleNamespace(name="shared") extra = types.SimpleNamespace(name="extra") - cleanup_clients = AsyncOmniEngine._collect_initialized_clients_for_cleanup( + cleanup_clients = StageRuntime._collect_initialized_clients_for_cleanup( stage_pools=[types.SimpleNamespace(clients=[shared, None])], initialized_clients_by_stage={0: [shared], 1: [extra]}, ) @@ -173,35 +172,36 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): assert cleanup_clients == [shared, extra] -def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod +def test_initialize_local_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.engine.stage_engine_startup import StageReplicaResources from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.num_stages = 1 - engine.diffusion_batch_size = 1 - engine.single_stage_mode = False - engine._omni_master_server = None - engine.stage_configs = [] + runtime = StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) - plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0] + plan = _make_diffusion_plan(0, stage_id=0).replicas[0] env_var = current_omni_platform.device_control_env_var old_env = os.environ.get(env_var) os.environ[env_var] = "0,1" + runtime._init_visible_devices_baseline = "0,1" - @contextmanager - def _fake_stage_runtime_setup(_stage_id, _runtime_cfg): - current_omni_platform.set_device_control_env_var("1") - yield - - monkeypatch.setattr(engine_mod, "stage_runtime_setup", _fake_stage_runtime_setup) - monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None) - monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace()) + monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None) + monkeypatch.setattr( + runtime_mod, + "launch_diffusion_stage_replica", + lambda **_: (types.SimpleNamespace(), StageReplicaResources()), + ) try: - engine._initialize_diffusion_replica(plan, stage_init_timeout=1, stage_launch_lock=threading.Lock()) + runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) assert os.environ.get(env_var) == "0,1" finally: if old_env is None: @@ -210,68 +210,63 @@ def test_initialize_diffusion_replica_restores_device_visibility_after_local_ini os.environ[env_var] = old_env -def test_initialize_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod +def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.engine.stage_engine_startup import StageReplicaResources - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.num_stages = 1 - engine.diffusion_batch_size = 4 - engine.single_stage_mode = False - engine._omni_master_server = None - engine.stage_configs = [] + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=4, + async_chunk=False, + ) - plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0] + plan = _make_diffusion_plan(0, stage_id=0).replicas[0] captured: dict[str, object] = {} - @contextmanager - def _noop_stage_runtime_setup(*_): - yield + monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None) - monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup) - monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None) + def _capture_launch_diffusion_stage_replica(**kwargs): + captured["stage_id"] = kwargs["metadata"].stage_id + captured["stage_init_timeout"] = kwargs["stage_init_timeout"] + captured["batch_size"] = kwargs["batch_size"] + captured["use_inline"] = kwargs["use_inline"] + captured["omni_master_server"] = kwargs["omni_master_server"] + return types.SimpleNamespace(), StageReplicaResources() - def _capture_initialize_diffusion_stage( - stage_id, _model, _stage_cfg, _metadata, *, stage_init_timeout, batch_size, use_inline - ): - captured["stage_id"] = stage_id - captured["stage_init_timeout"] = stage_init_timeout - captured["batch_size"] = batch_size - captured["use_inline"] = use_inline - return types.SimpleNamespace() + monkeypatch.setattr(runtime_mod, "launch_diffusion_stage_replica", _capture_launch_diffusion_stage_replica) - monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", _capture_initialize_diffusion_stage) - - engine._initialize_diffusion_replica(plan, stage_init_timeout=302, stage_launch_lock=threading.Lock()) + runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=302) assert captured == { "stage_id": 0, "stage_init_timeout": 302, "batch_size": 4, "use_inline": True, + "omni_master_server": None, } -def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_input_processor(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod +def test_stage_runtime_initializes_stage_pools(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.config_path = "dummy-config" - engine.num_stages = 2 - engine.async_chunk = False - engine.diffusion_batch_size = 1 - engine.single_stage_mode = False - engine._single_stage_id_filter = None - engine._omni_master_server = None - engine.stage_configs = [types.SimpleNamespace(), types.SimpleNamespace()] + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace(), types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) stage_plans = [ - _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True), - _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, final_output=True), + _make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True), + _make_llm_plan(1, stage_id=1, vllm_config=cfg1, final_output=True), ] stage0_client_r0 = types.SimpleNamespace( @@ -302,51 +297,43 @@ def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_inpu stage0_output_processor = object() stage1_output_processor = object() - top_level_input_processor = object() - - monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) - monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) - monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2, 1], {})) - monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None)) - monkeypatch.setattr(engine, "_initialize_stage_replicas", lambda *_: initialized_clients) + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans) + monkeypatch.setattr(runtime, "_initialize_stage_replicas", lambda *_: initialized_clients) monkeypatch.setattr( - engine_mod, + runtime_mod, "build_llm_stage_output_processor", lambda plan, _cfg, **_kw: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor, ) - monkeypatch.setattr(engine_mod, "build_stage0_input_processor", lambda _cfg: top_level_input_processor) - engine._initialize_stages(stage_init_timeout=1) + runtime.initialize() - assert len(engine.stage_pools) == 2 - assert engine.input_processor is top_level_input_processor - assert engine.stage_clients == [stage0_client_r0, stage1_client_r0] - assert engine.stage_vllm_configs == [cfg0, cfg1] - assert engine.output_processors == [stage0_output_processor, stage1_output_processor] - assert engine.default_sampling_params_list == [ - stage0_client_r0.default_sampling_params, - stage1_client_r0.default_sampling_params, - ] - assert engine.stage_metadata == [ - StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm"), - StageRuntimeInfo(final_output=True, final_output_type=None, stage_type="llm"), - ] + assert len(runtime.stage_pools) == 2 + assert runtime.stage_pools[0].stage_client is stage0_client_r0 + assert runtime.stage_pools[1].stage_client is stage1_client_r0 + assert runtime.stage_pools[0].stage_vllm_config is cfg0 + assert runtime.stage_pools[1].stage_vllm_config is cfg1 + assert runtime.stage_pools[0].output_processor is stage0_output_processor + assert runtime.stage_pools[1].output_processor is stage1_output_processor def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.async_chunk = False - engine.single_stage_mode = False - engine._single_stage_id_filter = None - engine.stage_configs = [ - types.SimpleNamespace(stage_id=0, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="0")), - types.SimpleNamespace( - stage_id=1, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="1,2,3") - ), - ] + runtime = StageRuntime( + stage_configs=[ + types.SimpleNamespace( + stage_id=0, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="0") + ), + types.SimpleNamespace( + stage_id=1, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="1,2,3") + ), + ], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) metadata_by_stage = { 0: _make_llm_metadata(0), @@ -354,78 +341,46 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc } monkeypatch.setattr( - engine_mod, + runtime_mod, "extract_stage_metadata", lambda cfg: types.SimpleNamespace(**metadata_by_stage[cfg.stage_id].__dict__), ) - monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) - monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) - monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {}) monkeypatch.setattr( - engine_mod, + runtime_mod, "build_vllm_config", lambda stage_cfg, *_args, **_kwargs: (types.SimpleNamespace(tag=f"cfg-{stage_cfg.stage_id}"), object), ) - stage_plans, prompt_expand_func = engine._build_logical_stage_init_plans( + stage_plans = runtime._build_logical_stage_init_plans( omni_transfer_config=None, replicas_per_stage=[1, 3], replica_devices_map={1: ["1", "2", "3"]}, ) - assert prompt_expand_func is None - assert [plan.configured_stage_id for plan in stage_plans] == [0, 1] + assert [plan.stage_id for plan in stage_plans] == [0, 1] assert [replica.stage_cfg.runtime.devices for replica in stage_plans[1].replicas] == ["1", "2", "3"] assert [replica.replica_id for replica in stage_plans[1].replicas] == [0, 1, 2] assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas) -def test_initialize_stage_replicas_reuses_long_lived_executor(monkeypatch): - engine = object.__new__(AsyncOmniEngine) - engine._stage_init_executor = None - - cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) - stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg, num_replicas=1)] - - created_executors = [] - - class FakeExecutor: - def __init__(self, *args, **kwargs): - self.shutdown_calls = [] - created_executors.append(self) - - def submit(self, fn, *args, **kwargs): - future = concurrent.futures.Future() - future.set_result(fn(*args, **kwargs)) - return future - - def shutdown(self, wait=False, cancel_futures=False): - self.shutdown_calls.append((wait, cancel_futures)) - - monkeypatch.setattr(concurrent.futures, "ThreadPoolExecutor", FakeExecutor) - monkeypatch.setattr( - engine, - "_initialize_replica", - lambda plan, _stage_init_timeout, _stage_launch_lock: f"client-{plan.replica_id}", - ) - - first = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1) - second = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1) - - assert first == {0: ["client-0"]} - assert second == {0: ["client-0"]} - assert len(created_executors) == 1 - assert engine._stage_init_executor is created_executors[0] - - def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch): - engine = object.__new__(AsyncOmniEngine) + runtime = StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=123, + diffusion_batch_size=1, + async_chunk=False, + ) cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) stage_plans = [ - _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2), - _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, num_replicas=2), + _make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2), + _make_llm_plan(1, stage_id=1, vllm_config=cfg1, num_replicas=2), ] clients = { @@ -435,13 +390,13 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk (1, 1): types.SimpleNamespace(name="stage1-replica1"), } - def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): + def _initialize_replica(plan, _stage_init_timeout): time.sleep(0.02 * (3 - plan.metadata.stage_id - plan.replica_id)) return clients[(plan.metadata.stage_id, plan.replica_id)] - monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica) + monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica) - initialized_clients = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=123) + initialized_clients = runtime._initialize_stage_replicas(stage_plans, stage_init_timeout=123) assert initialized_clients == { 0: [clients[(0, 0)], clients[(0, 1)]], @@ -449,148 +404,120 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk } -def test_shutdown_closes_long_lived_stage_init_executor(): - engine = object.__new__(AsyncOmniEngine) - engine._shutdown_called = False - engine._weak_finalizer = None - engine._omni_master_server = None - engine._coordinator_runtime = None - engine.orchestrator_thread = types.SimpleNamespace(join=lambda timeout: None, is_alive=lambda: False) +def test_remote_replicas_use_distinct_init_group_keys(): + runtime = StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=123, + diffusion_batch_size=1, + async_chunk=False, + ) + plan = _make_llm_plan( + 1, + stage_id=1, + vllm_config=types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)), + num_replicas=2, + ) - class FakeQueue: - def __init__(self): - self.closed = False - self.sync_q = types.SimpleNamespace(put_nowait=lambda _msg: None) + for replica in plan.replicas: + replica.launch_mode = "remote" + replica.metadata.runtime_cfg = None - def close(self): - self.closed = True - - class FakeExecutor: - def __init__(self): - self.shutdown_calls = [] - - def shutdown(self, wait=False, cancel_futures=False): - self.shutdown_calls.append((wait, cancel_futures)) - - engine.request_queue = FakeQueue() - engine.output_queue = FakeQueue() - engine.rpc_output_queue = FakeQueue() - engine._stage_init_executor = FakeExecutor() - - engine.shutdown() - - assert engine._stage_init_executor is None - assert engine.request_queue.closed is True - assert engine.output_queue.closed is True - assert engine.rpc_output_queue.closed is True + assert [runtime._replica_init_group_key(replica) for replica in plan.replicas] == [ + "remote:1:0", + "remote:1:1", + ] def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.config_path = "dummy-config" - engine.num_stages = 1 - engine.async_chunk = False - engine.diffusion_batch_size = 1 - 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()] + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) - stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)] + stage_plans = [_make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2)] initialized_client = types.SimpleNamespace(shutdown=lambda: None) - monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) - monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) - monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {})) - monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None)) + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans) - def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): + def _initialize_replica(plan, _stage_init_timeout): if plan.replica_id == 0: return initialized_client time.sleep(0.05) raise RuntimeError("replica launch failed") - monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica) + monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica) captured_cleanup: list[list[object]] = [] def _capture_shutdown(clients): captured_cleanup.append(list(clients)) - monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown) + monkeypatch.setattr(runtime, "_shutdown_initialized_clients", _capture_shutdown) with pytest.raises(RuntimeError, match="replica launch failed"): - engine._initialize_stages(stage_init_timeout=1) + runtime.initialize() assert captured_cleanup == [[initialized_client]] def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_replica_failure(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.config_path = "dummy-config" - engine.num_stages = 1 - engine.async_chunk = False - engine.diffusion_batch_size = 1 - 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()] + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) - stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)] + stage_plans = [_make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2)] initialized_client = types.SimpleNamespace(shutdown=lambda: None) - replica1_started = threading.Event() - monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) - monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) - monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {})) - monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None)) + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans) - def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): - if plan.replica_id == 0: - assert replica1_started.wait(timeout=1) - raise RuntimeError("replica launch failed") - replica1_started.set() - time.sleep(0.05) - return initialized_client + def _initialize_stage_replicas(_stage_plans, _stage_init_timeout): + exc = RuntimeError("replica launch failed") + exc._initialized_clients_by_stage = {0: [None, initialized_client]} + raise exc - monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica) + monkeypatch.setattr(runtime, "_initialize_stage_replicas", _initialize_stage_replicas) captured_cleanup: list[list[object]] = [] def _capture_shutdown(clients): captured_cleanup.append(list(clients)) - monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown) + monkeypatch.setattr(runtime, "_shutdown_initialized_clients", _capture_shutdown) with pytest.raises(RuntimeError, match="replica launch failed"): - engine._initialize_stages(stage_init_timeout=1) + runtime.initialize() assert captured_cleanup == [[initialized_client]] -def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod +def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.single_stage_mode = False - engine._omni_master_server = None - engine.stage_configs = [] + runtime = StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=302, + diffusion_batch_size=1, + async_chunk=False, + ) fake_vllm_config = types.SimpleNamespace() fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None) - fake_proc = types.SimpleNamespace() captured_timeout: int | None = None plan = ReplicaInitPlan( @@ -603,34 +530,38 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand omni_kv_connector=(None, None, None), stage_vllm_config=fake_vllm_config, executor_class=object, + engine_args_dict={}, ) device_env_var = current_omni_platform.device_control_env_var prev_device_env = os.environ.get(device_env_var) os.environ[device_env_var] = "0" - @contextmanager - def _noop_stage_runtime_setup(*_): - yield - - monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup) - monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) - monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: []) - monkeypatch.setattr(engine_mod, "spawn_stage_core", lambda **_: (fake_addresses, fake_proc, "ipc://handshake")) - - def _capture_stage_timeout(_proc, _handshake_addr, _addresses, _vllm_cfg, handshake_timeout): + def _capture_acquire_device_locks(*_args): nonlocal captured_timeout - captured_timeout = handshake_timeout + captured_timeout = _args[2] + return [] - monkeypatch.setattr(engine_mod, "complete_stage_handshake", _capture_stage_timeout) + monkeypatch.setattr(runtime_mod, "acquire_device_locks", _capture_acquire_device_locks) + + from vllm_omni.engine.stage_engine_startup import StageReplicaResources + + @contextlib.contextmanager + def _fake_launch_stage_replica(**_kwargs): + yield StageReplicaResources( + manager=types.SimpleNamespace(shutdown=lambda: None), + addresses=fake_addresses, + ) + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica) monkeypatch.setattr( - engine_mod.StageEngineCoreClientBase, + runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", staticmethod(lambda **_: types.SimpleNamespace(shutdown=lambda: None)), ) try: - engine._initialize_llm_replica(plan, 302, threading.Lock()) + runtime._initialize_local_llm_replica(plan, 302) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) diff --git a/tests/engine/test_membership_controller.py b/tests/engine/test_membership_controller.py new file mode 100644 index 000000000..c00a7f646 --- /dev/null +++ b/tests/engine/test_membership_controller.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus +from vllm_omni.engine.membership_controller import MembershipController + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeHub: + def __init__(self, snapshots=None): + self.snapshots = list(snapshots or []) + self.closed = False + + def get_replica_list(self): + if self.snapshots: + return self.snapshots.pop(0) + return SimpleNamespace(replicas=[]) + + def close(self): + self.closed = True + + +class FakePool: + def __init__(self, stage_id: int): + self.stage_id = stage_id + self.clients = [] + self.added = [] + self.removed = [] + self.invalidated = [] + self.hub = None + self.lb = None + + def attach_hub(self, hub): + self.hub = hub + + def attach_load_balancer(self, lb): + self.lb = lb + + def add_client(self, input_addr, client): + self.added.append((input_addr, client)) + self.clients.append(client) + return len(self.clients) - 1 + + def invalidate_addr(self, input_addr): + self.invalidated.append(input_addr) + return ["req-1", "req-2"] + + def remove_client(self, input_addr): + self.removed.append(input_addr) + return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown")) + + +def _snapshot(*replicas): + return SimpleNamespace(replicas=list(replicas)) + + +def _replica(stage_id: int, input_addr: str, status=ReplicaStatus.UP): + return SimpleNamespace(stage_id=stage_id, input_addr=input_addr, status=status) + + +def _controller(monkeypatch, pool, hub, remote_replica_factory=None): + import vllm_omni.engine.membership_controller as membership_mod + + monkeypatch.setattr(membership_mod, "OmniCoordClientForHub", lambda _addr: hub) + if remote_replica_factory is None: + + def remote_replica_factory(stage_id, replica_id): + return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"}) + + return MembershipController( + stage_pools=[pool], + coordinator_pub_address="tcp://127.0.0.1:12345", + load_balancer_factory=lambda: object(), + remote_replica_factory=remote_replica_factory, + ) + + +@pytest.mark.asyncio +async def test_watch_replica_list_unregisters_disappeared_replicas(monkeypatch): + pool = FakePool(stage_id=0) + hub = FakeHub( + snapshots=[ + _snapshot(_replica(0, "tcp://gone")), + _snapshot(), + ] + ) + controller = _controller(monkeypatch, pool, hub) + controller.WATCH_INTERVAL_S = 0 + unregistered = [] + + async def _handle_unregister(stage_id, input_addr, *args, **kwargs): + unregistered.append((stage_id, input_addr)) + controller._shutdown_event.set() + + controller.handle_unregister = _handle_unregister # type: ignore[method-assign] + + await asyncio.wait_for(controller._watch_replica_list(), timeout=1) + await controller.drain_tasks(timeout=1) + + assert unregistered == [(0, "tcp://gone")] + + +@pytest.mark.asyncio +async def test_shutdown_closes_hub_then_cancels_watcher(monkeypatch): + pool = FakePool(stage_id=0) + hub = FakeHub() + controller = _controller(monkeypatch, pool, hub) + watcher = asyncio.create_task(asyncio.sleep(10)) + controller._watcher_task = watcher + + controller.shutdown() + await asyncio.sleep(0) + + assert hub.closed is True + assert watcher.cancelled() + + +@pytest.mark.asyncio +async def test_drain_tasks_waits_for_membership_tasks(monkeypatch): + pool = FakePool(stage_id=0) + controller = _controller(monkeypatch, pool, FakeHub()) + completed = [] + + async def _task(): + await asyncio.sleep(0) + completed.append(True) + + controller._spawn_task(_task(), label="unit") + await controller.drain_tasks(timeout=1) + + assert completed == [True] + + +@pytest.mark.asyncio +async def test_do_register_offloads_remote_factory(monkeypatch): + pool = FakePool(stage_id=0) + factory_thread_ids = [] + + def _factory(stage_id, replica_id): + factory_thread_ids.append(threading.get_ident()) + return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"}) + + controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory) + + await controller._do_register(0, 1) + + assert factory_thread_ids + assert factory_thread_ids[0] != threading.get_ident() + assert pool.added[0][0] == "tcp://stage-0-replica-1" diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index b86cc0af9..029c55fd1 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -11,11 +11,9 @@ from types import SimpleNamespace from typing import Any import janus -import psutil import pytest 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, @@ -27,7 +25,6 @@ from vllm_omni.engine.messages import ( 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 from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -180,42 +177,6 @@ class FakeOutputProcessor: return None -class _FakeProc: - pid = 1234 - - def __init__(self): - self.terminated = False - self.killed = False - self.join_calls = [] - - def is_alive(self): - return not self.terminated and not self.killed - - def terminate(self): - self.terminated = True - - def kill(self): - self.killed = True - - def join(self, timeout=None): - self.join_calls.append(timeout) - - -class _FakeChildProc: - def __init__(self): - self.terminated = False - self.killed = False - - def is_running(self): - return not self.terminated and not self.killed - - def terminate(self): - self.terminated = True - - def kill(self): - self.killed = True - - def _sampling_params(max_tokens: int = 4) -> SamplingParams: return SamplingParams(max_tokens=max_tokens) @@ -437,61 +398,6 @@ async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, requ 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): - fake_proc = _FakeProc() - fake_child = _FakeChildProc() - - class FakePsutilProcess: - def __init__(self, pid): - assert pid == fake_proc.pid - - def children(self, recursive=True): - assert recursive - return [fake_child] - - def fail_base_shutdown(self, **kwargs): - raise RuntimeError("base shutdown failed") - - monkeypatch.setattr(psutil, "Process", FakePsutilProcess) - monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: (list(procs), [])) - monkeypatch.setattr(AsyncMPClient, "shutdown", fail_base_shutdown) - - client = object.__new__(StageEngineCoreClient) - client._proc = fake_proc - - with pytest.raises(RuntimeError, match="base shutdown failed"): - client.shutdown() - - assert fake_proc.terminated - assert fake_proc.join_calls == [5] - assert fake_child.terminated - - -def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch): - fake_proc = _FakeProc() - fake_child = _FakeChildProc() - - class FakePsutilProcess: - def __init__(self, pid): - assert pid == fake_proc.pid - - def children(self, recursive=True): - assert recursive - return [fake_child] - - monkeypatch.setattr(psutil, "Process", FakePsutilProcess) - monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs))) - monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self, **kwargs: None) - - client = object.__new__(StageEngineCoreClient) - client._proc = fake_proc - - client.shutdown() - - assert fake_child.terminated - assert fake_child.killed - - @pytest.fixture def orchestrator_factory(): fixtures: list[OrchestratorFixture] = [] diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 2be2e9fb2..dc06c0c8d 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -3,8 +3,7 @@ from __future__ import annotations import os -import threading -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager from types import SimpleNamespace from typing import Any @@ -18,10 +17,14 @@ from vllm_omni.engine.stage_engine_startup import ( OmniMasterServer, StageAllocation, StageCoordinatorAddresses, + StageRegistrationResponse, + StageReplicaResources, + _launch_omni_core_engines, + connect_remote_diffusion_proc, connect_remote_engine_cores, - launch_omni_core_engines, ) from vllm_omni.engine.stage_init_utils import LogicalStageInitPlan, ReplicaInitPlan +from vllm_omni.engine.stage_runtime import DistStageRuntime, StageRuntime pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -43,13 +46,13 @@ def _make_stage_cfg(stage_id: int, stage_type: str = "llm"): def _make_llm_plan( stage_idx: int, *, - configured_stage_id: int, + stage_id: int, launch_mode: str, vllm_config: Any | None = None, ) -> LogicalStageInitPlan: - stage_cfg = _make_stage_cfg(configured_stage_id) + stage_cfg = _make_stage_cfg(stage_id) metadata = SimpleNamespace( - stage_id=configured_stage_id, + stage_id=stage_id, stage_type="llm", runtime_cfg={"devices": "0"}, prompt_expand_func=None, @@ -63,7 +66,7 @@ def _make_llm_plan( ) return LogicalStageInitPlan( stage_idx=stage_idx, - configured_stage_id=configured_stage_id, + stage_id=stage_id, replicas=[ ReplicaInitPlan( replica_id=0, @@ -76,6 +79,7 @@ def _make_llm_plan( stage_vllm_config=vllm_config or SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)), executor_class=object, + engine_args_dict={}, ) ], ) @@ -84,12 +88,12 @@ def _make_llm_plan( def _make_diffusion_plan( stage_idx: int, *, - configured_stage_id: int, + stage_id: int, launch_mode: str, ) -> LogicalStageInitPlan: - stage_cfg = _make_stage_cfg(configured_stage_id, stage_type="diffusion") + stage_cfg = _make_stage_cfg(stage_id, stage_type="diffusion") metadata = SimpleNamespace( - stage_id=configured_stage_id, + stage_id=stage_id, stage_type="diffusion", runtime_cfg={"devices": "0"}, prompt_expand_func=None, @@ -103,7 +107,7 @@ def _make_diffusion_plan( ) return LogicalStageInitPlan( stage_idx=stage_idx, - configured_stage_id=configured_stage_id, + stage_id=stage_id, replicas=[ ReplicaInitPlan( replica_id=0, @@ -415,7 +419,7 @@ class TestSingleStageModeDetection: def test_omni_master_server_starts_as_none(self, mocker: MockerFixture): engine = self._make_engine_no_thread(mocker) - assert engine._omni_master_server is None + assert not hasattr(engine, "_omni_master_server") # --------------------------------------------------------------------------- @@ -424,34 +428,28 @@ class TestSingleStageModeDetection: class TestSingleStageInitialization: - def _build_engine( - self, stage_cfgs: list[Any], *, single_stage_mode: bool, stage_id_filter: int | None - ) -> AsyncOmniEngine: - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.config_path = "/fake/stages.yaml" - engine.num_stages = len(stage_cfgs) - engine.stage_configs = stage_cfgs - engine.single_stage_mode = single_stage_mode - engine._single_stage_id_filter = stage_id_filter - 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 + def _build_runtime(self, stage_cfgs: list[Any], *, stage_id_filter: int | None) -> DistStageRuntime: + return DistStageRuntime( + stage_configs=stage_cfgs, + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=stage_id_filter, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) def test_build_logical_stage_init_plans_marks_non_matching_stage_remote(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod - stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11)] - engine = self._build_engine(stage_cfgs, single_stage_mode=True, stage_id_filter=7) + stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] + runtime = self._build_runtime(stage_cfgs, stage_id_filter=0) monkeypatch = pytest.MonkeyPatch() monkeypatch.setattr( - engine_mod, + runtime_mod, "extract_stage_metadata", lambda cfg: SimpleNamespace( stage_id=cfg.stage_id, @@ -460,24 +458,46 @@ class TestSingleStageInitialization: runtime_cfg={}, ), ) - monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) - monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) - monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) - monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (SimpleNamespace(), object)) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(runtime_mod, "build_vllm_config", lambda *_, **__: (SimpleNamespace(), object)) try: - stage_plans, _ = engine._build_logical_stage_init_plans(None, [1, 1], {}) + stage_plans = runtime._build_logical_stage_init_plans(None, [1, 1], {}) finally: monkeypatch.undo() assert [plan.replicas[0].launch_mode for plan in stage_plans] == ["local", "remote"] - def test_start_omni_master_server_uses_configured_stage_ids(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + def test_build_logical_stage_init_plans_rejects_non_contiguous_stage_ids(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = self._build_runtime([_make_stage_cfg(7), _make_stage_cfg(11)], stage_id_filter=7) + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr( + runtime_mod, + "extract_stage_metadata", + lambda cfg: SimpleNamespace( + stage_id=cfg.stage_id, + stage_type=getattr(cfg, "stage_type", "llm"), + prompt_expand_func=None, + runtime_cfg={}, + ), + ) + try: + with pytest.raises(ValueError, match="stage_id must match its position"): + runtime._build_logical_stage_init_plans(None, [1, 1], {}) + finally: + monkeypatch.undo() + + def test_start_omni_master_server_uses_stage_ids(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.distributed import omni_coordinator as omni_coord_mod - engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) + runtime = self._build_runtime([], stage_id_filter=0) mock_oms = mocker.Mock(spec=OmniMasterServer) - mocker.patch.object(engine_mod, "OmniMasterServer", return_value=mock_oms) + mocker.patch.object(runtime_mod, "OmniMasterServer", return_value=mock_oms) mocker.patch.object( omni_coord_mod, "OmniCoordinatorRuntime", @@ -485,51 +505,51 @@ class TestSingleStageInitialization: ) stage_plans = [ - _make_llm_plan(0, configured_stage_id=7, launch_mode="local"), - _make_diffusion_plan(1, configured_stage_id=11, launch_mode="remote"), + _make_llm_plan(0, stage_id=0, launch_mode="local"), + _make_diffusion_plan(1, stage_id=1, launch_mode="remote"), ] - engine._start_omni_master_server(stage_plans) + runtime._start_omni_master_server(stage_plans) - call_kwargs = engine_mod.OmniMasterServer.call_args.kwargs + call_kwargs = runtime_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} + assert call_kwargs["stage_ids"] == [0, 1] + assert call_kwargs["stage_replica_counts"] == {0: 1, 1: 1} # head_local_replicas reserves slots that the head will register - # itself (launch_mode == "local"). Stage 11 is remote, so it must + # itself (launch_mode == "local"). Stage 1 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]} + assert call_kwargs["head_local_replicas"] == {0: [0]} mock_oms.start.assert_called_once() def test_start_omni_master_server_duplicate_stage_ids_raise(self): - engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) + runtime = self._build_runtime([], stage_id_filter=0) stage_plans = [ - _make_llm_plan(0, configured_stage_id=7, launch_mode="local"), - _make_llm_plan(1, configured_stage_id=7, launch_mode="remote"), + _make_llm_plan(0, stage_id=0, launch_mode="local"), + _make_llm_plan(1, stage_id=0, launch_mode="remote"), ] with pytest.raises(ValueError, match="Duplicate stage_id"): - engine._start_omni_master_server(stage_plans) + runtime._start_omni_master_server(stage_plans) def test_start_omni_master_server_missing_address_raises(self): - engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) - engine._omni_master_address = None + runtime = self._build_runtime([], stage_id_filter=0) + runtime._omni_master_address = None with pytest.raises(ValueError, match="requires both"): - engine._start_omni_master_server([_make_llm_plan(0, configured_stage_id=7, launch_mode="local")]) + runtime._start_omni_master_server([_make_llm_plan(0, stage_id=0, launch_mode="local")]) 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 + import vllm_omni.engine.stage_runtime as runtime_mod - engine = self._build_engine([_make_stage_cfg(7)], single_stage_mode=True, stage_id_filter=7) + runtime = self._build_runtime([_make_stage_cfg(0)], stage_id_filter=0) monkeypatch = pytest.MonkeyPatch() monkeypatch.setattr( - engine_mod, + runtime_mod, "extract_stage_metadata", lambda cfg: SimpleNamespace( stage_id=cfg.stage_id, @@ -538,16 +558,16 @@ class TestSingleStageInitialization: runtime_cfg={"devices": "0"}, ), ) - monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) - monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) - monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {}) monkeypatch.setattr( - engine_mod, + runtime_mod, "build_vllm_config", lambda *_, **__: (SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)), object), ) try: - stage_plans, _ = engine._build_logical_stage_init_plans(None, [1], {}) + stage_plans = runtime._build_logical_stage_init_plans(None, [1], {}) finally: monkeypatch.undo() @@ -556,22 +576,39 @@ class TestSingleStageInitialization: def test_validate_single_stage_mode_allows_diffusion_replicas(self): stage_cfg = _make_stage_cfg(0, stage_type="diffusion") stage_cfg.runtime.num_replicas = 2 - engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=0) + runtime = self._build_runtime([stage_cfg], stage_id_filter=0) - engine._validate_single_stage_mode_replica_constraints() + runtime._validate_single_stage_mode_replica_constraints() + + def test_validate_single_stage_mode_does_not_swallow_num_replica_validation_errors(self): + class RuntimeConfig: + @property + def num_replicas(self): + return 1 + + @num_replicas.setter + def num_replicas(self, value): + raise ValueError("invalid num_replicas") + + stage_cfg = _make_stage_cfg(0, stage_type="llm") + stage_cfg.runtime = RuntimeConfig() + runtime = self._build_runtime([stage_cfg], stage_id_filter=0) + + with pytest.raises(ValueError, match="invalid num_replicas"): + runtime._validate_single_stage_mode_replica_constraints() def test_build_logical_stage_init_plans_preserves_diffusion_runtime_cfg_in_single_stage_mode( self, mocker: MockerFixture ): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod - stage_cfg = _make_stage_cfg(7, stage_type="diffusion") + stage_cfg = _make_stage_cfg(0, stage_type="diffusion") stage_cfg.runtime.devices = "0,1" - engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=7) + runtime = self._build_runtime([stage_cfg], stage_id_filter=0) monkeypatch = pytest.MonkeyPatch() monkeypatch.setattr( - engine_mod, + runtime_mod, "extract_stage_metadata", lambda cfg: SimpleNamespace( stage_id=cfg.stage_id, @@ -587,10 +624,10 @@ class TestSingleStageInitialization: replica_id=0, ), ) - monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) - monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) try: - stage_plans, _ = engine._build_logical_stage_init_plans(None, [2], {0: ["0", "1"]}) + stage_plans = runtime._build_logical_stage_init_plans(None, [2], {0: ["0", "1"]}) finally: monkeypatch.undo() @@ -600,11 +637,11 @@ class TestSingleStageInitialization: assert [replica.metadata.runtime_cfg for replica in replicas] == [{"devices": "0"}, {"devices": "1"}] def test_initialize_stages_calls_master_server_only_in_single_stage_mode(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod stage_cfgs = [_make_stage_cfg(0)] - engine = self._build_engine(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - stage_plan = _make_llm_plan(0, configured_stage_id=0, launch_mode="local") + runtime = self._build_runtime(stage_cfgs, stage_id_filter=0) + stage_plan = _make_llm_plan(0, stage_id=0, launch_mode="local") client = SimpleNamespace( stage_type="llm", is_comprehension=False, @@ -613,59 +650,61 @@ class TestSingleStageInitialization: default_sampling_params=SimpleNamespace(), ) - mocker.patch.object(engine_mod, "prepare_engine_environment") - mocker.patch.object(engine_mod, "load_omni_transfer_config_for_model", return_value=None) - mocker.patch.object(engine_mod, "compute_replica_layout", return_value=([1], {})) - mocker.patch.object(engine, "_build_logical_stage_init_plans", return_value=([stage_plan], None)) - mock_start = mocker.patch.object(engine, "_start_omni_master_server") - mocker.patch.object(engine, "_initialize_stage_replicas", return_value={0: [client]}) - mocker.patch.object(engine_mod, "build_stage0_input_processor", return_value=object()) - mocker.patch.object(engine_mod, "build_llm_stage_output_processor", return_value=object()) + mocker.patch.object(runtime_mod, "prepare_engine_environment") + mocker.patch.object(runtime_mod, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(runtime_mod, "compute_replica_layout", return_value=([1], {})) + mocker.patch.object(runtime, "_prepare_stage_plans", return_value=[stage_plan]) + mock_start = mocker.patch.object(runtime, "_start_omni_master_server") + mocker.patch.object(runtime, "_initialize_stage_replicas", return_value={0: [client]}) + mocker.patch.object(runtime_mod, "build_llm_stage_output_processor", return_value=object()) - engine._initialize_stages(stage_init_timeout=60) + runtime.initialize() mock_start.assert_called_once() - engine = self._build_engine(stage_cfgs, single_stage_mode=False, stage_id_filter=None) - mocker.patch.object(engine_mod, "prepare_engine_environment") - mocker.patch.object(engine_mod, "load_omni_transfer_config_for_model", return_value=None) - mocker.patch.object(engine_mod, "compute_replica_layout", return_value=([1], {})) - mocker.patch.object(engine, "_build_logical_stage_init_plans", return_value=([stage_plan], None)) - mock_start = mocker.patch.object(engine, "_start_omni_master_server") - mocker.patch.object(engine, "_initialize_stage_replicas", return_value={0: [client]}) - mocker.patch.object(engine_mod, "build_stage0_input_processor", return_value=object()) - mocker.patch.object(engine_mod, "build_llm_stage_output_processor", return_value=object()) + plain_runtime = StageRuntime( + stage_configs=stage_cfgs, + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + ) + mocker.patch.object(runtime_mod, "prepare_engine_environment") + mocker.patch.object(runtime_mod, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(runtime_mod, "compute_replica_layout", return_value=([1], {})) + mocker.patch.object(plain_runtime, "_prepare_stage_plans", return_value=[stage_plan]) + mocker.patch.object(plain_runtime, "_initialize_stage_replicas", return_value={0: [client]}) + mocker.patch.object(runtime_mod, "build_llm_stage_output_processor", return_value=object()) - engine._initialize_stages(stage_init_timeout=60) - mock_start.assert_not_called() + plain_runtime.initialize() def test_initialize_stages_stops_master_server_and_shuts_down_initialized_clients_on_failure( self, mocker: MockerFixture, ): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod stage_cfgs = [_make_stage_cfg(0)] - engine = self._build_engine(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - stage_plan = _make_llm_plan(0, configured_stage_id=0, launch_mode="local") + runtime = self._build_runtime(stage_cfgs, stage_id_filter=0) + stage_plan = _make_llm_plan(0, stage_id=0, launch_mode="local") initialized_client = mocker.Mock() mock_master = mocker.Mock(spec=OmniMasterServer) - mocker.patch.object(engine_mod, "prepare_engine_environment") - mocker.patch.object(engine_mod, "load_omni_transfer_config_for_model", return_value=None) - mocker.patch.object(engine_mod, "compute_replica_layout", return_value=([1], {})) - mocker.patch.object(engine, "_build_logical_stage_init_plans", return_value=([stage_plan], None)) + mocker.patch.object(runtime_mod, "prepare_engine_environment") + mocker.patch.object(runtime_mod, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(runtime_mod, "compute_replica_layout", return_value=([1], {})) + mocker.patch.object(runtime, "_prepare_stage_plans", return_value=[stage_plan]) def _start_master(_plans): - engine._omni_master_server = mock_master + runtime._omni_master_server = mock_master - mocker.patch.object(engine, "_start_omni_master_server", side_effect=_start_master) - mocker.patch.object(engine, "_initialize_stage_replicas", return_value={0: [initialized_client]}) - mocker.patch.object(engine_mod, "build_stage0_input_processor", return_value=object()) - mocker.patch.object(engine, "_assemble_stage_pools", side_effect=RuntimeError("assemble failed")) - mock_shutdown = mocker.patch.object(engine, "_shutdown_initialized_clients") + mocker.patch.object(runtime, "_start_omni_master_server", side_effect=_start_master) + mocker.patch.object(runtime, "_initialize_stage_replicas", return_value={0: [initialized_client]}) + mocker.patch.object(runtime, "_finalize_initialized_stages", side_effect=RuntimeError("assemble failed")) + mock_shutdown = mocker.patch.object(runtime, "_shutdown_initialized_clients") with pytest.raises(RuntimeError, match="assemble failed"): - engine._initialize_stages(stage_init_timeout=60) + runtime.initialize() mock_shutdown.assert_called_once_with([initialized_client]) mock_master.stop.assert_called_once() @@ -673,13 +712,21 @@ class TestSingleStageInitialization: class TestSingleStageReplicaInitialization: def test_initialize_llm_replica_remote_uses_connect_remote_engine_cores(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = True - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.get_stage_config.return_value = {"stage_id": 7, "stage_type": "llm"} + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "llm"} fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)) fake_addresses = SimpleNamespace( @@ -693,52 +740,74 @@ class TestSingleStageReplicaInitialization: def _fake_connect(**kwargs): events.append("enter") try: - yield fake_manager, fake_coordinator, fake_addresses, None + yield StageReplicaResources( + manager=fake_manager, + coordinator=fake_coordinator, + addresses=fake_addresses, + ) finally: events.append("exit") plan = _make_llm_plan( 0, - configured_stage_id=7, + stage_id=1, launch_mode="remote", vllm_config=fake_vllm_config, ).replicas[0] sentinel_client = SimpleNamespace() - mock_connect = mocker.patch.object(engine_mod, "connect_remote_engine_cores", side_effect=_fake_connect) + mock_connect = mocker.patch.object(runtime_mod, "connect_remote_engine_cores", side_effect=_fake_connect) mocker.patch.object( StageEngineCoreClientBase, "make_async_mp_client", side_effect=lambda **_: (events.append("attach"), sentinel_client)[1], ) - result = engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + result = runtime._initialize_remote_replica(plan, stage_init_timeout=60) assert result is sentinel_client - engine._omni_master_server.get_stage_config.assert_called_once_with(7, timeout_s=60, replica_id=0) - assert fake_vllm_config.parallel_config.data_parallel_size_local == 0 - assert mock_connect.call_args.kwargs["stage_id"] == 7 + runtime._omni_master_server.get_stage_config.assert_called_once_with(1, timeout_s=60, replica_id=0) + assert mock_connect.call_args.kwargs["vllm_config"].parallel_config.data_parallel_size_local == 0 + assert mock_connect.call_args.kwargs["stage_id"] == 1 assert mock_connect.call_args.kwargs["replica_id"] == 0 assert events == ["enter", "exit", "attach"] def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture): - engine = object.__new__(AsyncOmniEngine) - engine.single_stage_mode = True - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.get_stage_config.return_value = None + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.get_stage_config.return_value = None - plan = _make_llm_plan(0, configured_stage_id=7, launch_mode="remote").replicas[0] + plan = _make_llm_plan(0, stage_id=1, launch_mode="remote").replicas[0] with pytest.raises(ValueError, match="registered without stage config"): - engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + runtime._initialize_remote_replica(plan, stage_init_timeout=60) def test_initialize_llm_replica_remote_attach_failure_cleans_up_started_resources(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod - engine = object.__new__(AsyncOmniEngine) - engine.single_stage_mode = True - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.get_stage_config.return_value = {"stage_id": 7, "stage_type": "llm"} + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "llm"} fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)) fake_addresses = SimpleNamespace( @@ -749,15 +818,19 @@ class TestSingleStageReplicaInitialization: @contextmanager def _fake_connect(**kwargs): - yield fake_manager, fake_coordinator, fake_addresses, None + yield StageReplicaResources( + manager=fake_manager, + coordinator=fake_coordinator, + addresses=fake_addresses, + ) plan = _make_llm_plan( 0, - configured_stage_id=7, + stage_id=1, launch_mode="remote", vllm_config=fake_vllm_config, ).replicas[0] - mocker.patch.object(engine_mod, "connect_remote_engine_cores", side_effect=_fake_connect) + mocker.patch.object(runtime_mod, "connect_remote_engine_cores", side_effect=_fake_connect) mocker.patch.object( StageEngineCoreClientBase, "make_async_mp_client", @@ -765,21 +838,29 @@ class TestSingleStageReplicaInitialization: ) with pytest.raises(RuntimeError, match="attach failed"): - engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + runtime._initialize_remote_replica(plan, stage_init_timeout=60) fake_manager.shutdown.assert_called_once() fake_coordinator.shutdown.assert_called_once() - def test_initialize_llm_replica_single_stage_local_uses_launch_omni_core_engines(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + def test_initialize_local_llm_replica_uses_shared_launcher(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = True - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._coordinator_runtime = None - engine.stage_configs = [] + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._coordinator_runtime = None + runtime._stage_configs = [] fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace()) fake_addresses = SimpleNamespace( @@ -788,11 +869,14 @@ class TestSingleStageReplicaInitialization: @contextmanager def _fake_launch(**kwargs): - yield mocker.Mock(), None, fake_addresses + yield StageReplicaResources( + manager=mocker.Mock(), + addresses=fake_addresses, + ) plan = _make_llm_plan( 0, - configured_stage_id=3, + stage_id=0, launch_mode="local", vllm_config=fake_vllm_config, ).replicas[0] @@ -801,19 +885,19 @@ class TestSingleStageReplicaInitialization: device_env_var = current_omni_platform.device_control_env_var prev_device_env = os.environ.get(device_env_var) os.environ[device_env_var] = "0" + runtime._init_visible_devices_baseline = "0" - mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext()) - mocker.patch.object(engine_mod, "build_engine_args_dict", return_value={}) - mocker.patch.object(engine_mod, "acquire_device_locks", return_value=[]) - mocker.patch.object(engine_mod, "release_device_locks") - mock_launch = mocker.patch.object(engine_mod, "launch_omni_core_engines", side_effect=_fake_launch) + mocker.patch.object(runtime_mod, "build_engine_args_dict", return_value={}) + mocker.patch.object(runtime_mod, "acquire_device_locks", return_value=[]) + mocker.patch.object(runtime_mod, "release_device_locks") + mock_launch = mocker.patch.object(runtime_mod, "launch_stage_replica", side_effect=_fake_launch) mocker.patch.object( StageEngineCoreClientBase, "make_async_mp_client", side_effect=lambda **_: sentinel_client, ) try: - result = engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + result = runtime._initialize_local_llm_replica(plan, stage_init_timeout=60) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) @@ -821,85 +905,121 @@ class TestSingleStageReplicaInitialization: os.environ[device_env_var] = prev_device_env assert result is sentinel_client - assert mock_launch.call_args.kwargs["stage_id"] == 3 + assert mock_launch.call_args.kwargs["stage_id"] == 0 assert mock_launch.call_args.kwargs["stage_config"] is plan.stage_cfg assert mock_launch.call_args.kwargs["replica_id"] == 0 + assert mock_launch.call_args.kwargs["omni_master_server"] is runtime._omni_master_server def test_initialize_diffusion_replica_remote_uses_from_addresses(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod - engine = object.__new__(AsyncOmniEngine) - engine.single_stage_mode = True - engine.diffusion_batch_size = 4 - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.get_stage_config.return_value = {"stage_id": 11, "stage_type": "diffusion"} - engine._omni_master_server.get_zmq_addresses.return_value = SimpleNamespace( - inputs=["tcp://in"], - outputs=["tcp://out"], + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=4, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "diffusion"} - remote_metadata = _make_diffusion_plan(1, configured_stage_id=11, launch_mode="remote").replicas[0].metadata - plan = _make_diffusion_plan(1, configured_stage_id=11, launch_mode="remote").replicas[0] + def _fake_connect(**kwargs): + @contextmanager + def _ctx(): + yield StageReplicaResources( + addresses=SimpleNamespace( + inputs=["tcp://in"], + outputs=["tcp://out"], + ) + ) + + return _ctx() + + remote_metadata = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0].metadata + plan = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0] sentinel_client = SimpleNamespace() - mocker.patch.object(engine_mod, "extract_stage_metadata", return_value=remote_metadata) - mock_from_addresses = mocker.patch.object( - engine_mod.StageDiffusionClient, "from_addresses", return_value=sentinel_client + mocker.patch.object(runtime_mod, "extract_stage_metadata", return_value=remote_metadata) + mock_connect = mocker.patch.object(runtime_mod, "connect_remote_diffusion_proc", side_effect=_fake_connect) + mock_from_addresses = mocker.patch( + "vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses", + return_value=sentinel_client, ) - result = engine._initialize_diffusion_replica(plan, stage_init_timeout=60, stage_launch_lock=threading.Lock()) + result = runtime._initialize_remote_replica(plan, stage_init_timeout=60) assert result is sentinel_client - engine._omni_master_server.get_stage_config.assert_called_once_with(11, timeout_s=60, replica_id=0) - engine._omni_master_server.get_zmq_addresses.assert_called_once_with(11, replica_id=0) + runtime._omni_master_server.get_stage_config.assert_called_once_with(1, timeout_s=60, replica_id=0) + mock_connect.assert_called_once_with( + omni_master_server=runtime._omni_master_server, + stage_id=1, + replica_id=0, + ) mock_from_addresses.assert_called_once() - def test_initialize_diffusion_replica_single_stage_local_registers_with_master(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + def test_initialize_local_diffusion_replica_registers_with_master(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = True - engine.diffusion_batch_size = 4 - engine.stage_configs = [] - 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 + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=4, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.address = "127.0.0.1" + runtime._omni_master_server.port = 25000 + runtime._coordinator_runtime = None - plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0] + plan = _make_diffusion_plan(0, stage_id=0, launch_mode="local").replicas[0] sentinel_client = SimpleNamespace() proc = mocker.Mock() device_env_var = current_omni_platform.device_control_env_var prev_device_env = os.environ.get(device_env_var) os.environ[device_env_var] = "0" + runtime._init_visible_devices_baseline = "0" - mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext()) - mocker.patch.object(engine_mod, "inject_kv_stage_info") - mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config") - mock_register = mocker.patch.object( - engine_mod, - "register_stage_with_omni_master", - return_value=("tcp://hs", "tcp://req", "tcp://resp"), + mocker.patch.object(runtime_mod, "inject_kv_stage_info") + mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config") + mock_register = mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + replica_id=0, + coordinator_router_address=None, + ), ) - mock_spawn = mocker.patch.object( - engine_mod, - "spawn_diffusion_proc", - return_value=(proc, None, None, None), + fake_manager = SimpleNamespace( + proc=proc, + addresses=SimpleNamespace( + inputs=["tcp://127.0.0.1:26002"], + outputs=["tcp://127.0.0.1:26003"], + ), ) - mock_handshake = mocker.patch.object(engine_mod, "complete_diffusion_handshake") - mock_from_addresses = mocker.patch.object( - engine_mod.StageDiffusionClient, - "from_addresses", + mock_manager = mocker.patch( + "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager", + return_value=fake_manager, + ) + mock_from_addresses = mocker.patch( + "vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses", return_value=sentinel_client, ) try: - result = engine._initialize_diffusion_replica( - plan, stage_init_timeout=60, stage_launch_lock=threading.Lock() - ) + result = runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=60) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) @@ -910,78 +1030,84 @@ class TestSingleStageReplicaInitialization: mock_register.assert_called_once_with( omni_master_address="127.0.0.1", omni_master_port=25000, - omni_stage_id=5, + omni_stage_id=0, omni_stage_config=plan.stage_cfg, - return_addresses=True, replica_id=0, ) - mock_spawn.assert_called_once_with( - "fake-model", - "diffusion-config", - handshake_address="tcp://hs", - request_address="tcp://req", - response_address="tcp://resp", + mock_manager.assert_called_once_with( + model="fake-model", + od_config="diffusion-config", + stage_init_timeout=60, + handshake_address="tcp://127.0.0.1:26001", + addresses=mocker.ANY, omni_coordinator_address=None, - omni_stage_id=5, + omni_stage_id=0, omni_replica_id=0, ) - mock_handshake.assert_called_once_with(proc, "tcp://hs", 60) + assert mock_manager.call_args.kwargs["addresses"].inputs == ["tcp://127.0.0.1:26002"] + assert mock_manager.call_args.kwargs["addresses"].outputs == ["tcp://127.0.0.1:26003"] mock_from_addresses.assert_called_once_with( plan.metadata, - request_address="tcp://req", - response_address="tcp://resp", - proc=proc, + request_address="tcp://127.0.0.1:26002", + response_address="tcp://127.0.0.1:26003", + proc_manager=mocker.ANY, batch_size=4, ) + assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc - def test_initialize_diffusion_replica_local_failure_terminates_proc(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + def test_initialize_local_diffusion_replica_failure_terminates_proc(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = True - engine.diffusion_batch_size = 4 - engine.stage_configs = [] - 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 + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=4, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.address = "127.0.0.1" + runtime._omni_master_server.port = 25000 + runtime._coordinator_runtime = None - plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0] - proc = mocker.Mock() + plan = _make_diffusion_plan(0, stage_id=0, launch_mode="local").replicas[0] device_env_var = current_omni_platform.device_control_env_var prev_device_env = os.environ.get(device_env_var) os.environ[device_env_var] = "0" + runtime._init_visible_devices_baseline = "0" - mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext()) - mocker.patch.object(engine_mod, "inject_kv_stage_info") - mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config") - mocker.patch.object( - engine_mod, - "register_stage_with_omni_master", - return_value=("tcp://hs", "tcp://req", "tcp://resp"), + mocker.patch.object(runtime_mod, "inject_kv_stage_info") + mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config") + mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + replica_id=0, + coordinator_router_address=None, + ), ) - mocker.patch.object( - engine_mod, - "spawn_diffusion_proc", - return_value=(proc, None, None, None), + mocker.patch( + "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager", + side_effect=RuntimeError("handshake failed"), ) - mocker.patch.object(engine_mod, "complete_diffusion_handshake", side_effect=RuntimeError("handshake failed")) - mock_terminate = mocker.patch.object(engine_mod, "terminate_alive_proc") try: with pytest.raises(RuntimeError, match="handshake failed"): - engine._initialize_diffusion_replica(plan, stage_init_timeout=60, stage_launch_lock=threading.Lock()) + runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=60) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) else: os.environ[device_env_var] = prev_device_env - mock_terminate.assert_called_once_with(proc) - # --------------------------------------------------------------------------- # Stage engine startup helpers @@ -1025,14 +1151,16 @@ class TestConnectRemoteEngineCoresCoordinator: yield mocker.Mock() mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) - mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") + mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") with connect_remote_engine_cores( vllm_config=vllm_config, omni_master_server=omni_master_server, stage_id=7, replica_id=2, - ) as (_, yielded_coordinator, yielded_addresses, _tensor_queue): - assert yielded_coordinator is None + ) as resources: + assert resources.coordinator is None + yielded_addresses = resources.addresses + assert yielded_addresses is not None assert yielded_addresses.coordinator_input == "tcp://coord-in" assert yielded_addresses.coordinator_output == "tcp://coord-out" assert yielded_addresses.frontend_stats_publish_address == "tcp://stats" @@ -1058,17 +1186,48 @@ class TestConnectRemoteEngineCoresCoordinator: yield mocker.Mock() mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) - mocker.patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") + mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") with connect_remote_engine_cores( vllm_config=vllm_config, omni_master_server=omni_master_server, stage_id=7, - ) as (_, yielded_coordinator, yielded_addresses, _tensor_queue): - assert yielded_coordinator is None + ) as resources: + assert resources.coordinator is None + yielded_addresses = resources.addresses + assert yielded_addresses is not None assert yielded_addresses.coordinator_input is None assert yielded_addresses.coordinator_output is None assert yielded_addresses.frontend_stats_publish_address is None + def test_connect_remote_diffusion_proc_waits_for_headless_remote(self, mocker: MockerFixture): + omni_master_server = mocker.Mock(spec=OmniMasterServer) + omni_master_server.get_zmq_addresses.return_value = EngineZmqAddresses( + inputs=["tcp://client-in"], + outputs=["tcp://client-out"], + ) + omni_master_server.get_allocation.return_value = mocker.Mock(handshake_bind_address="tcp://127.0.0.1:26001") + + @contextmanager + def fake_socket_ctx(*args, **kwargs): + yield mocker.Mock() + + mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) + mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") + + with connect_remote_diffusion_proc( + omni_master_server=omni_master_server, + stage_id=7, + replica_id=2, + ) as resources: + assert resources.addresses is omni_master_server.get_zmq_addresses.return_value + + omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2) + omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2) + mock_wait.assert_called_once() + _, _, core_engines, parallel_config, *_ = mock_wait.call_args.args + assert core_engines[0].local is False + assert parallel_config.data_parallel_size_local == 0 + class TestLaunchOmniCoreEngines: def test_registers_stage_once_and_reuses_handshake_for_all_local_engines(self, mocker: MockerFixture): @@ -1093,7 +1252,13 @@ class TestLaunchOmniCoreEngines: mock_register = mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://client-in", + output_address="tcp://client-out", + replica_id=2, + coordinator_router_address=None, + ), ) mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) mock_manager_cls = mocker.patch( @@ -1101,7 +1266,7 @@ class TestLaunchOmniCoreEngines: return_value=local_engine_manager, ) mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") - with launch_omni_core_engines( + with _launch_omni_core_engines( vllm_config=vllm_config, executor_class=mocker.Mock(), log_stats=False, @@ -1164,7 +1329,13 @@ class TestLaunchOmniCoreEngines: mocker.patch("vllm_omni.engine.stage_engine_startup.DPCoordinator", return_value=coordinator) mock_register = mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://client-in", + output_address="tcp://client-out", + replica_id=3, + coordinator_router_address=None, + ), ) mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) mocker.patch( @@ -1172,7 +1343,7 @@ class TestLaunchOmniCoreEngines: return_value=mocker.Mock(), ) mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") - with launch_omni_core_engines( + with _launch_omni_core_engines( vllm_config=vllm_config, executor_class=mocker.Mock(), log_stats=False, diff --git a/tests/engine/test_stage_engine_core_client.py b/tests/engine/test_stage_engine_core_client.py index dde0927af..36c95cde7 100644 --- a/tests/engine/test_stage_engine_core_client.py +++ b/tests/engine/test_stage_engine_core_client.py @@ -1,15 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for StageEngineCoreClient.check_health(). - -Uses object.__new__ to construct a minimal client — check_health only -touches self.resources, self.stage_id, and self._proc. -""" +"""Tests for StageEngineCoreClient.check_health().""" from __future__ import annotations from types import SimpleNamespace -from unittest.mock import MagicMock import pytest from vllm.v1.engine.exceptions import EngineDeadError @@ -19,28 +14,19 @@ from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _make_client(*, engine_dead=False, proc_alive=True): +def _make_client(*, engine_dead=False): client = object.__new__(StageEngineCoreClient) client.stage_id = 0 client.resources = SimpleNamespace(engine_dead=engine_dead) - client._proc = MagicMock(is_alive=MagicMock(return_value=proc_alive), exitcode=1) return client def test_check_health_passes_when_alive(): - client = _make_client(engine_dead=False, proc_alive=True) + client = _make_client(engine_dead=False) client.check_health() # no exception def test_check_health_raises_when_resources_engine_dead(): - client = _make_client(engine_dead=True, proc_alive=True) + client = _make_client(engine_dead=True) with pytest.raises(EngineDeadError, match="engine core is dead"): client.check_health() - - -def test_check_health_raises_when_proc_not_alive(): - client = _make_client(engine_dead=False, proc_alive=False) - with pytest.raises(EngineDeadError, match="not alive"): - client.check_health() - # Verify it set resources.engine_dead as a side effect - assert client.resources.engine_dead is True diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index e05983d45..6c4445306 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -119,7 +119,7 @@ def _make_stage_cfg(stage_id: int, stage_type: str) -> SimpleNamespace: def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: MockerFixture) -> None: """LLM headless: each loop iteration registers with auto-assigned replica_id (master picks a free slot) and spawns one - ``OmniCoreEngineProcManager`` per local replica.""" + ``StageEngineCoreProcManager`` per local replica.""" from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse stage_cfg = _make_stage_cfg(0, stage_type="llm") @@ -159,7 +159,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker ), ) mock_manager_cls = mocker.patch( - "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager", + "vllm_omni.engine.stage_engine_core_proc_manager.StageEngineCoreProcManager", return_value=engine_manager, ) mocker.patch("signal.signal") @@ -168,8 +168,8 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker # The launcher must request auto-assignment (replica_id=None) and the # full response so it can wire the master-allocated coordinator into the - # spawned subprocess. ``replica_binds_sockets=False`` is required for LLM - # because the head binds all three sockets (handshake, input, output). + # spawned subprocess. LLM uses head-owned sockets: the head binds all + # three sockets (handshake, input, output) and the worker connects. assert mock_register.call_count == 1 kwargs = mock_register.call_args.kwargs assert kwargs["omni_master_address"] == "127.0.0.1" @@ -177,8 +177,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker assert kwargs["omni_stage_id"] == 0 assert kwargs["omni_stage_config"] is stage_cfg assert kwargs["replica_id"] is None - assert kwargs["return_full_response"] is True - assert kwargs["replica_binds_sockets"] is False + assert "socket_ownership" not in kwargs assert mock_manager_cls.call_count == 1 mgr_kwargs = mock_manager_cls.call_args.kwargs @@ -239,7 +238,7 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo ], ) mock_manager_cls = mocker.patch( - "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager", + "vllm_omni.engine.stage_engine_core_proc_manager.StageEngineCoreProcManager", side_effect=[manager_a, manager_b], ) mocker.patch("signal.signal") @@ -294,18 +293,24 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) coordinator_router_address="tcp://127.0.0.1:26100", ), ) - mock_spawn = mocker.patch( - "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc", - return_value=(proc, None, None, None), + fake_manager = SimpleNamespace( + proc=proc, + addresses=SimpleNamespace( + inputs=["tcp://127.0.0.1:26002"], + outputs=["tcp://127.0.0.1:26003"], + ), + shutdown=mocker.Mock(), + ) + mock_manager = mocker.patch( + "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager.launch_headless", + return_value=fake_manager, ) - mock_handshake = mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") # Replace the blocking wait with one that returns the only proc's sentinel # immediately so the test does not hang. mocker.patch( "multiprocessing.connection.wait", side_effect=lambda sentinels: [sentinels[0]], ) - mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc") mocker.patch("signal.signal") run_headless(_make_headless_args(stage_id=1)) @@ -321,17 +326,15 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) assert reg_kwargs["omni_stage_id"] == 1 assert reg_kwargs["omni_stage_config"] is stage_cfg assert reg_kwargs["replica_id"] is None - assert reg_kwargs["return_full_response"] is True + assert "socket_ownership" not in reg_kwargs - spawn_kwargs = mock_spawn.call_args.kwargs - assert spawn_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" - assert spawn_kwargs["request_address"] == "tcp://127.0.0.1:26002" - assert spawn_kwargs["response_address"] == "tcp://127.0.0.1:26003" - assert spawn_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100" - assert spawn_kwargs["omni_stage_id"] == 1 - assert spawn_kwargs["omni_replica_id"] == 0 - - mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001", 600) + manager_kwargs = mock_manager.call_args.kwargs + assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" + assert manager_kwargs["addresses"].inputs == ["tcp://127.0.0.1:26002"] + assert manager_kwargs["addresses"].outputs == ["tcp://127.0.0.1:26003"] + assert manager_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100" + assert manager_kwargs["omni_stage_id"] == 1 + assert manager_kwargs["omni_replica_id"] == 0 def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None: @@ -371,15 +374,13 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur ), ) mocker.patch( - "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc", - return_value=(proc, None, None, None), + "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager.launch_headless", + return_value=SimpleNamespace(proc=proc, shutdown=mocker.Mock()), ) - mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") mocker.patch( "multiprocessing.connection.wait", side_effect=lambda sentinels: [sentinels[0]], ) - mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc") mocker.patch("signal.signal") with pytest.raises(RuntimeError, match=r"exited with code 137"): diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index 05593d647..c4c0667df 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -4,7 +4,7 @@ import sys import pytest from pytest_mock import MockerFixture -from vllm_omni.entrypoints.stage_utils import _map_device_list, set_stage_devices +from vllm_omni.entrypoints.stage_utils import _map_device_list, resolve_stage_physical_devices, set_stage_devices def _make_dummy_torch(call_log): @@ -207,6 +207,13 @@ def test_map_device_list_multi_replica_offset_cvd(): assert _map_device_list(1, ["2", "3"], visible) == ["4", "5"] +@pytest.mark.core_model +@pytest.mark.cpu +def test_resolve_stage_physical_devices_uses_visible_baseline(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3") + assert resolve_stage_physical_devices(1, "1", visible_baseline="3,4,5") == "4" + + @pytest.mark.core_model @pytest.mark.cpu def test_map_device_list_raises_on_non_numeric(): diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 613fb64c7..8edf78bfb 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -1,8 +1,7 @@ """Stage Diffusion Client for vLLM-Omni multi-stage runtime. -Spawns StageDiffusionProc in a subprocess and communicates via ZMQ -(PUSH/PULL) to expose the same interface the Orchestrator expects -from any stage client. +Owns the frontend-side ZMQ sockets for StageDiffusionProc and exposes the +interface the Orchestrator expects from a stage client. """ from __future__ import annotations @@ -23,15 +22,14 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.stage_diffusion_proc import ( StageDiffusionProc, - complete_diffusion_handshake, - spawn_diffusion_proc, + StageDiffusionProcManager, ) 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.engine.stage_init_utils import StageMetadata from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -54,8 +52,17 @@ def create_diffusion_client( from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size) - return StageDiffusionClient( - model, od_config, metadata, stage_init_timeout=stage_init_timeout, batch_size=batch_size + proc_manager = StageDiffusionProcManager( + model=model, + od_config=od_config, + stage_init_timeout=stage_init_timeout, + ) + return StageDiffusionClient.from_addresses( + metadata, + request_address=proc_manager.addresses.inputs[0], + response_address=proc_manager.addresses.outputs[0], + proc_manager=proc_manager, + batch_size=batch_size, ) @@ -74,16 +81,20 @@ class StageDiffusionClient(StageClientBase): def __init__( self, - model: str, - od_config: OmniDiffusionConfig, metadata: StageMetadata, - stage_init_timeout: int, + request_address: str, + response_address: str, + *, + proc_manager: StageDiffusionProcManager | None = None, batch_size: int = 1, ) -> None: - # Spawn StageDiffusionProc subprocess and wait for READY. - proc, handshake_address, request_address, response_address = spawn_diffusion_proc(model, od_config) - complete_diffusion_handshake(proc, handshake_address, stage_init_timeout) - self._initialize_client(metadata, request_address, response_address, proc=proc, batch_size=batch_size) + self._initialize_client( + metadata, + request_address, + response_address, + proc_manager=proc_manager, + batch_size=batch_size, + ) @classmethod def from_addresses( @@ -92,19 +103,17 @@ class StageDiffusionClient(StageClientBase): request_address: str, response_address: str, *, - proc: Any = None, + proc_manager: StageDiffusionProcManager | None = None, batch_size: int = 1, ) -> StageDiffusionClient: """Create a client for an already-running diffusion subprocess.""" - client = cls.__new__(cls) - client._initialize_client( + return cls( metadata, request_address, response_address, - proc=proc, + proc_manager=proc_manager, batch_size=batch_size, ) - return client def _initialize_client( self, @@ -112,37 +121,12 @@ class StageDiffusionClient(StageClientBase): request_address: str, response_address: str, *, - proc: Any, + proc_manager: StageDiffusionProcManager | None = None, batch_size: int, ) -> None: - self.stage_id = metadata.stage_id - self.replica_id = metadata.replica_id - self.final_output = metadata.final_output - self.final_output_type = metadata.final_output_type - self.model_stage = getattr(metadata, "model_stage", None) - self.default_sampling_params = metadata.default_sampling_params - self.requires_multimodal_data = getattr(metadata, "requires_multimodal_data", False) - self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None) - self.engine_input_source = getattr(metadata, "engine_input_source", []) - self._proc = proc - self._owns_process = proc is not None - # Expose the ZMQ addresses on the instance so callers (e.g. - # ``StagePool._client_input_addr``) can identify the diffusion - # replica by its bound address. - self.request_address = request_address - self.response_address = response_address - - self._zmq_ctx = zmq.Context() - self._request_socket = self._zmq_ctx.socket(zmq.PUSH) - self._request_socket.connect(request_address) - self._response_socket = self._zmq_ctx.socket(zmq.PULL) - self._response_socket.connect(response_address) - - self._response_poller = zmq.asyncio.Poller() - self._response_poller.register(self._response_socket, zmq.POLLIN) - - self._encoder = OmniMsgpackEncoder() - self._decoder = OmniMsgpackDecoder() + self._set_stage_metadata(metadata) + self._proc_manager = proc_manager + self._connect_transport(request_address, response_address) self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue() self._rpc_results: dict[str, Any] = {} @@ -151,19 +135,48 @@ class StageDiffusionClient(StageClientBase): self._shutting_down = False self._engine_dead: bool = False - # Background thread to detect silent process death (SIGKILL, segfault) - # where the subprocess cannot send the ZMQ death sentinel. - # Mirrors MPClient.start_engine_core_monitor() in vLLM. - self._start_proc_monitor() + if self._proc_manager is not None: + self._start_proc_monitor() logger.info( "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s, batch_size=%d)", self.stage_id, self.replica_id, - self._owns_process, + self._proc_manager is not None, batch_size, ) + def _set_stage_metadata(self, metadata: StageMetadata) -> None: + self.stage_id = metadata.stage_id + self.replica_id = metadata.replica_id + self.final_output = metadata.final_output + self.final_output_type = metadata.final_output_type + self.model_stage = metadata.model_stage + self.default_sampling_params = metadata.default_sampling_params + self.prompt_expand_func = metadata.prompt_expand_func + self.requires_multimodal_data = getattr(metadata, "requires_multimodal_data", False) + self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None) + self.engine_input_source = getattr(metadata, "engine_input_source", []) + + def _connect_transport(self, request_address: str, response_address: str) -> None: + # Expose the ZMQ addresses on the instance so callers (e.g. + # ``StagePool._client_input_addr``) can identify the diffusion + # replica by its bound address. + self.request_address = request_address + self.response_address = response_address + + self._zmq_ctx = zmq.Context() + self._request_socket = self._zmq_ctx.socket(zmq.PUSH) + self._request_socket.bind(request_address) + self._response_socket = self._zmq_ctx.socket(zmq.PULL) + self._response_socket.bind(response_address) + + self._response_poller = zmq.asyncio.Poller() + self._response_poller.register(self._response_socket, zmq.POLLIN) + + self._encoder = OmniMsgpackEncoder() + self._decoder = OmniMsgpackDecoder() + # ------------------------------------------------------------------ # Process monitor (mirrors vLLM's MPClient.start_engine_core_monitor) # ------------------------------------------------------------------ @@ -175,7 +188,10 @@ class StageDiffusionClient(StageClientBase): (e.g. SIGKILL, segfault), this thread sets ``_engine_dead`` so subsequent calls raise ``EngineDeadError``. """ - proc = self._proc + # Background thread to detect silent process death (SIGKILL, segfault) + # where the subprocess cannot send the ZMQ death sentinel. + # Mirrors MPClient.start_engine_core_monitor() in vLLM. + proc = self._proc_manager.proc self_ref = weakref.ref(self) def _monitor() -> None: @@ -413,9 +429,12 @@ class StageDiffusionClient(StageClientBase): if self._shutting_down: return None raise EngineDeadError() - if not self._shutting_down and self._owns_process and self._proc is not None and not self._proc.is_alive(): + if self._proc_manager is None: + return None + proc = self._proc_manager.proc + if not self._shutting_down and not proc.is_alive(): self._engine_dead = True - exitcode = self._proc.exitcode + exitcode = proc.exitcode # One final drain – the last ZMQ frame may have arrived # between the first drain and the is_alive() check. self._drain_responses() @@ -490,11 +509,12 @@ class StageDiffusionClient(StageClientBase): result = self._rpc_results.pop(rpc_id, None) if result is not None: return result - if self._engine_dead or (self._owns_process and self._proc is not None and not self._proc.is_alive()): + proc = self._proc_manager.proc + if self._engine_dead or not proc.is_alive(): self._engine_dead = True raise EngineDeadError( f"StageDiffusionProc died while waiting for " - f"collective_rpc '{method}' (exit code {self._proc.exitcode})" + f"collective_rpc '{method}' (exit code {proc.exitcode})" ) if deadline is not None and time.monotonic() > deadline: raise TimeoutError(f"collective_rpc_async '{method}' timed out after {timeout}s") @@ -518,10 +538,13 @@ class StageDiffusionClient(StageClientBase): """ if self._engine_dead: raise EngineDeadError(f"Stage-{self.stage_id} diffusion subprocess is dead") - if self._proc is not None and not self._proc.is_alive(): + if self._proc_manager is None: + return + proc = self._proc_manager.proc + if not proc.is_alive(): self._engine_dead = True raise EngineDeadError( - f"Stage-{self.stage_id} diffusion subprocess is not alive (exit code: {self._proc.exitcode})." + f"Stage-{self.stage_id} diffusion subprocess is not alive (exit code: {proc.exitcode})." ) def shutdown(self) -> None: @@ -531,9 +554,8 @@ class StageDiffusionClient(StageClientBase): except Exception: pass - if self._owns_process and self._proc is not None and self._proc.is_alive(): - self._proc.join(timeout=10) - terminate_alive_proc(self._proc) + if self._proc_manager is not None and self._proc_manager.proc.is_alive(): + self._proc_manager.shutdown(timeout=10) self._request_socket.close(linger=0) self._response_socket.close(linger=0) diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 27c6f16a0..7530e61c9 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -8,10 +8,10 @@ from __future__ import annotations import asyncio import contextlib +import multiprocessing.connection import signal -import time from concurrent.futures import ThreadPoolExecutor -from multiprocessing.process import BaseProcess +from types import SimpleNamespace from typing import TYPE_CHECKING, Any import msgspec @@ -22,6 +22,8 @@ from PIL import Image from vllm.logger import init_logger from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx from vllm.utils.system_utils import get_mp_context +from vllm.v1.engine.core import EngineCoreProc +from vllm.v1.engine.utils import CoreEngine, EngineZmqAddresses, wait_for_engine_startup from vllm.v1.utils import shutdown from vllm_omni.diffusion.data import DiffusionRequestAbortedError @@ -360,10 +362,10 @@ class StageDiffusionProc: ctx = zmq.asyncio.Context() request_socket = ctx.socket(zmq.PULL) - request_socket.bind(request_address) + request_socket.connect(request_address) response_socket = ctx.socket(zmq.PUSH) - response_socket.bind(response_address) + response_socket.connect(response_address) encoder = OmniMsgpackEncoder() decoder = OmniMsgpackDecoder() @@ -612,15 +614,51 @@ class StageDiffusionProc: # Subprocess entry point # ------------------------------------------------------------------ + @staticmethod + def _open_startup_handshake( + handshake_address: str, + *, + local_client: bool, + headless: bool, + ) -> tuple[zmq.Context, zmq.Socket, EngineZmqAddresses]: + ctx = zmq.Context() + socket = ctx.socket(zmq.DEALER) + socket.setsockopt(zmq.IDENTITY, (0).to_bytes(2, "little")) + socket.connect(handshake_address) + addresses = EngineCoreProc.startup_handshake( + socket, + local_client=local_client, + headless=headless, + parallel_config=None, + ) + return ctx, socket, addresses + + @staticmethod + def _send_startup_ready( + handshake_socket: zmq.Socket, + *, + local_client: bool, + headless: bool, + ) -> None: + handshake_socket.send( + msgspec.msgpack.encode( + { + "status": "READY", + "local": local_client, + "headless": headless, + } + ) + ) + @classmethod def run_diffusion_proc( cls, model: str, od_config: OmniDiffusionConfig, handshake_address: str, - request_address: str, - response_address: str, *, + local_client: bool, + headless: bool, omni_coordinator_address: str | None = None, omni_stage_id: int | None = None, omni_replica_id: int = 0, @@ -651,21 +689,32 @@ class StageDiffusionProc: proc = cls(model, od_config) coord_client: OmniCoordClientForStage | None = None + handshake_ctx: zmq.Context | None = None + handshake_socket: zmq.Socket | None = None try: + handshake_ctx, handshake_socket, addresses = cls._open_startup_handshake( + handshake_address, + local_client=local_client, + headless=headless, + ) + request_address = addresses.inputs[0] + response_address = addresses.outputs[0] + proc.initialize() - # Send READY via handshake socket - handshake_ctx = zmq.Context() - handshake_socket = handshake_ctx.socket(zmq.DEALER) - handshake_socket.connect(handshake_address) - handshake_socket.send(msgspec.msgpack.encode({"status": "READY"})) + cls._send_startup_ready( + handshake_socket, + local_client=local_client, + headless=headless, + ) handshake_socket.close() handshake_ctx.term() + handshake_socket = None + handshake_ctx = None - # Wire OmniCoordClientForStage *after* READY so that the head - # has bound its head-side request/response sockets — the - # address pair we report is the same pair this proc binds to - # (request/response addresses passed in). + # Wire OmniCoordClientForStage *after* READY. The address pair is + # owned by the frontend client; this proc connects to it as the + # backend runtime. if omni_coordinator_address is not None: if omni_stage_id is None: raise ValueError("omni_stage_id must be provided when omni_coordinator_address is set") @@ -688,7 +737,6 @@ class StageDiffusionProc: omni_coordinator_address, ) - # Run async event loop asyncio.run(proc.run_loop(request_address, response_address)) except SystemExit: @@ -698,105 +746,138 @@ class StageDiffusionProc: logger.exception("StageDiffusionProc encountered a fatal error.") raise finally: + if handshake_socket is not None: + handshake_socket.close(linger=0) + if handshake_ctx is not None: + handshake_ctx.term() if coord_client is not None: with contextlib.suppress(RuntimeError): coord_client.close() proc.close() -# -- Free functions for backward compatibility with StageDiffusionClient ------ +class StageDiffusionProcManager: + """Owns a StageDiffusionProc subprocess. - -def spawn_diffusion_proc( - model: str, - od_config: OmniDiffusionConfig, - handshake_address: str | None = None, - request_address: str | None = None, - response_address: str | None = None, - *, - omni_coordinator_address: str | None = None, - omni_stage_id: int | None = None, - omni_replica_id: int = 0, -) -> tuple[BaseProcess, str, str, str]: - """Spawn a StageDiffusionProc subprocess. - - Returns ``(proc, handshake_address, request_address, response_address)``. - - Pass ``omni_coordinator_address`` / ``omni_stage_id`` / ``omni_replica_id`` - to have the subprocess publish heartbeats to an OmniCoordinator. + Mirrors the small process-lifecycle surface used by vLLM's + CoreEngineProcManager while keeping diffusion's custom wire protocol. """ - handshake_address = handshake_address or get_open_zmq_ipc_path() - request_address = request_address or get_open_zmq_ipc_path() - response_address = response_address or get_open_zmq_ipc_path() - ctx = get_mp_context() - proc = ctx.Process( - target=StageDiffusionProc.run_diffusion_proc, - name="StageDiffusionProc", - kwargs={ - "model": model, - "od_config": od_config, - "handshake_address": handshake_address, - "request_address": request_address, - "response_address": response_address, - "omni_coordinator_address": omni_coordinator_address, - "omni_stage_id": omni_stage_id, - "omni_replica_id": omni_replica_id, - }, - ) - proc.start() - # Wait for the process to become alive before returning. - deadline = time.monotonic() + 10 - while not proc.is_alive(): - if proc.exitcode is not None: - raise RuntimeError(f"StageDiffusionProc failed to start (exit code {proc.exitcode})") - if time.monotonic() > deadline: - raise TimeoutError("StageDiffusionProc did not become alive within 10s") - time.sleep(0.01) - return proc, handshake_address, request_address, response_address + def __init__( + self, + *, + model: str, + od_config: OmniDiffusionConfig, + stage_init_timeout: int, + handshake_address: str | None = None, + addresses: EngineZmqAddresses | None = None, + omni_coordinator_address: str | None = None, + omni_stage_id: int | None = None, + omni_replica_id: int = 0, + ) -> None: + handshake_address = handshake_address or get_open_zmq_ipc_path() + addresses = addresses or EngineZmqAddresses( + inputs=[get_open_zmq_ipc_path()], + outputs=[get_open_zmq_ipc_path()], + ) + ctx = get_mp_context() + proc = ctx.Process( + target=StageDiffusionProc.run_diffusion_proc, + name="StageDiffusionProc", + kwargs={ + "model": model, + "od_config": od_config, + "handshake_address": handshake_address, + "local_client": True, + "headless": False, + "omni_coordinator_address": omni_coordinator_address, + "omni_stage_id": omni_stage_id, + "omni_replica_id": omni_replica_id, + }, + ) + proc.start() + self.proc = proc + self.addresses = addresses + self.manager_stopped = False + self.failed_proc_name: str | None = None -def complete_diffusion_handshake( - proc: BaseProcess, - handshake_address: str, - handshake_timeout: int, -) -> None: - """Wait for the diffusion subprocess to signal READY. + self._wait_until_started(handshake_address, stage_init_timeout) - On failure the process is terminated before re-raising. - """ - try: - _perform_diffusion_handshake(proc, handshake_address, handshake_timeout) - except Exception: - shutdown([proc]) - raise + @classmethod + def launch_headless( + cls, + *, + model: str, + od_config: OmniDiffusionConfig, + handshake_address: str, + addresses: EngineZmqAddresses, + omni_coordinator_address: str | None, + omni_stage_id: int, + omni_replica_id: int, + ) -> StageDiffusionProcManager: + """Launch a headless diffusion backend that connects to head-owned sockets.""" + self = cls.__new__(cls) + ctx = get_mp_context() + proc = ctx.Process( + target=StageDiffusionProc.run_diffusion_proc, + name="StageDiffusionProc", + kwargs={ + "model": model, + "od_config": od_config, + "handshake_address": handshake_address, + "local_client": False, + "headless": True, + "omni_coordinator_address": omni_coordinator_address, + "omni_stage_id": omni_stage_id, + "omni_replica_id": omni_replica_id, + }, + ) + proc.start() + self.proc = proc + self.addresses = addresses + self.manager_stopped = False + self.failed_proc_name = None + return self - -def _perform_diffusion_handshake( - proc: BaseProcess, - handshake_address: str, - handshake_timeout: int, -) -> None: - """Run the handshake with the diffusion subprocess.""" - with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: - poller = zmq.Poller() - poller.register(handshake_socket, zmq.POLLIN) - poller.register(proc.sentinel, zmq.POLLIN) - - timeout_ms = handshake_timeout * 1000 - while True: - events = dict(poller.poll(timeout=timeout_ms)) - if not events: - raise TimeoutError( - f"Timed out waiting for READY from StageDiffusionProc after {handshake_timeout}s. " - f"This typically indicates model loading or warmup is taking too long. " - f"Consider increasing `stage_init_timeout` for large models." + def _wait_until_started(self, handshake_address: str, stage_init_timeout: int) -> None: + try: + with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: + wait_for_engine_startup( + handshake_socket, + self.addresses, + [CoreEngine(index=0, local=True)], + SimpleNamespace( + data_parallel_size_local=1, + data_parallel_hybrid_lb=False, + data_parallel_external_lb=False, + ), + False, + None, + self, + None, ) - if handshake_socket in events: - identity, raw = handshake_socket.recv_multipart() - msg = msgspec.msgpack.decode(raw) - if msg.get("status") == "READY": - return - raise RuntimeError(f"Expected READY, got: {msg}") - if proc.exitcode is not None: - raise RuntimeError(f"StageDiffusionProc died during handshake (exit code {proc.exitcode})") + except Exception: + shutdown([self.proc]) + raise + + def shutdown(self, timeout: float | None = None) -> None: + self.manager_stopped = True + shutdown([self.proc], timeout=timeout) + + def sentinels(self) -> list[int]: + return [self.proc.sentinel] + + def finished_procs(self) -> dict[str, int]: + if self.proc.exitcode is None: + return {} + return {self.proc.name: self.proc.exitcode} + + def monitor_engine_liveness(self) -> None: + try: + multiprocessing.connection.wait([self.proc.sentinel]) + except Exception: + return + if self.proc.exitcode not in (None, 0) and not self.manager_stopped: + self.failed_proc_name = self.proc.name + self.shutdown() diff --git a/vllm_omni/distributed/omni_coordinator/__init__.py b/vllm_omni/distributed/omni_coordinator/__init__.py index 326fb9b9a..b37c2bfb2 100644 --- a/vllm_omni/distributed/omni_coordinator/__init__.py +++ b/vllm_omni/distributed/omni_coordinator/__init__.py @@ -8,11 +8,13 @@ from .load_balancer import ( RandomBalancer, RoundRobinBalancer, Task, - build_load_balancer_factory, ) from .messages import ReplicaEvent, ReplicaInfo, ReplicaList, ReplicaStatus from .omni_coord_client_for_hub import OmniCoordClientForHub -from .omni_coord_client_for_stage import OmniCoordClientForStage +from .omni_coord_client_for_stage import ( + OmniCoordClientForStage, + create_stage_coord_client, +) from .omni_coordinator import OmniCoordinator from .runtime import OmniCoordinatorRuntime @@ -24,6 +26,7 @@ __all__ = [ "ReplicaInfo", "ReplicaList", "OmniCoordClientForStage", + "create_stage_coord_client", "OmniCoordClientForHub", "Task", "LoadBalancer", @@ -31,5 +34,4 @@ __all__ = [ "RandomBalancer", "RoundRobinBalancer", "LeastQueueLengthBalancer", - "build_load_balancer_factory", ] diff --git a/vllm_omni/distributed/omni_coordinator/load_balancer.py b/vllm_omni/distributed/omni_coordinator/load_balancer.py index bcd9822ec..9f21d0b97 100644 --- a/vllm_omni/distributed/omni_coordinator/load_balancer.py +++ b/vllm_omni/distributed/omni_coordinator/load_balancer.py @@ -6,7 +6,6 @@ from __future__ import annotations import random import threading from abc import ABC, abstractmethod -from collections.abc import Callable from enum import Enum from typing import Any, TypedDict @@ -122,22 +121,6 @@ class LeastQueueLengthBalancer(LoadBalancer): return random.choice(candidates) -def build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]: - """Translate ``--omni-lb-policy`` (string) into a per-pool LB factory.""" - try: - normalized = LoadBalancingPolicy(policy) - except ValueError as exc: - valid = ", ".join(p.value for p in LoadBalancingPolicy) - raise ValueError(f"unknown --omni-lb-policy {policy!r} (valid: {valid})") from exc - if normalized is LoadBalancingPolicy.RANDOM: - return RandomBalancer - if normalized is LoadBalancingPolicy.ROUND_ROBIN: - return RoundRobinBalancer - if normalized is LoadBalancingPolicy.LEAST_QUEUE_LENGTH: - return LeastQueueLengthBalancer - raise ValueError(f"unhandled load balancing policy {normalized!r}") - - __all__ = [ "Task", "LoadBalancingPolicy", @@ -145,5 +128,4 @@ __all__ = [ "RandomBalancer", "RoundRobinBalancer", "LeastQueueLengthBalancer", - "build_load_balancer_factory", ] diff --git a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py index aee25e584..fb9a9cc7c 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py @@ -230,3 +230,30 @@ class OmniCoordClientForStage: except zmq.ZMQError: pass self._closed = True + + +def create_stage_coord_client( + *, + coord_zmq_addr: str, + input_addr: str, + output_addr: str, + stage_id: int, + queue_length_getter: Callable[[], int] | None = None, +) -> OmniCoordClientForStage: + """Create a stage coordinator client with an optional heartbeat hook.""" + client = OmniCoordClientForStage( + coord_zmq_addr=coord_zmq_addr, + input_addr=input_addr, + output_addr=output_addr, + stage_id=stage_id, + ) + if queue_length_getter is not None: + + def _refresh_queue_length() -> None: + try: + client._queue_length = max(int(queue_length_getter()), 0) + except Exception: + pass + + client._on_heartbeat = _refresh_queue_length + return client diff --git a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py index bbedcfeaa..224204f20 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py @@ -190,6 +190,15 @@ class OmniCoordinator: except zmq.ZMQError: pass + def wait_for_shutdown(self) -> None: + """Block until the coordinator is asked to stop. + + ``OmniCoordinatorRuntime`` runs the coordinator in a child process. + The child needs a small, explicit blocking lifecycle API so it does + not return immediately after construction. + """ + self._stop_event.wait() + def _parse_replica_event(self, data: dict[str, Any]) -> ReplicaEvent | None: """Parse wire payload dict into ReplicaEvent. Returns None if invalid.""" try: diff --git a/vllm_omni/distributed/omni_coordinator/runtime.py b/vllm_omni/distributed/omni_coordinator/runtime.py index 71df7c052..2cc0d2d27 100644 --- a/vllm_omni/distributed/omni_coordinator/runtime.py +++ b/vllm_omni/distributed/omni_coordinator/runtime.py @@ -3,35 +3,91 @@ """Lifecycle wrapper around :class:`OmniCoordinator`. -``OmniCoordinatorRuntime`` is the single-purpose owner of the head-side -coordinator process artifacts: it picks two free TCP ports, constructs an -:class:`OmniCoordinator` bound to them, exposes the resulting addresses, and -provides a single ``close()`` method to tear everything down. +``OmniCoordinatorRuntime`` spawns the coordinator as an independent process +(matching vLLM's DPCoordinator pattern). Physical isolation prevents GIL +contention and makes direct-object coupling impossible. The ROUTER address is later handed to :class:`OmniMasterServer` so it can be published to registering replicas; the PUB address is handed to the -``Orchestrator``, which constructs its :class:`OmniCoordClientForHub` against -it. +``MembershipController``, which constructs its :class:`OmniCoordClientForHub` +against it. """ from __future__ import annotations import logging +import multiprocessing +import multiprocessing.connection +import os +import signal +import weakref +from typing import Any from vllm.utils.network_utils import get_open_ports_list -from .omni_coordinator import OmniCoordinator +from vllm_omni.distributed.omni_coordinator.omni_coordinator import OmniCoordinator logger = logging.getLogger(__name__) -class OmniCoordinatorRuntime: - """Own one :class:`OmniCoordinator` and the two ports it binds. +def run_omni_coordinator_proc( + router_zmq_addr: str, + pub_zmq_addr: str, + heartbeat_timeout: float, + ready_pipe: Any, +) -> None: + """Main loop running inside the coordinator child process.""" + signal.signal(signal.SIGINT, signal.SIG_IGN) - Constructor binds; :meth:`close` tears down. The class deliberately does - not expose the coordinator instance — callers should consume the - coordinator only via its wire protocol through - :class:`OmniCoordClientForStage` and :class:`OmniCoordClientForHub`. + coordinator = OmniCoordinator( + router_zmq_addr=router_zmq_addr, + pub_zmq_addr=pub_zmq_addr, + heartbeat_timeout=heartbeat_timeout, + ) + + ready_pipe.send("ready") + ready_pipe.close() + + coordinator.wait_for_shutdown() + + +def _get_coordinator_mp_context() -> multiprocessing.context.BaseContext: + """Return the multiprocessing context used for OmniCoordinator. + + For the current vllm-omni startup path, ``spawn`` is too expensive: the + child re-imports a heavy CLI / model stack before it can acknowledge the + ready pipe, which can exceed the coordinator startup timeout. Prefer + ``fork`` on platforms that support it. + + ``fork`` must happen before the parent initializes CUDA or owns long-lived + ZMQ sockets; otherwise the child inherits unsafe state. If coordinator + startup moves later, switch this to ``spawn``/``forkserver``. + + TODO: make the coordinator child entry cheap and spawn-safe, then revisit + whether ``spawn`` is still needed here. + """ + if os.name != "nt" and "fork" in multiprocessing.get_all_start_methods(): + return multiprocessing.get_context("fork") + return multiprocessing.get_context("spawn") + + +def _shutdown_proc(proc: multiprocessing.Process) -> None: + """Best-effort process termination for weakref finalizer.""" + if proc.is_alive(): + proc.terminate() + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join(timeout=2) + + +class OmniCoordinatorRuntime: + """Own one :class:`OmniCoordinator` running in a child process. + + Constructor spawns the process; :meth:`close` tears it down. + The class deliberately does not expose the coordinator instance — + callers consume it only via ZMQ through :class:`OmniCoordClientForStage` + and :class:`OmniCoordClientForHub`. """ def __init__( @@ -50,25 +106,54 @@ class OmniCoordinatorRuntime: self.pub_address: str = f"tcp://{host}:{pub_port}" self._closed = False - self._coordinator = OmniCoordinator( - router_zmq_addr=self.router_address, - pub_zmq_addr=self.pub_address, - heartbeat_timeout=heartbeat_timeout, + + ctx = _get_coordinator_mp_context() + parent_conn, child_conn = ctx.Pipe(duplex=False) + + self._proc: multiprocessing.Process = ctx.Process( + target=run_omni_coordinator_proc, + kwargs={ + "router_zmq_addr": self.router_address, + "pub_zmq_addr": self.pub_address, + "heartbeat_timeout": heartbeat_timeout, + "ready_pipe": child_conn, + }, + daemon=True, + name="OmniCoordinator", ) + self._proc.start() + child_conn.close() + + ready = multiprocessing.connection.wait([parent_conn, self._proc.sentinel], timeout=30) + if not ready: + self._proc.terminate() + self._proc.join(timeout=5) + raise RuntimeError("OmniCoordinator process failed to start within 30s") + + try: + status = parent_conn.recv() + except EOFError: + raise RuntimeError("OmniCoordinator process died during startup") from None + finally: + parent_conn.close() + + if status != "ready": + raise RuntimeError(f"OmniCoordinator unexpected status: {status}") + + self._finalizer = weakref.finalize(self, _shutdown_proc, self._proc) logger.info( - "[OmniCoordinatorRuntime] Started (router=%s pub=%s heartbeat_timeout=%.1fs)", + "[OmniCoordinatorRuntime] Started (pid=%d router=%s pub=%s heartbeat_timeout=%.1fs)", + self._proc.pid, self.router_address, self.pub_address, heartbeat_timeout, ) def close(self) -> None: - """Tear down the underlying coordinator. Idempotent.""" + """Tear down the coordinator process. Idempotent.""" if self._closed: return self._closed = True - try: - self._coordinator.close() - except Exception: - logger.exception("[OmniCoordinatorRuntime] coordinator close failed") + _shutdown_proc(self._proc) + self._finalizer.detach() diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 21e3c0e72..dc1385e38 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -9,18 +9,15 @@ from __future__ import annotations import asyncio import concurrent.futures -import copy import dataclasses import json -import os import queue import threading import time import uuid import weakref -from collections.abc import Awaitable, Callable, Mapping, Sequence -from contextlib import ExitStack -from dataclasses import asdict, dataclass +from collections.abc import Mapping, Sequence +from dataclasses import asdict from typing import Any, Literal, cast import janus @@ -35,20 +32,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.inline_stage_diffusion_client import InlineStageDiffusionClient -from vllm_omni.diffusion.io_support import supports_audio_output -from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient -from vllm_omni.diffusion.stage_diffusion_proc import ( - complete_diffusion_handshake, - spawn_diffusion_proc, -) -from vllm_omni.distributed.omni_connectors.utils.initialization import ( - resolve_omni_kv_config_for_stage, -) -from vllm_omni.distributed.omni_coordinator import ( - LoadBalancer, - build_load_balancer_factory, -) +from vllm_omni.diffusion.diffusion_engine import supports_audio_output from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.messages import ( AbortRequestMessage, @@ -57,73 +41,30 @@ from vllm_omni.engine.messages import ( CollectiveRPCResultMessage, EngineQueueMessage, ErrorMessage, - RegisterRemoteReplicaMessage, 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, - spawn_stage_core, +from vllm_omni.engine.stage_init_utils import build_stage0_input_processor +from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.engine.stage_runtime import ( + StageRuntimeInfo, + create_stage_runtime, ) -from vllm_omni.engine.stage_engine_startup import ( - OmniMasterServer, - connect_remote_engine_cores, - launch_omni_core_engines, - register_stage_with_omni_master, -) -from vllm_omni.engine.stage_init_utils import ( - LogicalStageInitPlan, - ReplicaInitPlan, - StageRemoteFactoryContext, - _inject_inferred_kv_tp_topology, - acquire_device_locks, - acquire_diffusion_device_locks, - build_diffusion_config, - build_engine_args_dict, - build_llm_stage_output_processor, - build_stage0_input_processor, - build_vllm_config, - capture_stage_factory_contexts, - compute_replica_layout, - extract_stage_metadata, - get_stage_connector_spec, - initialize_diffusion_stage, - inject_kv_stage_info, - load_omni_transfer_config_for_model, - prepare_engine_environment, - release_device_locks, - stage_runtime_setup, - terminate_alive_proc, -) -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.entrypoints.utils import load_and_resolve_stage_configs from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.prometheus import OmniRequestCounter -from vllm_omni.platforms import current_omni_platform 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 - model_stage: str | None = None +_REQUEST_QUEUE_MAXSIZE = 256 # ============================================================================ @@ -231,7 +172,7 @@ def _weak_shutdown_async_omni_engine( try: if orchestrator_thread is not None and orchestrator_thread.is_alive(): - orchestrator_thread.join(timeout=10) + orchestrator_thread.join() except Exception: pass @@ -308,7 +249,6 @@ class AsyncOmniEngine: ) self._omni_master_address: str | None = kwargs.get("omni_master_address") self._omni_master_port: int | None = kwargs.get("omni_master_port") - self._omni_master_server: OmniMasterServer | None = None # New omni-coordinator flags. Consumed only in single_stage_mode. # ``omni_dp_size_local`` is process-local: each invocation (head and @@ -320,12 +260,6 @@ class AsyncOmniEngine: self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0) if self._omni_heartbeat_timeout <= 0: raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}") - # Coordinator runtime (head-distributed only). - self._coordinator_runtime: Any | None = None - # Per-stage construction context, captured after _initialize_stages - # and used by ``_build_remote_replica`` (the RemoteReplicaFactory - # passed to Orchestrator) when a headless replica registers. - self._stage_remote_factory_contexts: dict[int, StageRemoteFactoryContext] = {} if single_stage_mode: logger.info( @@ -336,7 +270,6 @@ class AsyncOmniEngine: ) self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs) - self._validate_single_stage_mode_replica_constraints() self.num_stages = len(self.stage_configs) stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None @@ -344,6 +277,7 @@ class AsyncOmniEngine: self.stage_pools: list[StagePool] = [] self.stage_clients: list[StageClient] = [] # logical-stage view for external readers self.input_processor: InputProcessor | None = None + self.prompt_expand_func: Any | None = None self.supported_tasks: tuple[str, ...] = ("generate",) self.default_sampling_params_list: list[OmniSamplingParams] = [] self.stage_metadata: list[StageRuntimeInfo] = [] @@ -352,11 +286,10 @@ class AsyncOmniEngine: # sees a non-None ``self.request_queue`` when on_register fires. # ``async_q`` lazily binds to whatever event loop first awaits on # it (the orchestrator loop), so cross-thread use stays correct. - self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue() + self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue(maxsize=_REQUEST_QUEUE_MAXSIZE) self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() self._shutdown_called = False - self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None self._weak_finalizer: weakref.finalize | None = None self._rpc_lock = threading.Lock() self._running_counter = OmniRequestCounter() @@ -390,979 +323,61 @@ class AsyncOmniEngine: logger.info(f"[AsyncOmniEngine] Orchestrator ready with {self.num_stages} stages") - @staticmethod - def _cleanup_launched_llm_resources( - *, - stage_id: int, - proc: Any = None, - engine_manager: Any = None, - coordinator: Any = None, - ) -> None: - """Release launch-only LLM resources when client creation never completed.""" - - if proc is not None: - try: - terminate_alive_proc(proc) - except Exception as cleanup_error: - logger.warning( - "[AsyncOmniEngine] Failed to terminate process for stage %s: %s", - stage_id, - cleanup_error, - ) - - for resource, resource_name in ( - (engine_manager, "engine manager"), - (coordinator, "coordinator"), - ): - if resource is None: - continue - shutdown = getattr(resource, "shutdown", None) - close = getattr(resource, "close", None) - try: - if callable(shutdown): - shutdown() - elif callable(close): - close() - except Exception as cleanup_error: - logger.warning( - "[AsyncOmniEngine] Failed to cleanup launched %s for stage %s: %s", - resource_name, - stage_id, - cleanup_error, - ) - - @staticmethod - def _collect_initialized_clients_for_cleanup( - 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[StageClient] = [] - seen: set[int] = set() - - def _add_client(client: StageClient | None) -> None: - if client is None: - return - client_id = id(client) - if client_id in seen: - return - seen.add(client_id) - collected.append(client) - - for pool in stage_pools: - for client in getattr(pool, "clients", ()): - _add_client(client) - - for clients in initialized_clients_by_stage.values(): - for client in clients: - _add_client(client) - - return collected - - @staticmethod - def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> None: - """Best-effort shutdown for attached clients after init failure.""" - - for client in reversed(list(clients)): - if client is None: - continue - try: - client.shutdown() - except Exception as cleanup_error: - logger.warning( - "[AsyncOmniEngine] Failed to shutdown initialized client after init failure: %s", - cleanup_error, - ) - - def _validate_single_stage_mode_replica_constraints(self) -> None: - """Apply --omni-dp-size-local to the local stage's runtime.num_replicas. - - In the previous revision this method rejected LLM stages with - ``num_replicas > 1``. The whole point of ``--omni-dp-size-local`` is - to lift that restriction for the *local* stage, so the rejection is - gone. We now use this hook to write ``--omni-dp-size-local`` onto - the self-stage's runtime config so downstream code - (``compute_replica_layout`` → ``_build_logical_stage_init_plans``) - sees a consistent view. - """ - if not self.single_stage_mode: - return - target_stage_id = self._single_stage_id_filter - if target_stage_id is None: - return - - for idx, stage_cfg in enumerate(self.stage_configs): - stage_id = int(getattr(stage_cfg, "stage_id", idx)) - runtime_cfg = getattr(stage_cfg, "runtime", None) - if runtime_cfg is None: - continue - if stage_id == target_stage_id: - # Self stage: take --omni-dp-size-local from this process. - try: - runtime_cfg.num_replicas = self._omni_dp_size_local - except Exception: - if hasattr(runtime_cfg, "__setitem__"): - runtime_cfg["num_replicas"] = self._omni_dp_size_local - # Other stages keep their config-declared num_replicas; in - # head-distributed mode they will be launched as ``launch_mode - # == "remote"`` with the configured count. - - def _build_logical_stage_init_plans( - self, - omni_transfer_config: Any, - replicas_per_stage: Sequence[int], - replica_devices_map: Mapping[int, Sequence[str]], - ) -> tuple[list[LogicalStageInitPlan], Any]: - """Build startup plans for every logical stage and replica.""" - - prompt_expand_func = None - stage_plans: list[LogicalStageInitPlan] = [] - - for stage_idx, stage_cfg in enumerate(self.stage_configs): - base_metadata = extract_stage_metadata(stage_cfg) - configured_stage_id = base_metadata.stage_id - if base_metadata.prompt_expand_func is not None: - prompt_expand_func = base_metadata.prompt_expand_func - - stage_connector_spec = get_stage_connector_spec( - omni_transfer_config=omni_transfer_config, - stage_id=configured_stage_id, - async_chunk=self.async_chunk, - ) - omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, configured_stage_id) - num_replicas = replicas_per_stage[stage_idx] - launch_mode = "local" - if ( - self.single_stage_mode - and self._single_stage_id_filter is not None - and configured_stage_id != self._single_stage_id_filter - ): - launch_mode = "remote" - - replicas: list[ReplicaInitPlan] = [] - stage_vllm_config = None - executor_class = None - if base_metadata.stage_type != "diffusion": - engine_args_dict = build_engine_args_dict( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - cli_tokenizer=getattr(self, "tokenizer", None), - ) - omni_conn_cfg, omni_from, omni_to = omni_kv_connector - if omni_conn_cfg: - omni_kv = engine_args_dict.get("omni_kv_config") or {} - if not isinstance(omni_kv, dict): - omni_kv = dict(omni_kv) - omni_kv["connector_config"] = omni_conn_cfg - omni_kv["omni_from_stage"] = omni_from - omni_kv["omni_to_stage"] = omni_to - omni_kv.setdefault("stage_id", configured_stage_id) - engine_args_dict["omni_kv_config"] = omni_kv - if self.stage_configs: - _inject_inferred_kv_tp_topology( - engine_args_dict.get("omni_kv_config"), - configured_stage_id, - self.stage_configs, - ) - stage_vllm_config, executor_class = build_vllm_config( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - engine_args_dict=engine_args_dict, - ) - - for replica_id in range(num_replicas): - replica_cfg = copy.deepcopy(stage_cfg) if replica_id > 0 else stage_cfg - if stage_idx in replica_devices_map: - replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id] - - replica_metadata = extract_stage_metadata(replica_cfg) - replica_metadata.replica_id = replica_id - # In single_stage_mode the head only owns its self stage's - # replicas; the remote-stage metadata exists only so the - # orchestrator can route requests through StagePool. Wiping - # ``runtime_cfg`` there makes sense (we don't manage those - # devices). For the *self stage* we MUST keep the - # per-replica runtime so ``setup_stage_devices`` can apply - # the device split from ``compute_replica_layout`` — - # otherwise every replica inherits the parent's full - # CUDA_VISIBLE_DEVICES and stacks on cuda:0 (OOM with any - # model whose footprint exceeds ~1/(2N) of the card). - if launch_mode == "remote" and replica_metadata.stage_type != "diffusion": - replica_metadata.runtime_cfg = None - - replicas.append( - ReplicaInitPlan( - replica_id=replica_id, - num_replicas=num_replicas, - launch_mode=launch_mode, - stage_cfg=replica_cfg, - metadata=replica_metadata, - stage_connector_spec=stage_connector_spec, - omni_kv_connector=omni_kv_connector, - stage_vllm_config=stage_vllm_config, - executor_class=executor_class, - ) - ) - - stage_plans.append( - LogicalStageInitPlan( - stage_idx=stage_idx, - configured_stage_id=configured_stage_id, - replicas=replicas, - ) - ) - - return stage_plans, prompt_expand_func - - def _start_omni_master_server(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: - """Start OmniMasterServer for single-stage mode.""" - - if not self._omni_master_address or not self._omni_master_port: - raise ValueError( - "AsyncOmniEngine single_stage_mode requires both omni_master_address and omni_master_port to be set." - ) - - all_stage_ids: list[int] = [] - stage_replica_counts: dict[int, int] = {} - # Slots that the head itself will register (launch_mode == "local") - # — auto-assign on the master must not hand these out to remote - # headless registrations even within the race window before the - # head's own register_stage_with_omni_master call completes. - head_local_replicas: dict[int, list[int]] = {} - seen_stage_ids: set[int] = set() - for plan in stage_plans: - stage_id = plan.configured_stage_id - if stage_id in seen_stage_ids: - raise ValueError( - f"Duplicate stage_id {stage_id!r} detected among configured stages; stage_ids must be unique." - ) - seen_stage_ids.add(stage_id) - all_stage_ids.append(stage_id) - stage_replica_counts[stage_id] = len(plan.replicas) - local_rids = [rep.replica_id for rep in plan.replicas if rep.launch_mode == "local"] - if local_rids: - head_local_replicas[stage_id] = local_rids - - # Start the OmniCoordinator runtime first so its router address is - # available to publish in every registration reply. - from vllm_omni.distributed.omni_coordinator import OmniCoordinatorRuntime - - self._coordinator_runtime = OmniCoordinatorRuntime( - host=self._omni_master_address, - heartbeat_timeout=self._omni_heartbeat_timeout, - ) - - self._omni_master_server = OmniMasterServer( - master_address=self._omni_master_address, - master_port=self._omni_master_port, - stage_ids=all_stage_ids, - stage_replica_counts=stage_replica_counts, - coordinator_router_address=self._coordinator_runtime.router_address, - on_register=self._dispatch_master_register, - head_local_replicas=head_local_replicas, - ) - self._omni_master_server.start() - logger.info( - "[AsyncOmniEngine] OmniMasterServer started for stages %s", - all_stage_ids, - ) - - # ------------------------------------------------------------------ - # Remote replica factory (head-side client construction) - # ------------------------------------------------------------------ - - async def _build_remote_replica(self, stage_id: int, replica_id: int) -> Any: - """Construct a head-side stage client for a newly-registered remote replica. - - Used by :class:`Orchestrator` as its ``remote_replica_factory``. - The orchestrator awaits this from its own asyncio loop, so the - client is created in the same loop that owns ZMQ sockets — no - cross-thread setup is required. - - Raises if the stage is not known or if its construction context - was not captured (e.g. the stage was empty at bring-up time). - """ - ctx = self._stage_remote_factory_contexts.get(stage_id) - if ctx is None: - raise RuntimeError( - f"no factory context captured for stage {stage_id}; " - f"known stages: {sorted(self._stage_remote_factory_contexts.keys())}" - ) - if self._omni_master_server is None: - raise RuntimeError("OmniMasterServer is not running; cannot build remote replica") - - alloc = self._omni_master_server.get_allocation(stage_id, replica_id=replica_id) - - # Build a per-replica copy of the base metadata so ``replica_id`` - # is correct (StageMetadata is a plain dataclass-like). - metadata = copy.copy(ctx.base_metadata) - try: - metadata.replica_id = replica_id - except Exception: - # Best-effort: if the metadata object is frozen / unusual, the - # downstream client will fall back to ``replica_id = 0``. - pass - - if ctx.stage_type == "diffusion": - client = StageDiffusionClient.from_addresses( - metadata, - request_address=alloc.input_bind_address, - response_address=alloc.output_bind_address, - batch_size=ctx.diffusion_batch_size, - ) - logger.info( - "[AsyncOmniEngine] Built remote diffusion client for stage=%d replica=%d (req=%s resp=%s)", - stage_id, - replica_id, - alloc.input_bind_address, - alloc.output_bind_address, - ) - return client - - # LLM path - if ctx.vllm_config is None or ctx.executor_class is None: - raise RuntimeError(f"stage {stage_id} factory context is missing vllm_config / executor_class") - - # The headless's StageEngineCoreProc subprocess calls - # vllm.v1.engine.core.startup_handshake at boot and blocks until the - # head's handshake ROUTER answers — without this step it hits the - # built-in 5-minute timeout and exits. The bootstrap (pre-allocated) - # path runs `connect_remote_engine_cores` to perform that - # handshake; dynamic attach must do the same, otherwise every - # replica that comes in via `on_register` (auto-assigned or - # explicit-id-beyond-pre-alloc) deadlocks. Run the blocking - # handshake in a thread so the orchestrator loop stays responsive, - # then build the async client on this loop where `make_async_mp_client` - # expects to be invoked. - master_server = self._omni_master_server - ctx_vllm_config = ctx.vllm_config - ctx_executor_class = ctx.executor_class - - def _run_handshake() -> Any: - with connect_remote_engine_cores( - vllm_config=ctx_vllm_config, - omni_master_server=master_server, - stage_id=stage_id, - replica_id=replica_id, - ) as remote_resources: - _engine_manager, _coordinator, addresses, _ = remote_resources - return _engine_manager, _coordinator, addresses - - engine_manager, coordinator, addresses = await asyncio.to_thread(_run_handshake) - - client_addresses: dict[str, str] = { - "input_address": addresses.inputs[0], - "output_address": addresses.outputs[0], - } - if addresses.frontend_stats_publish_address is not None: - client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address - - client = StageEngineCoreClientBase.make_async_mp_client( - vllm_config=ctx_vllm_config, - executor_class=ctx_executor_class, - metadata=metadata, - client_addresses=client_addresses, - proc=None, - engine_manager=engine_manager, - coordinator=coordinator, - ) - logger.info( - "[AsyncOmniEngine] Built remote LLM client for stage=%d replica=%d (input=%s)", - stage_id, - replica_id, - client_addresses["input_address"], - ) - return client - - # ------------------------------------------------------------------ - # OmniCoordinator on_register proxy - # ------------------------------------------------------------------ - - def _dispatch_master_register(self, stage_id: int, replica_id: int, alloc: Any) -> None: - """Forward a master-server registration to the orchestrator queue. - - Called on :class:`OmniMasterServer`'s ROUTER thread (not the - orchestrator loop). Must return promptly. ``self.request_queue`` - is initialized in :meth:`_initialize_janus_queues` (running inside - the orchestrator thread) before the master server starts, so it - is non-None here; ``janus.Queue.sync_q`` is thread-safe by - construction. - """ - if self.request_queue is None: - logger.warning( - "[AsyncOmniEngine] request_queue not initialized; dropping register_remote_replica stage=%d replica=%d", - stage_id, - replica_id, - ) - return - msg = RegisterRemoteReplicaMessage( - stage_id=int(stage_id), - replica_id=int(replica_id), - ) - try: - self.request_queue.sync_q.put_nowait(msg) - except Exception: - logger.exception( - "[AsyncOmniEngine] Failed to enqueue register_remote_replica for stage=%d replica=%d", - stage_id, - replica_id, - ) - - def _initialize_llm_replica( - self, - plan: ReplicaInitPlan, - stage_init_timeout: int, - llm_stage_launch_lock: threading.Lock, - ) -> StageEngineCoreClientBase: - """Initialize one LLM replica end-to-end.""" - - proc = None - engine_manager = None - coordinator = None - stage_client = None - lock_fds: list[int] = [] - device_control_env = current_omni_platform.device_control_env_var - stage_cfg = plan.stage_cfg - - try: - if plan.launch_mode == "remote": - assert self._omni_master_server is not None - raw_stage_cfg = self._omni_master_server.get_stage_config( - plan.metadata.stage_id, - timeout_s=stage_init_timeout, - replica_id=plan.replica_id, - ) - if raw_stage_cfg is None: - raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config") - vllm_config = plan.stage_vllm_config - executor_class = plan.executor_class - assert vllm_config is not None - assert executor_class is not None - vllm_config.parallel_config.data_parallel_size_local = 0 - launch_cm = connect_remote_engine_cores( - vllm_config=vllm_config, - omni_master_server=self._omni_master_server, - stage_id=plan.metadata.stage_id, - replica_id=plan.replica_id, - ) - logger.info( - "[AsyncOmniEngine] Stage %s remote engine handshake started", - plan.metadata.stage_id, - ) - with launch_cm as remote_resources: - engine_manager, coordinator, addresses, _tensor_queue = remote_resources - - logger.info( - "[AsyncOmniEngine] Stage %s remote engine startup completed", - plan.metadata.stage_id, - ) - client_addresses: dict[str, str] = { - "input_address": addresses.inputs[0], - "output_address": addresses.outputs[0], - } - if addresses.frontend_stats_publish_address is not None: - client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address - stage_client = StageEngineCoreClientBase.make_async_mp_client( - vllm_config=vllm_config, - executor_class=executor_class, - metadata=plan.metadata, - client_addresses=client_addresses, - engine_manager=engine_manager, - coordinator=coordinator, - ) - else: - handshake_address = None - with ExitStack() as launch_stack: - with llm_stage_launch_lock: - previous_visible_devices = os.environ.get(device_control_env) - try: - with stage_runtime_setup( - plan.metadata.stage_id, - plan.metadata.runtime_cfg, - ): - vllm_config = plan.stage_vllm_config - executor_class = plan.executor_class - assert vllm_config is not None - assert executor_class is not None - engine_args_dict = build_engine_args_dict( - stage_cfg, - self.model, - stage_connector_spec=plan.stage_connector_spec, - cli_tokenizer=getattr(self, "tokenizer", None), - ) - lock_fds = acquire_device_locks( - plan.metadata.stage_id, - engine_args_dict, - stage_init_timeout, - ) - if self.single_stage_mode and self._omni_master_server is not None: - coord_router_addr: str | None = ( - self._coordinator_runtime.router_address - if self._coordinator_runtime is not None - else None - ) - engine_manager, coordinator, addresses = launch_stack.enter_context( - launch_omni_core_engines( - vllm_config=vllm_config, - executor_class=executor_class, - log_stats=self._log_stats, - omni_master_server=self._omni_master_server, - stage_id=plan.metadata.stage_id, - stage_config=stage_cfg, - replica_id=plan.replica_id, - omni_coordinator_address=coord_router_addr, - ) - ) - else: - addresses, proc, handshake_address = spawn_stage_core( - vllm_config=vllm_config, - executor_class=executor_class, - log_stats=self._log_stats, - omni_stage_id=plan.metadata.stage_id, - omni_replica_id=plan.replica_id, - ) - logger.info( - "[AsyncOmniEngine] Stage %s engine launch started", - plan.metadata.stage_id, - ) - finally: - if previous_visible_devices is None: - current_omni_platform.unset_device_control_env_var() - else: - current_omni_platform.set_device_control_env_var(previous_visible_devices) - - if self.single_stage_mode and self._omni_master_server is not None: - launch_stack.close() - else: - assert proc is not None - assert handshake_address is not None - complete_stage_handshake(proc, handshake_address, addresses, vllm_config, stage_init_timeout) - logger.info( - "[AsyncOmniEngine] Stage %s engine startup completed", - plan.metadata.stage_id, - ) - - client_addresses: dict[str, str] = { - "input_address": addresses.inputs[0], - "output_address": addresses.outputs[0], - } - if addresses.frontend_stats_publish_address is not None: - client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address - stage_client = StageEngineCoreClientBase.make_async_mp_client( - vllm_config=vllm_config, - executor_class=executor_class, - metadata=plan.metadata, - client_addresses=client_addresses, - proc=proc, - engine_manager=engine_manager, - coordinator=coordinator, - ) - - logger.info("[AsyncOmniEngine] Stage %s initialized", plan.metadata.stage_id) - return stage_client - except Exception: - if stage_client is not None: - try: - stage_client.shutdown() - except Exception as cleanup_error: - logger.warning( - "[AsyncOmniEngine] Failed to cleanup stage %s after attach failure: %s", - plan.metadata.stage_id, - cleanup_error, - ) - else: - self._cleanup_launched_llm_resources( - stage_id=plan.metadata.stage_id, - proc=proc, - engine_manager=engine_manager, - coordinator=coordinator, - ) - raise - finally: - if lock_fds: - release_device_locks(lock_fds) - - def _initialize_diffusion_replica( - self, - plan: ReplicaInitPlan, - stage_init_timeout: int, - stage_launch_lock: threading.Lock, - ) -> StageDiffusionClient | InlineStageDiffusionClient: - """Initialize one diffusion replica end-to-end.""" - - client = None - proc = None - lock_fds: list[int] = [] - try: - if plan.launch_mode == "remote": - assert self._omni_master_server is not None - remote_stage_cfg = OmegaConf.create( - self._omni_master_server.get_stage_config( - plan.metadata.stage_id, - timeout_s=stage_init_timeout, - replica_id=plan.replica_id, - ) - ) - remote_metadata = extract_stage_metadata(remote_stage_cfg) - addresses = self._omni_master_server.get_zmq_addresses( - plan.metadata.stage_id, - replica_id=plan.replica_id, - ) - logger.info( - "[AsyncOmniEngine] Stage %s remote diffusion startup completed", - plan.metadata.stage_id, - ) - client = StageDiffusionClient.from_addresses( - remote_metadata, - request_address=addresses.inputs[0], - response_address=addresses.outputs[0], - batch_size=self.diffusion_batch_size, - ) - else: - device_control_env = current_omni_platform.device_control_env_var - with stage_launch_lock: - previous_visible_devices = os.environ.get(device_control_env) - try: - with stage_runtime_setup( - plan.metadata.stage_id, - plan.metadata.runtime_cfg, - ): - omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector - if omni_conn_cfg: - inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to) - inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self.stage_configs) - if self.single_stage_mode: - assert self._omni_master_server is not None - od_config = build_diffusion_config(self.model, plan.stage_cfg, plan.metadata) - lock_fds = acquire_diffusion_device_locks( - plan.metadata.stage_id, - od_config, - stage_init_timeout, - ) - handshake_address, request_address, response_address = register_stage_with_omni_master( - omni_master_address=self._omni_master_server.address, - omni_master_port=self._omni_master_server.port, - omni_stage_id=plan.metadata.stage_id, - omni_stage_config=plan.stage_cfg, - return_addresses=True, - replica_id=plan.replica_id, - ) - logger.info( - "[AsyncOmniEngine] Stage %s diffusion registration completed", - plan.metadata.stage_id, - ) - coord_router_addr: str | None = ( - self._coordinator_runtime.router_address - if self._coordinator_runtime is not None - else None - ) - proc, _, _, _ = spawn_diffusion_proc( - self.model, - od_config, - handshake_address=handshake_address, - request_address=request_address, - response_address=response_address, - omni_coordinator_address=coord_router_addr, - omni_stage_id=plan.metadata.stage_id, - omni_replica_id=plan.replica_id, - ) - complete_diffusion_handshake(proc, handshake_address, stage_init_timeout) - logger.info( - "[AsyncOmniEngine] Stage %s diffusion startup completed", - plan.metadata.stage_id, - ) - client = StageDiffusionClient.from_addresses( - plan.metadata, - request_address=request_address, - response_address=response_address, - proc=proc, - batch_size=self.diffusion_batch_size, - ) - else: - client = initialize_diffusion_stage( - plan.metadata.stage_id, - self.model, - plan.stage_cfg, - plan.metadata, - stage_init_timeout=stage_init_timeout, - batch_size=self.diffusion_batch_size, - use_inline=self.num_stages == 1 and plan.num_replicas == 1, - ) - finally: - if previous_visible_devices is None: - current_omni_platform.unset_device_control_env_var() - else: - current_omni_platform.set_device_control_env_var(previous_visible_devices) - - logger.info( - "[AsyncOmniEngine] Stage %s replica %s initialized (diffusion, batch_size=%d, devices=%s)", - plan.metadata.stage_id, - plan.replica_id, - self.diffusion_batch_size, - getattr(getattr(plan.stage_cfg, "runtime", None), "devices", "default"), - ) - return client - except Exception: - if proc is not None: - terminate_alive_proc(proc) - raise - finally: - if lock_fds: - release_device_locks(lock_fds) - - def _initialize_replica( - self, - plan: ReplicaInitPlan, - stage_init_timeout: int, - stage_launch_lock: threading.Lock, - ) -> StagePoolClient: - """Initialize one replica, regardless of backend type.""" - - if plan.metadata.stage_type == "diffusion": - return self._initialize_diffusion_replica(plan, stage_init_timeout, stage_launch_lock) - return self._initialize_llm_replica(plan, stage_init_timeout, stage_launch_lock) - - def _initialize_stage_replicas( - self, - stage_plans: Sequence[LogicalStageInitPlan], - stage_init_timeout: int, - ) -> dict[int, list[StagePoolClient | None]]: - """Initialize all stage replicas with parallel startup. - - Diffusion replicas are launched **inline on the orchestrator thread** - (the long-lived daemon thread created in ``__init__``). Their - ``mp.Process`` workers are therefore parented by a thread whose - lifetime equals the engine's lifetime. Submitting diffusion init to a - scoped ``ThreadPoolExecutor`` causes the clone-parent Python thread to - be destroyed at the end of init, which under Ray's actor subreaper - leads the spawned ``DiffusionWorker`` processes to be silently - ``SIGKILL``ed (exitcode -9). - - LLM replicas are submitted to a **long-lived** engine-level executor - (``_stage_init_executor``) whose threads outlive the init phase, - preventing ``prctl(PR_SET_PDEATHSIG, SIGTERM)`` from firing when the - spawning thread exits. - - To achieve parallel startup, LLM replicas are submitted to the - executor first so they begin loading in background, then diffusion - replicas run inline on the orchestrator thread concurrently, and - finally LLM futures are awaited. - """ - - stage_launch_lock = threading.Lock() - 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 - - # Partition replicas: diffusion runs inline on the orchestrator thread; - # LLM replicas are submitted to the engine-level executor. - diffusion_replicas: list[tuple[int, ReplicaInitPlan]] = [] - llm_replicas: list[tuple[int, ReplicaInitPlan]] = [] - for plan in stage_plans: - for replica in plan.replicas: - if replica.metadata.stage_type == "diffusion": - diffusion_replicas.append((plan.stage_idx, replica)) - else: - llm_replicas.append((plan.stage_idx, replica)) - - # Lazily create the engine-level executor (once per engine lifetime). - if llm_replicas and getattr(self, "_stage_init_executor", None) is None: - self._stage_init_executor = concurrent.futures.ThreadPoolExecutor( - max_workers=max(1, len(llm_replicas)), - thread_name_prefix="stage-init", - ) - - # --- 1) Submit LLM replicas to the long-lived executor (starts - # loading in background). --- - future_to_replica: dict[concurrent.futures.Future[StagePoolClient], tuple[int, int]] = {} - for stage_idx, replica in llm_replicas: - future = self._stage_init_executor.submit( - self._initialize_replica, - replica, - stage_init_timeout, - stage_launch_lock, - ) - future_to_replica[future] = (stage_idx, replica.replica_id) - - # --- 2) Diffusion replicas: inline on the orchestrator thread. - # LLM is already loading in background — both run concurrently. --- - if primary_exc is None: - for stage_idx, replica in diffusion_replicas: - try: - initialized_clients_by_stage[stage_idx][replica.replica_id] = self._initialize_replica( - replica, - stage_init_timeout, - stage_launch_lock, - ) - except Exception as exc: - primary_exc = exc - break - - # --- 3) Await LLM futures. --- - if primary_exc is None and future_to_replica: - for future in concurrent.futures.as_completed(future_to_replica): - stage_idx, replica_id = future_to_replica[future] - try: - initialized_clients_by_stage[stage_idx][replica_id] = future.result() - except concurrent.futures.CancelledError: - continue - except Exception as exc: - if primary_exc is None: - primary_exc = exc - for other_future in future_to_replica: - if other_future is future: - continue - other_future.cancel() - - if primary_exc is not None: - setattr(primary_exc, "_initialized_clients_by_stage", initialized_clients_by_stage) - raise primary_exc - - return initialized_clients_by_stage - - def _assemble_stage_pools( - self, - stage_plans: Sequence[LogicalStageInitPlan], - 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[OmniSamplingParams] = [] - stage_metadata_list: list[StageRuntimeInfo] = [] - - for plan in stage_plans: - replica_clients = initialized_clients_by_stage[plan.stage_idx] - first_client = replica_clients[0] if replica_clients else None - if first_client is None: - raise RuntimeError(f"Stage {plan.stage_idx} initialization completed with a missing client") - - 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": - stage_vllm_config = plan.replicas[0].stage_vllm_config - assert stage_vllm_config is not None - output_processor = build_llm_stage_output_processor(plan, stage_vllm_config, log_stats=self._log_stats) - - stage_pools.append( - StagePool( - plan.stage_idx, - clients, - output_processor=output_processor, - stage_vllm_config=stage_vllm_config, - ) - ) - default_sampling_params_list.append(first_client.default_sampling_params) - stage_metadata_list.append( - StageRuntimeInfo( - final_output=first_client.final_output, - final_output_type=first_client.final_output_type, - stage_type=first_client.stage_type, - model_stage=getattr(first_client, "model_stage", None), - ) - ) - - self.default_sampling_params_list = list(default_sampling_params_list) - self.stage_metadata = list(stage_metadata_list) - return stage_pools - def _initialize_stages(self, stage_init_timeout: int) -> None: - """Initialize stage clients/processors in orchestrator thread and assign to self. - - Phases: - 1. Compute replica layout (counts + device splits). - 2. Build per-stage/per-replica startup plans. - 3. Initialize all replicas in parallel via backend-specific launchers. - 4. Build logical StagePools and finalize runtime metadata. - - TODO(stage-pool): move per-stage launch + attach logic into a - StagePool.build_from_config() classmethod so this method only - iterates stage_configs, collects pools, and finalizes metadata. - """ - num_stages = len(self.stage_configs) - self.num_stages = num_stages - self._validate_single_stage_mode_replica_constraints() - - replicas_per_stage, replica_devices_map = compute_replica_layout(self.stage_configs) - - prepare_engine_environment() - omni_transfer_config = load_omni_transfer_config_for_model(self.model, self.config_path) - stage_plans, prompt_expand_func = self._build_logical_stage_init_plans( - omni_transfer_config, - replicas_per_stage, - replica_devices_map, + """Initialize stage clients/processors via StageRuntime and assign to self.""" + self._runtime = create_stage_runtime( + stage_configs=self.stage_configs, + model=self.model, + config_path=self.config_path, + single_stage_mode=self.single_stage_mode, + stage_init_timeout=stage_init_timeout, + diffusion_batch_size=self.diffusion_batch_size, + async_chunk=self.async_chunk, + tokenizer=self.tokenizer, + single_stage_id_filter=self._single_stage_id_filter, + omni_master_address=self._omni_master_address, + omni_master_port=self._omni_master_port, + omni_dp_size_local=self._omni_dp_size_local, + omni_heartbeat_timeout=self._omni_heartbeat_timeout, + omni_lb_policy=self._omni_lb_policy, + request_queue=self.request_queue, ) - # Capture per-stage context now (before _start_omni_master_server) - # so the on_register proxy can build head-side clients for - # registrations that arrive immediately after the server starts. - if self.single_stage_mode: - self._stage_remote_factory_contexts = capture_stage_factory_contexts( - stage_plans, diffusion_batch_size=self.diffusion_batch_size - ) - self._start_omni_master_server(stage_plans) + self._runtime.initialize() - stage_pools: list[StagePool] = [] - input_processor: InputProcessor | None = None - initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = { - plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans - } - - try: - initialized_clients_by_stage = self._initialize_stage_replicas(stage_plans, stage_init_timeout) - if stage_plans and stage_plans[0].replicas[0].metadata.stage_type != "diffusion": - stage0_vllm_config = stage_plans[0].replicas[0].stage_vllm_config - assert stage0_vllm_config is not None - input_processor = build_stage0_input_processor(stage0_vllm_config) - stage_pools = self._assemble_stage_pools(stage_plans, initialized_clients_by_stage) - except Exception as exc: - initialized_clients_by_stage = getattr( - exc, - "_initialized_clients_by_stage", - initialized_clients_by_stage, - ) - cleanup_clients = self._collect_initialized_clients_for_cleanup( - stage_pools, - initialized_clients_by_stage, - ) - logger.exception( - "[AsyncOmniEngine] Stage initialization failed; shutting down %s initialized client(s)", - len(cleanup_clients), - ) - self._shutdown_initialized_clients(cleanup_clients) - if self._omni_master_server is not None: - try: - self._omni_master_server.stop() - except Exception: - logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during stage-init cleanup") - if self._coordinator_runtime is not None: - try: - self._coordinator_runtime.close() - except Exception: - logger.exception( - "[AsyncOmniEngine] Failed to close OmniCoordinatorRuntime during stage-init cleanup" - ) - self._coordinator_runtime = None - raise - - self.stage_pools = stage_pools - self.input_processor = input_processor - self.prompt_expand_func = prompt_expand_func - - # Derive logical-stage views for external readers (entrypoints/async_omni.py). - self.stage_clients = [cast(StageClient, pool.stage_client) for pool in self.stage_pools] + self.num_stages = len(self.stage_configs) + self.stage_pools = self._runtime.stage_pools + self.stage_clients = [ + cast(StageClient, pool.stage_client) for pool in self.stage_pools if pool.stage_client is not None + ] 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] - - # TODO(Peiqi): Hack here + self.input_processor = ( + build_stage0_input_processor(self.stage_vllm_configs[0]) + if self.stage_vllm_configs and self.stage_vllm_configs[0] is not None + else None + ) + self.prompt_expand_func = next( + ( + getattr(client, "prompt_expand_func", None) + for client in self.stage_clients + if getattr(client, "prompt_expand_func", None) is not None + ), + None, + ) + self.default_sampling_params_list = [client.default_sampling_params for client in self.stage_clients] + self.stage_metadata = [ + StageRuntimeInfo( + final_output=client.final_output, + final_output_type=client.final_output_type, + stage_type=client.stage_type, + model_stage=getattr(client, "model_stage", None), + ) + for client in self.stage_clients + ] supported_tasks: set[str] = set() - if any(getattr(pool.stage_client, "is_comprehension", False) for pool in self.stage_pools): + if any(getattr(client, "is_comprehension", False) for client in self.stage_clients): supported_tasks.add("generate") - if any(m.final_output_type == "audio" for m in self.stage_metadata): + if any(meta.final_output_type == "audio" for meta in self.stage_metadata): supported_tasks.add("speech") self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",) @@ -1380,13 +395,9 @@ class AsyncOmniEngine: self._initialize_stages(stage_init_timeout) pd_config = self._detect_pd_config() - coordinator_pub_address: str | None = None - load_balancer_factory: Callable[[], LoadBalancer] | None = None - remote_replica_factory: Callable[[int, int], Awaitable[Any]] | None = None - if self._coordinator_runtime is not None: - coordinator_pub_address = self._coordinator_runtime.pub_address - load_balancer_factory = build_load_balancer_factory(self._omni_lb_policy) - remote_replica_factory = self._build_remote_replica + + membership_controller = self._runtime.create_membership_controller() + orchestrator = Orchestrator( request_async_queue=self.request_queue.async_q, output_async_queue=self.output_queue.async_q, @@ -1394,10 +405,8 @@ class AsyncOmniEngine: stage_pools=self.stage_pools, async_chunk=self.async_chunk, pd_config=pd_config, + membership_controller=membership_controller, running_counter=self._running_counter, - coordinator_pub_address=coordinator_pub_address, - load_balancer_factory=load_balancer_factory, - remote_replica_factory=remote_replica_factory, transfer_emitter=self._transfer_emitter, log_stats=self._log_stats, ) @@ -1776,7 +785,7 @@ class AsyncOmniEngine: # Registration of this companion on stage-0's output processor is # deferred to Orchestrator._handle_add_companion, which routes # admission through StagePool.submit_initial(..., affinity_request_id=...). - self.request_queue.sync_q.put_nowait( + self.request_queue.sync_q.put( AddCompanionRequestMessage( companion_id=cid, parent_id=parent_id, @@ -1921,9 +930,9 @@ class AsyncOmniEngine: pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False - use_hsdp = normalized_kwargs.get("use_hsdp") or False - hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size") or -1 - hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size") or 1 + use_hsdp = normalized_kwargs.get("use_hsdp", False) + hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1) + hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1) if sequence_parallel_size is None: sequence_parallel_size = ulysses_degree * ring_degree @@ -2242,7 +1251,7 @@ class AsyncOmniEngine: reasoning_ended=reasoning_ended, resumable=resumable, ) - self.request_queue.sync_q.put_nowait(msg) + self.request_queue.sync_q.put(msg) # CFG companion expansion: create and enqueue companion requests # so the AR stage also generates their KV caches. @@ -2313,7 +1322,7 @@ class AsyncOmniEngine: resumable=resumable, message_type="streaming_update", ) - self.request_queue.sync_q.put_nowait(msg) + self.request_queue.sync_q.put(msg) async def add_streaming_update_async( self, @@ -2365,7 +1374,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(AbortRequestMessage(request_ids=request_ids)) + self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) async def abort_async(self, request_ids: list[str]) -> None: """Async abort API.""" @@ -2395,7 +1404,7 @@ class AsyncOmniEngine: ) with self._rpc_lock: - self.request_queue.sync_q.put_nowait(msg) + self.request_queue.sync_q.put(msg) deadline = None if timeout is None else time.monotonic() + timeout while True: @@ -2460,20 +1469,10 @@ class AsyncOmniEngine: finalizer.detach() logger.info("[AsyncOmniEngine] Shutting down Orchestrator") - try: - if self.request_queue is not None: - self.request_queue.sync_q.put_nowait(ShutdownRequestMessage()) - except Exception: - pass + if self.request_queue is not None: + self.request_queue.sync_q.put_nowait(ShutdownRequestMessage()) if self.is_alive(): - self.orchestrator_thread.join(timeout=10) - if self.orchestrator_thread.is_alive(): - logger.warning("[AsyncOmniEngine] Orchestrator thread did not exit in time") - - stage_init_executor = getattr(self, "_stage_init_executor", None) - if stage_init_executor is not None: - stage_init_executor.shutdown(wait=False, cancel_futures=True) - self._stage_init_executor = None + self.orchestrator_thread.join() for q in (self.request_queue, self.output_queue, self.rpc_output_queue): try: @@ -2481,19 +1480,11 @@ class AsyncOmniEngine: except Exception: pass - if self._omni_master_server is not None: + if hasattr(self, "_runtime") and self._runtime is not None: try: - self._omni_master_server.stop() + self._runtime.shutdown() except Exception: - logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during shutdown") - self._omni_master_server = None - - if self._coordinator_runtime is not None: - try: - self._coordinator_runtime.close() - except Exception: - logger.exception("[AsyncOmniEngine] Failed to close OmniCoordinatorRuntime during shutdown") - self._coordinator_runtime = None + logger.exception("[AsyncOmniEngine] Failed to shutdown StageRuntime") # ── Release CuMem allocator memory pool ────────────────────────────── # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory diff --git a/vllm_omni/engine/membership_controller.py b/vllm_omni/engine/membership_controller.py new file mode 100644 index 000000000..48b6fe559 --- /dev/null +++ b/vllm_omni/engine/membership_controller.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""MembershipController: distributed replica lifecycle management. + +Extracted from Orchestrator to keep request-flow code free of distributed +concerns. Owns the OmniCoordClientForHub, watches for replica +disappearances, and handles register/unregister by building head-side +clients via an injected factory and mutating StagePool membership. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.distributed.omni_coordinator import ( + LoadBalancer, +) +from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus +from vllm_omni.distributed.omni_coordinator.omni_coord_client_for_hub import OmniCoordClientForHub +from vllm_omni.engine.messages import EngineQueueMessage, ErrorMessage +from vllm_omni.engine.stage_pool import StagePool + +logger = init_logger(__name__) + +RemoteReplicaFactory = Callable[[int, int], Any] + + +class MembershipController: + """Manages dynamic replica attach/detach for distributed mode. + + Constructed by DistStageRuntime and passed to Orchestrator. + The Orchestrator delegates register/unregister messages here. + """ + + WATCH_INTERVAL_S: float = 0.5 + + def __init__( + self, + stage_pools: list[StagePool], + coordinator_pub_address: str, + load_balancer_factory: Callable[[], LoadBalancer], + remote_replica_factory: RemoteReplicaFactory, + ) -> None: + self._stage_pools = stage_pools + self._remote_replica_factory = remote_replica_factory + self._membership_tasks: set[asyncio.Task[None]] = set() + self._shutdown_event = asyncio.Event() + self._watcher_task: asyncio.Task[None] | None = None + self._output_queue: asyncio.Queue[EngineQueueMessage] | None = None + self._cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None + + self._hub = OmniCoordClientForHub(coordinator_pub_address) + factory = load_balancer_factory + for pool in self._stage_pools: + pool.attach_hub(self._hub) + pool.attach_load_balancer(factory()) + + def start(self) -> asyncio.Task[None]: + """Start the replica watcher as a background task. Returns the task.""" + self._watcher_task = asyncio.create_task(self._watch_replica_list(), name="membership-watcher") + return self._watcher_task + + async def handle_register(self, stage_id: int, replica_id: int) -> None: + """Handle a register_remote_replica message (fire-and-forget).""" + self._spawn_task( + self._do_register(stage_id, replica_id), + label=f"register-s{stage_id}-r{replica_id}", + ) + + async def handle_unregister( + self, + stage_id: int, + input_addr: str, + output_queue: asyncio.Queue[EngineQueueMessage] | None = None, + cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None, + ) -> None: + """Handle an unregister_remote_replica message.""" + pool = self._pool_for_stage_id(stage_id) + if pool is None: + return + effective_output_queue = output_queue if output_queue is not None else self._output_queue + effective_cleanup_callback = cleanup_callback if cleanup_callback is not None else self._cleanup_callback + affected = pool.invalidate_addr(input_addr) + self._detach_replica(stage_id, input_addr) + if affected and effective_cleanup_callback is not None: + await effective_cleanup_callback(affected) + if affected and effective_output_queue is not None: + for req_id in affected: + await effective_output_queue.put( + ErrorMessage(error="stage replica disappeared", request_id=req_id, stage_id=stage_id) + ) + + def shutdown(self) -> None: + """Signal stop and close the hub.""" + self._shutdown_event.set() + if self._hub is not None: + self._hub.close() + self._hub = None + if self._watcher_task is not None and not self._watcher_task.done(): + self._watcher_task.cancel() + + async def drain_tasks(self, timeout: float = 10.0) -> None: + """Wait for in-flight membership tasks to complete.""" + if self._membership_tasks: + await asyncio.wait(self._membership_tasks, timeout=timeout) + + # ---- Internal ---- + + def _spawn_task(self, coro: Awaitable[None], *, label: str) -> None: + task = asyncio.create_task(coro, name=f"membership-{label}") + self._membership_tasks.add(task) + + def _on_done(t: asyncio.Task[None]) -> None: + self._membership_tasks.discard(t) + if t.cancelled(): + return + exc = t.exception() + if exc is not None: + logger.error("[MembershipController] %s task crashed", label, exc_info=exc) + + task.add_done_callback(_on_done) + + async def _watch_replica_list(self) -> None: + """Convert hub replica disappearances into unregister actions.""" + last_up: set[tuple[int, str]] = set() + while not self._shutdown_event.is_set(): + try: + snap = self._hub.get_replica_list() + current = {(rep.stage_id, rep.input_addr) for rep in snap.replicas if rep.status == ReplicaStatus.UP} + for stage_id, addr in last_up - current: + self._spawn_task( + self.handle_unregister(stage_id, addr), + label=f"unregister-s{stage_id}", + ) + last_up = current + except asyncio.CancelledError: + raise + except Exception: + logger.exception("[MembershipController] _watch_replica_list iteration failed") + + try: + await asyncio.sleep(self.WATCH_INTERVAL_S) + except asyncio.CancelledError: + raise + + async def _do_register(self, stage_id: int, replica_id: int) -> None: + pool = self._pool_for_stage_id(stage_id) + if pool is None: + logger.warning("[MembershipController] register: stage_id %d out of range", stage_id) + return + client = await asyncio.to_thread(self._remote_replica_factory, stage_id, replica_id) + input_addr = StagePool._client_input_addr(client) + if input_addr is None: + raise RuntimeError(f"remote replica factory for stage {stage_id} produced a client without input address") + pool.add_client(input_addr, client) + logger.info( + "[MembershipController] attached remote replica stage=%d replica=%d addr=%s", + stage_id, + replica_id, + input_addr, + ) + + def _detach_replica(self, stage_id: int, input_addr: str) -> None: + pool = self._pool_for_stage_id(stage_id) + if pool is None: + return + client = pool.remove_client(input_addr) + if client is None: + return + try: + client.shutdown() + except Exception: + logger.exception("[MembershipController] failed to shutdown client stage=%d addr=%s", stage_id, input_addr) + logger.info("[MembershipController] detached replica stage=%d addr=%s", stage_id, input_addr) + + def install_unregister_handlers( + self, + *, + output_queue: asyncio.Queue[EngineQueueMessage], + cleanup_callback: Callable[[list[str]], Awaitable[None]], + ) -> None: + """Install shared cleanup sinks for watcher-driven unregister events.""" + self._output_queue = output_queue + self._cleanup_callback = cleanup_callback + + def _pool_for_stage_id(self, stage_id: int) -> StagePool | None: + if not (0 <= stage_id < len(self._stage_pools)): + return None + return self._stage_pools[stage_id] diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index f8ff84de8..3876414fe 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -5,19 +5,14 @@ Runs inside a background thread with its own asyncio event loop. Owns logical request progression across stage pools and handles stage-to-stage transfer logic. -In distributed mode (``coordinator_pub_address`` provided), it also -owns the single :class:`OmniCoordClientForHub`, runs a -:meth:`_watch_replica_list` task that converts replica disappearances -into ``unregister_remote_replica`` control messages, and handles the -``register_remote_replica`` / ``unregister_remote_replica`` flow that -attaches / detaches head-side stage clients for headless replicas. +Distributed membership (replica attach/detach, hub monitoring) is +handled by :class:`MembershipController`, which is injected optionally. """ from __future__ import annotations import asyncio import time as _time -from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from typing import Any @@ -31,14 +26,9 @@ from vllm.v1.engine import EngineCoreOutputs from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats -from vllm_omni.distributed.omni_coordinator import ( - LoadBalancer, - OmniCoordClientForHub, - RandomBalancer, - ReplicaStatus, -) from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker +from vllm_omni.engine.membership_controller import MembershipController from vllm_omni.engine.messages import ( AbortRequestMessage, AddCompanionRequestMessage, @@ -59,16 +49,6 @@ from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger from vllm_omni.outputs import OmniRequestOutput -# Factory signature for building a head-side stage client for a -# *dynamically attached* (auto-assigned) remote replica. -# -# Receives ``(stage_id, replica_id)`` and returns an awaitable yielding the -# constructed client (any type — it must satisfy the shape expected by the -# matching :class:`StagePool`, i.e. expose ``client_addresses["input_address"]`` -# or ``request_address``, plus the usual ``add_request_async`` / -# ``get_output_async`` / ``shutdown`` surface). -RemoteReplicaFactory = Callable[[int, int], Awaitable[Any]] - logger = init_logger(__name__) @@ -159,10 +139,6 @@ class StreamingInputState: class Orchestrator: """Runs inside a background thread's asyncio event loop.""" - # Cadence at which the replica-list watcher polls for disappearances. - _WATCH_REPLICA_INTERVAL_S: float = 0.5 - _WATCH_REPLICA_IDLE_INTERVAL_S: float = 1.0 - # Class-level defaults so tests that bypass __init__ via object.__new__ # don't AttributeError when transfer / counter emit paths access them. _running_counter: OmniRequestCounter | None = None @@ -178,10 +154,8 @@ class Orchestrator: *, async_chunk: bool = False, pd_config: dict[str, Any] | None = None, + membership_controller: MembershipController | None = None, running_counter: OmniRequestCounter | None = None, - coordinator_pub_address: str | None = None, - load_balancer_factory: Callable[[], LoadBalancer] | None = None, - remote_replica_factory: RemoteReplicaFactory | None = None, transfer_emitter: Any = None, log_stats: bool = False, ) -> None: @@ -212,27 +186,8 @@ class Orchestrator: self._fatal_error: str | None = None self._fatal_error_stage_id: int | None = None - # Background tasks for fire-and-forget message handlers (currently - # only ``register_remote_replica`` and ``unregister_remote_replica``). - # Held as a set so each task's reference survives the loop and the - # task can self-deregister on completion. - self._membership_tasks: set[asyncio.Task[None]] = set() - - # Distributed-mode wiring. The hub is constructed on the - # orchestrator's asyncio loop because it spawns a SUB background - # thread; building it from another thread would race the - # ``_init_done`` event. - self._hub: OmniCoordClientForHub | None = ( - OmniCoordClientForHub(coordinator_pub_address) if coordinator_pub_address is not None else None - ) - self._remote_replica_factory = remote_replica_factory - # Inject hub + per-pool LB into each StagePool so they can run - # distributed dispatch via ``StagePool.pick``. - if self._hub is not None: - factory = load_balancer_factory or RandomBalancer - for pool in self.stage_pools: - pool.attach_hub(self._hub) - pool.attach_load_balancer(factory()) + # Distributed membership (optional, injected by DistStageRuntime) + self._membership = membership_controller def _init_metrics_state( self, @@ -308,16 +263,22 @@ class Orchestrator: self._orchestration_output_handler(), name="orchestrator-stage-output-handler", ) - # The replica watcher only runs in distributed mode. It's still - # created in both cases so ``run()`` has a uniform task graph; - # ``_watch_replica_list`` is a no-op poll when ``self._hub`` is None. - watch_task = asyncio.create_task( - self._watch_replica_list(), - name="orchestrator-replica-watcher", - ) + + # Start membership watcher if distributed mode is active. + membership_watcher: asyncio.Task[None] | None = None + if self._membership is not None: + self._membership.install_unregister_handlers( + output_queue=self.output_async_queue, + cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True), + ) + membership_watcher = self._membership.start() + + tasks = [request_task, output_task] + if membership_watcher is not None: + tasks.append(membership_watcher) try: - await asyncio.gather(request_task, output_task, watch_task) + await asyncio.gather(*tasks) except asyncio.CancelledError: raise except EngineDeadError as e: @@ -332,46 +293,23 @@ class Orchestrator: raise finally: self._shutdown_event.set() - for task in (request_task, output_task, watch_task): + for task in tasks: if not task.done(): task.cancel() try: - await asyncio.gather(request_task, output_task, watch_task, return_exceptions=True) + await asyncio.gather(*tasks, return_exceptions=True) except Exception: pass - # If a fatal error caused the shutdown, drain any pending - # add_request messages that were never processed and broadcast - # fatal error responses so callers are not left hanging. if self._fatal_error is not None: await self._drain_pending_requests_on_fatal() - # Wait briefly for any in-flight membership handlers (register / - # unregister remote replica) to finish so they don't leave the - # head-side pool in a half-attached state. Cancel anything that - # hasn't completed in time; the generic pending-task sweep below - # will collect the cancellations. - if self._membership_tasks: - try: - await asyncio.wait_for( - asyncio.gather(*self._membership_tasks, return_exceptions=True), - timeout=10.0, - ) - except (asyncio.TimeoutError, Exception): - for t in self._membership_tasks: - if not t.done(): - t.cancel() + if self._membership is not None: + await self._membership.drain_tasks(timeout=10.0) + self._membership.shutdown() self._shutdown_stages() - # Close the hub last so any in-flight dispatch still has access. - if self._hub is not None: - try: - self._hub.close() - except RuntimeError: - pass - self._hub = None - loop = asyncio.get_running_loop() pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task() and not t.done()] for task in pending: @@ -379,27 +317,6 @@ class Orchestrator: if pending: await asyncio.gather(*pending, return_exceptions=True) - # ---- Background task helpers ---- - - def _spawn_membership_task(self, coro: Awaitable[None], *, label: str) -> None: - """Run a fire-and-forget membership-change coroutine. - - Holds a strong reference until completion (asyncio would otherwise - garbage-collect a bare task), and logs any uncaught exception. - """ - task = asyncio.create_task(coro, name=f"orchestrator-{label}") - self._membership_tasks.add(task) - - def _on_done(t: asyncio.Task[None]) -> None: - self._membership_tasks.discard(t) - if t.cancelled(): - return - exc = t.exception() - if exc is not None: - logger.error("[Orchestrator] %s task crashed", label, exc_info=exc) - - task.add_done_callback(_on_done) - # ---- Request handling ---- async def _request_handler(self) -> None: @@ -419,18 +336,11 @@ class Orchestrator: elif msg_type == "collective_rpc": await self._handle_collective_rpc(msg) elif isinstance(msg, RegisterRemoteReplicaMessage): - # Dynamic-attach involves a ~5s blocking handshake (run in a - # thread by ``_build_remote_replica``); ``await`` here would - # block the queue and stall the next ``add_request`` until - # the attach finishes. Dispatch as a background task so the - # main message loop keeps draining. - self._spawn_membership_task(self._handle_register_remote_replica(msg), label="register_remote_replica") + if self._membership is not None: + await self._membership.handle_register(msg.stage_id, msg.replica_id) elif isinstance(msg, UnregisterRemoteReplicaMessage): - # Symmetric with register: keep the main queue flowing. - self._spawn_membership_task( - self._handle_unregister_remote_replica(msg), - label="unregister_remote_replica", - ) + if self._membership is not None: + await self._membership.handle_unregister(msg.stage_id, msg.input_addr) elif isinstance(msg, ShutdownRequestMessage): logger.info("[Orchestrator] Received shutdown signal") self._shutdown_event.set() @@ -441,7 +351,8 @@ class Orchestrator: for client in pool.clients: if hasattr(client, "_shutting_down"): client._shutting_down = True - self._shutdown_stages() + # Stage teardown runs once in run()'s finally after the + # orchestration loop observes _shutdown_event and exits. break else: logger.warning("[Orchestrator] Unknown message type: %s", msg_type) @@ -1300,7 +1211,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 = getattr(getattr(next_pool.stage_vllm_config, "model_config", None), "model_stage", None) mm_features = req_state.mm_features if model_stage == "thinker" else None request = build_engine_core_request_from_tokens( request_id=req_id, @@ -1495,129 +1406,6 @@ class Orchestrator: ) self.request_states.pop(req_id, None) - # ---- Distributed-mode replica attach / detach ---- - - async def _watch_replica_list(self) -> None: - """Convert hub replica disappearances into unregister control messages.""" - last_up: set[tuple[int, str]] = set() - while not self._shutdown_event.is_set(): - if self._hub is None: - # No coordinator wired up; sleep coarsely and re-check shutdown. - try: - await asyncio.sleep(self._WATCH_REPLICA_IDLE_INTERVAL_S) - except asyncio.CancelledError: - raise - continue - - try: - snap = self._hub.get_replica_list() - current = {(rep.stage_id, rep.input_addr) for rep in snap.replicas if rep.status == ReplicaStatus.UP} - for stage_id, addr in last_up - current: - await self.request_async_queue.put( - UnregisterRemoteReplicaMessage( - stage_id=stage_id, - input_addr=addr, - ) - ) - last_up = current - except asyncio.CancelledError: - raise - except Exception: - logger.exception("[Orchestrator] _watch_replica_list iteration failed") - - try: - await asyncio.sleep(self._WATCH_REPLICA_INTERVAL_S) - except asyncio.CancelledError: - raise - - async def _handle_register_remote_replica(self, msg: RegisterRemoteReplicaMessage) -> None: - """Bind a head-side client for a newly registered remote replica.""" - stage_id = int(msg.stage_id) - replica_id = int(msg.replica_id) - if not (0 <= stage_id < self.num_stages): - logger.warning( - "[Orchestrator] register_remote_replica: stage_id %d out of range (num_stages=%d)", - stage_id, - self.num_stages, - ) - return - if self._remote_replica_factory is None: - logger.warning( - "[Orchestrator] register_remote_replica received for stage=%d replica=%d but no factory installed", - stage_id, - replica_id, - ) - return - - try: - await self._attach_remote_replica(stage_id, replica_id) - except Exception: - logger.exception( - "[Orchestrator] failed to attach remote replica stage=%d replica=%d", - stage_id, - replica_id, - ) - - async def _handle_unregister_remote_replica(self, msg: UnregisterRemoteReplicaMessage) -> None: - """Tear down the head-side client for a vanished remote replica.""" - stage_id = int(msg.stage_id) - input_addr = str(msg.input_addr) - if not (0 <= stage_id < self.num_stages): - return - pool = self.stage_pools[stage_id] - affected = pool.invalidate_addr(input_addr) - self._detach_remote_replica(stage_id, input_addr) - if affected: - await self._cleanup_request_ids(affected, abort=True) - for req_id in affected: - await self.output_async_queue.put( - ErrorMessage( - error="stage replica disappeared", - request_id=req_id, - stage_id=stage_id, - ) - ) - - async def _attach_remote_replica(self, stage_id: int, replica_id: int) -> None: - """Build a head-side stage client via the injected factory and register it.""" - factory = self._remote_replica_factory - if factory is None: - return - pool = self.stage_pools[stage_id] - client = await factory(stage_id, replica_id) - input_addr = StagePool._client_input_addr(client) - if input_addr is None: - raise RuntimeError( - f"remote replica factory for stage {stage_id} produced a client without a discoverable input address" - ) - pool.add_client(input_addr, client) - logger.info( - "[Orchestrator] attached remote replica stage=%d replica=%d addr=%s", - stage_id, - replica_id, - input_addr, - ) - - def _detach_remote_replica(self, stage_id: int, input_addr: str) -> None: - """Shut down + remove the head-side client at ``input_addr``.""" - pool = self.stage_pools[stage_id] - client = pool.remove_client(input_addr) - if client is None: - return - try: - client.shutdown() - except Exception: - logger.exception( - "[Orchestrator] failed to shutdown client for stage=%d addr=%s", - stage_id, - input_addr, - ) - logger.info( - "[Orchestrator] detached remote replica stage=%d addr=%s", - stage_id, - input_addr, - ) - def _shutdown_stages(self) -> None: """Shutdown all stage pools.""" if self._stages_shutdown: diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 1f10a219c..9236fbc18 100644 --- a/vllm_omni/engine/stage_client.py +++ b/vllm_omni/engine/stage_client.py @@ -26,9 +26,11 @@ class StageClient(Protocol): stage_id: int replica_id: int stage_type: str + model_stage: str | None final_output: bool final_output_type: FinalOutputModalityType | None default_sampling_params: OmniSamplingParams + prompt_expand_func: Any | None requires_multimodal_data: bool custom_process_input_func: Any | None engine_input_source: Sequence[int] @@ -66,8 +68,6 @@ class StagePoolClient(StageClient, Protocol): 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: ... diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 2c6c00665..78c209b0b 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -6,18 +6,12 @@ Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture. from __future__ import annotations -import contextlib import inspect -import multiprocessing.connection import os import socket -import threading -import time -import weakref from typing import TYPE_CHECKING, Any from urllib.parse import urlparse -import psutil import vllm.v1.engine as _vllm_engine_module import vllm.v1.engine.core_client as _vllm_core_client_module from vllm.logger import init_logger @@ -43,8 +37,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -SHUTDOWN_TIMEOUT_S = 5 - def _default_process_engine_inputs( source_outputs: list[Any], @@ -78,10 +70,8 @@ class StageEngineCoreClientBase(StageClientBase): - outputs_queue, output_queue_task - All utility methods (get_output_async, abort_requests_async, etc.) - The subprocess is spawned externally via ``spawn_stage_core`` / - ``complete_stage_handshake`` from *stage_engine_core_proc.py*. - In single-stage CLI mode, the client may instead attach to an - ``engine_manager`` / ``coordinator`` pair created elsewhere. + The stage engine subprocesses are owned by vLLM-style engine managers + (`resources.engine_manager`), matching `MPClient` ownership. """ replica_id: int = 0 @@ -93,7 +83,6 @@ class StageEngineCoreClientBase(StageClientBase): log_stats: bool = False, metadata: StageMetadata | None = None, client_addresses: dict[str, str] | None = None, - proc: Any = None, engine_manager: Any = None, coordinator: Any = None, client_count: int = 1, @@ -107,7 +96,6 @@ class StageEngineCoreClientBase(StageClientBase): log_stats=log_stats, metadata=metadata, client_addresses=client_addresses, - proc=proc, engine_manager=engine_manager, coordinator=coordinator, client_count=client_count, @@ -125,7 +113,6 @@ class StageEngineCoreClientBase(StageClientBase): executor_class: type, log_stats: bool = False, client_addresses: dict[str, str] | None = None, - proc: Any = None, client_count: int = 1, client_index: int = 0, *, @@ -140,10 +127,8 @@ class StageEngineCoreClientBase(StageClientBase): via helpers in stage_init_utils.py. This constructor just stores metadata and calls super().__init__(). - The subprocess is spawned externally via ``spawn_stage_core`` / - ``complete_stage_handshake`` (see *stage_engine_core_proc.py*). - The resulting ``proc`` handle is passed in so this client can - manage the process lifecycle on shutdown. + The subprocess lifecycle is owned by the engine manager attached to + vLLM's background resources, not by this client directly. """ # -------- Stage metadata (public fields used at runtime) -------- self.replica_id = 0 @@ -151,19 +136,17 @@ class StageEngineCoreClientBase(StageClientBase): self.stage_id = metadata.stage_id self.replica_id = getattr(metadata, "replica_id", 0) self.stage_type = metadata.stage_type - self.engine_output_type = metadata.engine_output_type + self.model_stage = metadata.model_stage self.is_comprehension = metadata.is_comprehension self.requires_multimodal_data = metadata.requires_multimodal_data self.engine_input_source = metadata.engine_input_source self.final_output = metadata.final_output self.final_output_type = metadata.final_output_type self.default_sampling_params = metadata.default_sampling_params + self.prompt_expand_func = metadata.prompt_expand_func self.custom_process_input_func = metadata.custom_process_input_func - self.model_stage = metadata.model_stage self.engine_outputs: Any = None - self._proc = proc - self._shutting_down = False self.client_addresses = dict(client_addresses or {}) self._omni_kv_config = getattr(getattr(vllm_config, "model_config", None), "omni_kv_config", None) self._kv_sender_host = self._resolve_contact_host() @@ -203,6 +186,7 @@ class StageEngineCoreClientBase(StageClientBase): ) if engine_manager is not None: self.resources.engine_manager = engine_manager + self.start_engine_core_monitor() if coordinator is not None: self.resources.coordinator = coordinator except Exception: @@ -226,9 +210,6 @@ class StageEngineCoreClientBase(StageClientBase): self._initialize_kv_sender_endpoint() - if self._proc is not None: - self._start_proc_monitor() - logger.info( "[%s] stage-%s [rep-%s] EngineCore running", client_name, @@ -236,64 +217,6 @@ class StageEngineCoreClientBase(StageClientBase): self.replica_id, ) - def _start_proc_monitor(self) -> None: - """Start a daemon thread that watches the subprocess sentinel. - - When the subprocess dies without sending the ZMQ ``ENGINE_CORE_DEAD`` - sentinel (e.g. SIGKILL, segfault, OOM-killer), this thread sets - ``resources.engine_dead`` so subsequent calls raise - ``EngineDeadError``. - """ - proc = self._proc - resources_ref = weakref.ref(self.resources) - stage_id = self.stage_id - replica_id = self.replica_id - - def _monitor() -> None: - try: - multiprocessing.connection.wait([proc.sentinel]) - except Exception: - return - - # Give multiprocessing a brief chance to publish a stable exitcode - # before emitting the monitor log. Without this, the sentinel can - # fire first and the log may misleadingly print ``None``. - with contextlib.suppress(Exception): - proc.join(timeout=0) - exitcode = proc.exitcode - if exitcode is None: - deadline = time.monotonic() + 0.2 - while exitcode is None and time.monotonic() < deadline: - time.sleep(0.01) - with contextlib.suppress(Exception): - proc.join(timeout=0) - exitcode = proc.exitcode - - resources = resources_ref() - if resources is None or resources.engine_dead or self._shutting_down: - return - if exitcode == 0: - logger.info( - "[StageEngineCoreClient] stage-%s [rep-%s] subprocess exited cleanly.", - stage_id, - replica_id, - ) - return - resources.engine_dead = True - logger.error( - "[StageEngineCoreClient] stage-%s [rep-%s] subprocess died unexpectedly (exit code %s).", - stage_id, - replica_id, - exitcode, - ) - - t = threading.Thread( - target=_monitor, - daemon=True, - name=f"StageCoreProcMonitor-{stage_id}", - ) - t.start() - def check_health(self) -> None: """Raise ``EngineDeadError`` if the stage subprocess is dead. @@ -302,9 +225,6 @@ class StageEngineCoreClientBase(StageClientBase): """ if self.resources.engine_dead: raise EngineDeadError(f"Stage-{self.stage_id} engine core is dead") - if self._proc is not None and not self._proc.is_alive(): - self.resources.engine_dead = True - raise EngineDeadError(f"Stage-{self.stage_id} subprocess is not alive (exit code {self._proc.exitcode})") # ==================== Overrides ==================== @@ -335,7 +255,16 @@ class StageEngineCoreClientBase(StageClientBase): return None def _resolve_contact_host(self) -> str | None: - """Resolve a routable host for this stage from its client addresses.""" + """Resolve a routable host for this stage from its client addresses. + + For remote LLM replicas the ZMQ sockets are bound on the head node, + so the ZMQ addresses contain the head's IP. The ``replica_host`` key + (injected by DistStageRuntime) carries the actual replica IP where + the KV connector binds — prefer it when available. + """ + replica_host = self.client_addresses.get("replica_host") + if replica_host: + return replica_host for key in ("input_address", "output_address", "stats_update_address"): address = self.client_addresses.get(key) if not address: @@ -503,41 +432,6 @@ class StageEngineCoreClientBase(StageClientBase): kwargs=kwargs, ) - def shutdown(self, timeout: float | None = None) -> None: - """Shutdown managed resources and any externally spawned subprocess.""" - self._shutting_down = True - child_procs: list[psutil.Process] = [] - if self._proc is not None and self._proc.pid is not None: - try: - child_procs = psutil.Process(self._proc.pid).children(recursive=True) - except psutil.Error: - child_procs = [] - - try: - super().shutdown(timeout=timeout) - finally: - if self._proc is not None and self._proc.is_alive(): - self._proc.terminate() - self._proc.join(timeout=SHUTDOWN_TIMEOUT_S) - if self._proc.is_alive(): - self._proc.kill() - self._proc.join(timeout=SHUTDOWN_TIMEOUT_S) - - alive_children = [proc for proc in child_procs if proc.is_running()] - for proc in alive_children: - try: - proc.terminate() - except psutil.Error: - pass - _, still_alive = psutil.wait_procs(alive_children, timeout=SHUTDOWN_TIMEOUT_S) - for proc in still_alive: - try: - proc.kill() - except psutil.Error: - pass - # The process handle is no longer reliable after best-effort cleanup. - self._proc = None - class StageEngineCoreClient(StageEngineCoreClientBase, AsyncMPClient): """Stage async client backed by vLLM's ``AsyncMPClient``.""" diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 43a8460f7..e3f75d9e3 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -10,40 +10,28 @@ from __future__ import annotations import contextlib import os import signal -from multiprocessing.process import BaseProcess -from typing import TYPE_CHECKING, Any +from typing import Any -import msgspec import vllm.v1.engine.core as _vllm_engine_core_module -import zmq from vllm.logger import init_logger from vllm.transformers_utils.config import ( maybe_register_config_serialize_by_value, ) -from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx from vllm.utils.system_utils import ( decorate_logs, - get_mp_context, set_process_title, ) from vllm.v1.engine import EngineCoreRequestType from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState from vllm.v1.engine.utils import ( - EngineHandshakeMetadata, EngineZmqAddresses, SignalCallback, - get_engine_zmq_addresses, ) -from vllm.v1.utils import shutdown -from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage +from vllm_omni.distributed.omni_coordinator import create_stage_coord_client from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.stage_init_utils import set_death_signal -if TYPE_CHECKING: - from vllm.config import VllmConfig - from vllm.v1.executor import Executor - logger = init_logger(__name__) @@ -106,7 +94,7 @@ class StageEngineCoreProc(EngineCoreProc): ) engine_core: StageEngineCoreProc | None = None - coord_client: OmniCoordClientForStage | None = None + coord_client = None try: # NOTE: previous revisions hardcoded data_parallel_size=1 here # (TODO referencing issue #984). The hardcoding has been removed @@ -154,29 +142,17 @@ class StageEngineCoreProc(EngineCoreProc): "EngineCore handshake did not populate input/output addresses; " "cannot start OmniCoordClientForStage" ) - coord_client = OmniCoordClientForStage( + scheduler = getattr(engine_core, "scheduler", None) + if scheduler is None: + raise RuntimeError("EngineCore scheduler is not initialized") + coord_client = create_stage_coord_client( coord_zmq_addr=omni_coordinator_address, input_addr=addresses.inputs[0], output_addr=addresses.outputs[0], stage_id=int(omni_stage_id), + queue_length_getter=scheduler.get_num_unfinished_requests, ) - def _refresh_queue_length() -> None: - """Pre-heartbeat hook: refresh queue_length from scheduler.""" - scheduler = getattr(engine_core, "scheduler", None) - if scheduler is None: - return - try: - coord_client._queue_length = int( # type: ignore[union-attr] - scheduler.get_num_unfinished_requests() - ) - except Exception: - # Live scheduler stats are best-effort — heartbeats - # must not fail because of a stats lookup error. - pass - - coord_client._on_heartbeat = _refresh_queue_length - def wakeup_engine() -> None: engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None)) @@ -212,113 +188,3 @@ class StageEngineCoreProc(EngineCoreProc): coord_client.close() if engine_core is not None: engine_core.shutdown() - - -def spawn_stage_core( - vllm_config: VllmConfig, - executor_class: type[Executor], - log_stats: bool = False, - omni_stage_id: int | None = None, - omni_replica_id: int = 0, -) -> tuple[EngineZmqAddresses, BaseProcess, str]: - """Spawn a *StageEngineCoreProc* subprocess without performing the handshake. - - Must be called while the correct device env vars are set (e.g. under - the stage-launch lock). Call ``complete_stage_handshake`` afterwards. - - Returns ``(addresses, process, handshake_address)``. - """ - addresses = get_engine_zmq_addresses(vllm_config) - handshake_address = get_open_zmq_ipc_path() - - ctx = get_mp_context() - proc = ctx.Process( - target=StageEngineCoreProc.run_stage_core, - name="StageEngineCoreProc", - kwargs={ - "vllm_config": vllm_config, - "local_client": True, - "handshake_address": handshake_address, - "executor_class": executor_class, - "log_stats": log_stats, - "dp_rank": 0, - "local_dp_rank": 0, - "omni_stage_id": omni_stage_id, - "omni_replica_id": omni_replica_id, - }, - ) - proc.start() - return addresses, proc, handshake_address - - -def complete_stage_handshake( - proc: BaseProcess, - handshake_address: str, - addresses: EngineZmqAddresses, - vllm_config: VllmConfig, - handshake_timeout: int, -) -> None: - """Perform the HELLO/INIT/READY handshake with an already-spawned proc. - - On failure the process is terminated before re-raising. - """ - try: - _perform_handshake(proc, handshake_address, addresses, vllm_config, handshake_timeout) - except Exception: - shutdown([proc]) - raise - - -def _perform_handshake( - proc: BaseProcess, - handshake_address: str, - addresses: EngineZmqAddresses, - vllm_config: VllmConfig, - handshake_timeout: int, -) -> None: - """Run the HELLO / INIT / READY handshake with the subprocess.""" - with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: - poller = zmq.Poller() - poller.register(handshake_socket, zmq.POLLIN) - poller.register(proc.sentinel, zmq.POLLIN) - - identity, msg = _recv(poller, handshake_socket, proc, "HELLO", handshake_timeout) - if msg.get("status") != "HELLO": - raise RuntimeError(f"Expected HELLO, got: {msg}") - - init_payload = EngineHandshakeMetadata( - addresses=addresses, - parallel_config={}, - ) - handshake_socket.send_multipart([identity, msgspec.msgpack.encode(init_payload)]) - - identity, msg = _recv(poller, handshake_socket, proc, "READY", handshake_timeout) - if msg.get("status") != "READY": - raise RuntimeError(f"Expected READY, got: {msg}") - num_gpu_blocks = msg.get("num_gpu_blocks") - if num_gpu_blocks is not None: - vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks - - -def _recv( - poller: zmq.Poller, - handshake_socket: zmq.Socket, - proc: BaseProcess, - expected: str, - timeout_s: int = 600, -) -> tuple[bytes, dict]: - """Wait for one handshake message; raise if the process dies first.""" - timeout_ms = timeout_s * 1000 - while True: - events = dict(poller.poll(timeout=timeout_ms)) - if not events: - raise TimeoutError( - f"Timed out waiting for {expected} from StageEngineCoreProc after {timeout_s}s. " - f"This typically indicates model loading or initialization is taking too long. " - f"Consider increasing `stage_init_timeout` for large models." - ) - if handshake_socket in events: - identity, raw = handshake_socket.recv_multipart() - return identity, msgspec.msgpack.decode(raw) - if proc.exitcode is not None: - raise RuntimeError(f"StageEngineCoreProc died during {expected} (exit code {proc.exitcode})") diff --git a/vllm_omni/engine/omni_core_engine_proc_manager.py b/vllm_omni/engine/stage_engine_core_proc_manager.py similarity index 97% rename from vllm_omni/engine/omni_core_engine_proc_manager.py rename to vllm_omni/engine/stage_engine_core_proc_manager.py index 8df5fd053..454f7496f 100644 --- a/vllm_omni/engine/omni_core_engine_proc_manager.py +++ b/vllm_omni/engine/stage_engine_core_proc_manager.py @@ -47,7 +47,7 @@ except ImportError: # pragma: no cover - depends on vLLM build set_device_control_env_var = None # type: ignore[assignment] -class OmniCoreEngineProcManager(CoreEngineProcManager): +class StageEngineCoreProcManager(CoreEngineProcManager): """Spawn :class:`StageEngineCoreProc` subprocesses with omni kwargs. The body mirrors :class:`CoreEngineProcManager.__init__` because the @@ -102,7 +102,7 @@ class OmniCoreEngineProcManager(CoreEngineProcManager): # Intra-replica vLLM DP mesh (i.e. ``data_parallel_size`` ranks sharing # one engine, one DPCoordinator, one set of weights). Distinct from # the omni-level notion of multiple independent replicas of a stage — - # those each spawn their own OmniCoreEngineProcManager and never join + # those each spawn their own StageEngineCoreProcManager and never join # a vLLM DP group across replicas. has_intra_replica_dp = vllm_config.parallel_config.data_parallel_size > 1 diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index a2744eea4..8f6c199ec 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -4,30 +4,42 @@ from __future__ import annotations import contextlib import dataclasses +import os import socket import threading from collections.abc import Callable, Iterator -from dataclasses import dataclass +from dataclasses import dataclass, field +from multiprocessing import connection +from types import SimpleNamespace from typing import Any import msgspec import zmq from omegaconf import OmegaConf -from vllm.config import CacheConfig, VllmConfig +from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.network_utils import get_open_ports_list, zmq_socket_ctx from vllm.v1.engine.coordinator import DPCoordinator from vllm.v1.engine.utils import ( - STARTUP_POLL_PERIOD_MS, CoreEngine, CoreEngineProcManager, - CoreEngineState, - EngineHandshakeMetadata, EngineZmqAddresses, + get_engine_zmq_addresses, wait_for_engine_startup, ) from vllm.v1.executor import Executor +from vllm_omni.distributed.omni_connectors.utils import initialization +from vllm_omni.engine import stage_init_utils +from vllm_omni.engine.stage_init_utils import ( + acquire_device_locks, + build_diffusion_config, + initialize_diffusion_stage, + release_device_locks, +) +from vllm_omni.entrypoints.utils import inject_omni_kv_config +from vllm_omni.platforms import current_omni_platform + logger = init_logger(__name__) StageRoute = tuple[int, int] @@ -95,6 +107,11 @@ class StageAllocation: # Output channel: client binds PULL, engine connects PUSH output_bind_address: str output_connect_address: str + # The replica's routable IP (from registration). For LLM remote replicas + # the ZMQ sockets are bound on the head, but the KV connector runs on + # the replica host — this field preserves the replica's actual IP so the + # orchestrator can advertise the correct KV sender endpoint. + replica_host: str | None = None @dataclass(frozen=True) @@ -106,6 +123,16 @@ class StageCoordinatorAddresses: frontend_stats_publish_address: str | None = None +@dataclass +class StageReplicaResources: + """Resources created while launching one stage replica.""" + + manager: Any | None = None + coordinator: Any | None = None + addresses: EngineZmqAddresses | None = None + lock_fds: list[int] = field(default_factory=list) + + # --------------------------------------------------------------------------- # OmniMasterServer # --------------------------------------------------------------------------- @@ -142,12 +169,11 @@ class OmniMasterServer: # registrations from multiple headless processes for the same stage # don't race on the routing table. self._alloc_lock = threading.Lock() - self._stage_ids_known: set[int] = set(int(sid) for sid in stage_ids) self._thread: threading.Thread | None = None self._stop_event = threading.Event() stage_replica_counts = dict(stage_replica_counts or {}) - # Slots the *head* itself will fill via ``launch_omni_core_engines`` + # Slots the *head* itself will fill via ``_launch_omni_core_engines`` # / its own ``register_stage_with_omni_master`` call. Auto-assigning # headless registrations must skip these even when they appear # ``_stage_configs``-unfilled — otherwise a fast headless on the same @@ -187,11 +213,6 @@ class OmniMasterServer: """Return the registration port exposed to stage launchers.""" return self._port - @property - def coordinator_router_address(self) -> str | None: - """Return the OmniCoordinator ROUTER address echoed to replicas.""" - return self._coordinator_router_address - def get_allocation(self, stage_id: int, replica_id: int = 0) -> StageAllocation: """Return the full address allocation for *stage_id*.""" return self._stage_routes[(stage_id, replica_id)] @@ -237,7 +258,7 @@ class OmniMasterServer: this stage has been filled do we allocate a fresh id. Slots in ``_head_local_slots`` are reserved for the head's own - ``launch_omni_core_engines`` registration. Auto-assign must skip + ``_launch_omni_core_engines`` registration. Auto-assign must skip them even when ``_stage_configs`` shows them unfilled — otherwise a same-host headless that registers before the head's own ``register_stage_with_omni_master`` call would steal slot 0. @@ -314,14 +335,6 @@ class OmniMasterServer: return self._stage_coordinator_addresses[key] - def get_client_addresses(self, stage_id: int, replica_id: int = 0) -> dict[str, str]: - """Return the addresses the client-side sockets should *bind* to.""" - alloc = self.get_allocation(stage_id, replica_id) - return { - "input_address": alloc.input_bind_address, - "output_address": alloc.output_bind_address, - } - def get_zmq_addresses(self, stage_id: int, replica_id: int = 0) -> EngineZmqAddresses: """Return EngineZmqAddresses using the *bind* (client) side addresses.""" alloc = self.get_allocation(stage_id, replica_id) @@ -338,6 +351,13 @@ class OmniMasterServer: outputs=[alloc.output_connect_address], ) + def get_replica_host(self, stage_id: int, replica_id: int = 0) -> str | None: + """Return the replica's routable IP if it registered one.""" + alloc = self._stage_routes.get((stage_id, replica_id)) + if alloc is None: + return None + return alloc.replica_host + # ------------------------------------------------------------------ # Lifecycle # ------------------------------------------------------------------ @@ -473,59 +493,16 @@ class OmniMasterServer: else: alloc = self._stage_routes[(stage_id, replica_id)] - # Cross-host override: when the registering replica advertised - # its own bind address + ports, rewrite the StageAllocation so - # each socket is rooted on the host that actually binds it - # (the master's pre-allocated ports are unreachable from a - # remote replica's host). - # - # Diffusion and LLM stages have different binder ownership: - # - # Diffusion remote replica (StageDiffusionProc): - # - handshake: replica binds -> rewrite to replica IP - # - input : replica binds -> rewrite to replica IP - # - output : replica binds -> rewrite to replica IP - # - # LLM remote replica (CoreClient on head): - # - handshake: head binds (``connect_remote_engine_cores``) - # -> keep on master IP, worker TCP-connects - # - input : head binds (``CoreClient`` ROUTER) - # -> keep on master IP, worker TCP-connects - # - output : head binds (``CoreClient`` PULL — default - # bind=True for PULL in ``make_zmq_socket``) - # -> keep on master IP, worker TCP-connects - # - # The registrant indicates which case via the boolean - # ``replica_binds_sockets`` payload flag. It defaults to - # True (the diffusion / single-host case) so older callers - # still get the previous full-rewrite semantics. For LLM - # remote replicas, the master keeps every address on its - # own host and the remote worker establishes 3 outbound - # TCP connections to the master. + # Cross-host remote replicas connect to head-owned sockets, but + # still advertise their routable host for KV connector endpoints. new_bind_address = msg.get("replica_bind_address") if new_bind_address: - replica_binds_sockets = bool(msg.get("replica_binds_sockets", True)) - if replica_binds_sockets: - hs_port = int(msg["replica_handshake_port"]) - inp_port = int(msg["replica_input_port"]) - out_port = int(msg["replica_output_port"]) - hs_bind_addr = f"tcp://{new_bind_address}:{hs_port}" - inp_bind_addr = f"tcp://{new_bind_address}:{inp_port}" - out_bind_addr = f"tcp://{new_bind_address}:{out_port}" - alloc = StageAllocation( - handshake_bind_address=hs_bind_addr, - handshake_connect_address=hs_bind_addr, - input_bind_address=inp_bind_addr, - input_connect_address=inp_bind_addr, - output_bind_address=out_bind_addr, - output_connect_address=out_bind_addr, - ) - self._stage_routes[(stage_id, replica_id)] = alloc + alloc.replica_host = new_bind_address logger.info( - "[OmniMasterServer] Stage %d replica %d cross-host bind (sockets bound on %s; replica_ip=%s)", + "[OmniMasterServer] Stage %d replica %d registered from host %s " + "(serving sockets remain head-owned)", stage_id, replica_id, - "replica" if replica_binds_sockets else "master", new_bind_address, ) @@ -599,11 +576,9 @@ def _detect_local_bind_address(master_address: str, master_port: int) -> str: Uses a connected UDP socket as a routing-table probe: ``connect()`` on SOCK_DGRAM sends no packets but forces a route lookup, after which ``getsockname()[0]`` exposes the source IP that an outbound packet to - ``(master_address, master_port)`` would carry. For a co-located master - this returns the loopback or eth0 IP (same effect as the legacy - ``self._address`` behaviour); for a remote master it returns the - NIC IP that's actually reachable from the master — which is exactly - the address the headless's per-stage ZMQ sockets must bind on. + ``(master_address, master_port)`` would carry. For a remote master this + returns the NIC IP that's reachable from the master, which is exactly the + address the headless's per-stage ZMQ sockets must bind on. """ s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) try: @@ -613,6 +588,21 @@ def _detect_local_bind_address(master_address: str, master_port: int) -> str: s.close() +def _engine_count(parallel_config: Any) -> int: + data_parallel_size_local = getattr(parallel_config, "data_parallel_size_local", None) + if data_parallel_size_local is not None and data_parallel_size_local > 0: + return int(data_parallel_size_local) + return max(1, int(parallel_config.data_parallel_size)) + + +def _single_diffusion_parallel_config(*, local_client: bool) -> Any: + return SimpleNamespace( + data_parallel_size_local=1 if local_client else 0, + data_parallel_hybrid_lb=False, + data_parallel_external_lb=False, + ) + + def register_stage_with_omni_master( *, omni_master_address: str, @@ -620,21 +610,11 @@ def register_stage_with_omni_master( omni_stage_id: int, omni_stage_config: Any = None, coordinator: DPCoordinator | None = None, - return_addresses: bool = False, replica_id: int | None = 0, - return_full_response: bool = False, replica_bind_address: str | None = None, - replica_binds_sockets: bool = True, -) -> str | tuple[str, str, str] | StageRegistrationResponse: +) -> StageRegistrationResponse: """Register a stage with the omni master server. - Returns the per-stage handshake address by default. When - ``return_addresses`` is true, also returns the stage input/output - addresses allocated by the master. When ``return_full_response`` is - true, returns the full :class:`StageRegistrationResponse` including the - assigned ``replica_id`` and the OmniCoordinator ROUTER address (if - published by the master). - Pass ``replica_id=None`` to request auto-assignment of a free replica id by the master (used by headless launchers). """ @@ -660,30 +640,11 @@ def register_stage_with_omni_master( payload["coordinator_output"] = coordinator_output payload["frontend_stats_publish_address"] = coordinator.get_stats_publish_address() - # Always advertise THIS host's local bind address + 3 locally - # free ports so the master can root the per-stage socket - # allocation on the replica's own interface. For a co-located - # replica the detected IP matches the master's address and - # the override is a no-op semantically; for a cross-host - # replica it's what makes the headless's ROUTER bind succeed - # (otherwise the master would hand back ``tcp://:port`` - # and ``zmq.bind`` would EADDRNOTAVAIL on the remote host). + # Advertise this host for KV connector routing. Serving/control + # sockets stay head-owned; remote workers only connect to them. if replica_bind_address is None: replica_bind_address = _detect_local_bind_address(omni_master_address, omni_master_port) - hs_port, inp_port, out_port = get_open_ports_list(count=3) payload["replica_bind_address"] = replica_bind_address - payload["replica_handshake_port"] = hs_port - payload["replica_input_port"] = inp_port - payload["replica_output_port"] = out_port - # ``False`` only for LLM headless replicas: the head's - # ``connect_remote_engine_cores`` is the binder for the - # handshake ROUTER, and ``CoreClient.__init__`` binds the - # input ROUTER and the output PULL (``make_zmq_socket`` - # defaults bind=True for PULL). The master must keep all - # three addresses on the master's host so the head can - # ``bind`` them; the remote LLM worker TCP-connects across - # hosts on all three. - payload["replica_binds_sockets"] = bool(replica_binds_sockets) reg_sock.send(msgspec.msgpack.encode(payload)) timeout_ms = _DEFAULT_STARTUP_TIMEOUT_S * 1_000 @@ -712,81 +673,13 @@ def register_stage_with_omni_master( finally: reg_ctx.term() - if return_full_response: - return StageRegistrationResponse( - handshake_address=handshake_address, - input_address=input_address, - output_address=output_address, - replica_id=assigned_replica_id, - coordinator_router_address=coord_router_addr, - ) - if return_addresses: - return handshake_address, input_address, output_address - return handshake_address - - -def _wait_for_omni_engine_startup( - handshake_socket: zmq.Socket, - engine_addresses: EngineZmqAddresses, - engines: list[CoreEngine], - cache_config: CacheConfig, -) -> None: - """Wait for omni-managed engines to finish the HELLO/READY handshake.""" - conn_pending = len(engines) - start_pending = 0 - - poller = zmq.Poller() - poller.register(handshake_socket, zmq.POLLIN) - - while conn_pending or start_pending: - events = poller.poll(STARTUP_POLL_PERIOD_MS) - if not events: - logger.debug( - "[omni] Waiting for %d engine(s) to connect, %d to start.", - conn_pending, - start_pending, - ) - continue - - eng_identity, msg_bytes = handshake_socket.recv_multipart() - eng_index = int.from_bytes(eng_identity, "little") - engine = next((e for e in engines if e.identity == eng_identity), None) - if engine is None: - raise RuntimeError(f"[omni] Handshake message from unexpected engine rank: {eng_index}") - - msg = msgspec.msgpack.decode(msg_bytes) - status: str = msg["status"] - - if status == "HELLO" and engine.state == CoreEngineState.NEW: - init_message = msgspec.msgpack.encode( - EngineHandshakeMetadata(addresses=engine_addresses, parallel_config={}) - ) - handshake_socket.send_multipart((eng_identity, init_message), copy=False) - conn_pending -= 1 - start_pending += 1 - engine.state = CoreEngineState.CONNECTED - logger.debug("[omni] HELLO from engine %d", eng_index) - - elif status == "READY" and engine.state == CoreEngineState.CONNECTED: - # Upstream vllm >=0.19 dropped `num_gpu_blocks` from the READY - # handshake payload (the field is now communicated out-of-band via - # stats/health topics); tolerate both legacy and new message - # shapes so the omni handshake keeps working across rebases. - reported_blocks = msg.get("num_gpu_blocks") - if reported_blocks is not None: - cache_config.num_gpu_blocks = (cache_config.num_gpu_blocks or 0) + int(reported_blocks) - if engine_addresses.frontend_stats_publish_address is None: - engine_addresses.frontend_stats_publish_address = msg.get("dp_stats_address") - start_pending -= 1 - engine.state = CoreEngineState.READY - logger.debug( - "[omni] READY from engine %d (num_gpu_blocks=%s)", - eng_index, - "unknown" if reported_blocks is None else reported_blocks, - ) - - else: - raise RuntimeError(f"[omni] Unexpected status '{status}' from engine {eng_index} in state {engine.state}.") + return StageRegistrationResponse( + handshake_address=handshake_address, + input_address=input_address, + output_address=output_address, + replica_id=assigned_replica_id, + coordinator_router_address=coord_router_addr, + ) @contextlib.contextmanager @@ -795,16 +688,11 @@ def connect_remote_engine_cores( omni_master_server: OmniMasterServer, stage_id: int, replica_id: int = 0, -) -> Iterator[tuple[None, DPCoordinator | None, EngineZmqAddresses, None]]: +) -> Iterator[StageReplicaResources]: """Wait for remote engine cores to connect through the omni handshake.""" addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id) parallel_config = vllm_config.parallel_config - # Mirror the engine-count logic from launch_omni_core_engines. - remote_engine_count = ( - parallel_config.data_parallel_size_local - if parallel_config.data_parallel_size_local is not None and parallel_config.data_parallel_size_local > 0 - else max(1, parallel_config.data_parallel_size) - ) + remote_engine_count = _engine_count(parallel_config) start_index = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 coordinator = None @@ -828,18 +716,75 @@ def connect_remote_engine_cores( handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: - yield None, coordinator, addresses, None - - _wait_for_omni_engine_startup( + yield StageReplicaResources( + coordinator=coordinator, + addresses=addresses, + ) + wait_for_engine_startup( handshake_socket, addresses, engines_to_handshake, + vllm_config.parallel_config, + False, # coordinated_dp vllm_config.cache_config, + None, # proc_manager (remote — no local procs) + None, # coord_process ) @contextlib.contextmanager -def launch_omni_core_engines( +def connect_remote_diffusion_proc( + omni_master_server: OmniMasterServer, + stage_id: int, + replica_id: int = 0, +) -> Iterator[StageReplicaResources]: + """Wait for a remote headless diffusion proc to connect to head-owned sockets.""" + addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id) + handshake_bind_address = omni_master_server.get_allocation( + stage_id, + replica_id=replica_id, + ).handshake_bind_address + + logger.info("Waiting for remote diffusion proc for stage %d replica %d", stage_id, replica_id) + with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: + yield StageReplicaResources(addresses=addresses) + wait_for_engine_startup( + handshake_socket, + addresses, + [CoreEngine(index=0, local=False)], + _single_diffusion_parallel_config(local_client=False), + False, + None, + None, + None, + ) + + +@contextlib.contextmanager +def scoped_spawn_device_env( + stage_visible_devices: str | None, + spawn_device_lock: threading.Lock | None, +) -> Iterator[None]: + """Briefly scope device visibility while spawning a stage subprocess.""" + if stage_visible_devices is None or spawn_device_lock is None: + yield + return + + device_control_env = current_omni_platform.device_control_env_var + with spawn_device_lock: + previous_visible_devices = os.environ.get(device_control_env) + try: + current_omni_platform.set_device_control_env_var(stage_visible_devices) + yield + finally: + if previous_visible_devices is None: + current_omni_platform.unset_device_control_env_var() + else: + current_omni_platform.set_device_control_env_var(previous_visible_devices) + + +@contextlib.contextmanager +def _launch_omni_core_engines( vllm_config: VllmConfig, executor_class: type[Executor], log_stats: bool, @@ -849,22 +794,19 @@ def launch_omni_core_engines( replica_id: int = 0, *, omni_coordinator_address: str | None = None, + stage_visible_devices: str | None = None, + spawn_device_lock: threading.Lock | None = None, ) -> Iterator[tuple[CoreEngineProcManager, DPCoordinator | None, EngineZmqAddresses]]: """Launch local engine cores using the omni registration flow. When ``omni_coordinator_address`` is provided, the spawned engine - subprocesses use :class:`OmniCoreEngineProcManager` and each + subprocesses use :class:`StageEngineCoreProcManager` and each instantiates an :class:`OmniCoordClientForStage` after the handshake completes so the head's :class:`OmniCoordinator` knows about them. """ addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id) parallel_config = vllm_config.parallel_config - # Determine the number of local engines and their ranks. - local_engine_count = ( - parallel_config.data_parallel_size_local - if parallel_config.data_parallel_size_local is not None and parallel_config.data_parallel_size_local > 0 - else max(1, parallel_config.data_parallel_size) - ) + local_engine_count = _engine_count(parallel_config) dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 local_start_index = 0 start_index = dp_rank @@ -903,7 +845,7 @@ def launch_omni_core_engines( # Register the stage once and reuse the returned per-stage handshake # address for all local engine-core processes. - handshake_address = register_stage_with_omni_master( + registration = register_stage_with_omni_master( omni_master_address=omni_master_server.address, omni_master_port=omni_master_server.port, omni_stage_id=stage_id, @@ -911,6 +853,7 @@ def launch_omni_core_engines( coordinator=coordinator, replica_id=replica_id, ) + handshake_address = registration.handshake_address # One CoreEngine entry per local engine so wait_for_engine_startup can # track the HELLO/READY handshake for each of them. @@ -923,45 +866,557 @@ def launch_omni_core_engines( if omni_coordinator_address is not None: # Use the omni subclass so each spawned subprocess instantiates # an OmniCoordClientForStage and heartbeats to the coordinator. - from vllm_omni.engine.omni_core_engine_proc_manager import OmniCoreEngineProcManager + from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager - local_engine_manager: CoreEngineProcManager = OmniCoreEngineProcManager( - local_engine_count=local_engine_count, - start_index=start_index, - local_start_index=local_start_index, - vllm_config=vllm_config, - local_client=True, - handshake_address=handshake_address, - executor_class=executor_class, - log_stats=log_stats, - omni_stage_id=stage_id, - omni_coordinator_address=omni_coordinator_address, - omni_replica_base_id=replica_id, - ) + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + local_engine_manager: CoreEngineProcManager = StageEngineCoreProcManager( + local_engine_count=local_engine_count, + start_index=start_index, + local_start_index=local_start_index, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + omni_stage_id=stage_id, + omni_coordinator_address=omni_coordinator_address, + omni_replica_base_id=replica_id, + ) else: - local_engine_manager = CoreEngineProcManager( - local_engine_count=local_engine_count, - start_index=start_index, - local_start_index=local_start_index, - vllm_config=vllm_config, - local_client=True, - handshake_address=handshake_address, - executor_class=executor_class, - log_stats=log_stats, - ) + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + local_engine_manager = CoreEngineProcManager( + local_engine_count=local_engine_count, + start_index=start_index, + local_start_index=local_start_index, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + ) yield local_engine_manager, coordinator, addresses - - # Wait for all local engine-core processes to complete the - # standard HELLO/READY handshake — mirrors launch_core_engines. - coordinated_dp = parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe wait_for_engine_startup( handshake_socket, addresses, engines_to_handshake, parallel_config, - coordinated_dp, + parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe, vllm_config.cache_config, local_engine_manager, coordinator.proc if coordinator else None, ) + + +@contextlib.contextmanager +def launch_stage_replica( + vllm_config: VllmConfig, + executor_class: type[Executor], + log_stats: bool, + stage_id: int, + *, + replica_id: int = 0, + stage_config: Any = None, + omni_master_server: OmniMasterServer | None = None, + omni_coordinator_address: str | None = None, + stage_visible_devices: str | None = None, + spawn_device_lock: threading.Lock | None = None, +) -> Iterator[StageReplicaResources]: + """Launch a local LLM stage replica. + + This is the common entry point for colocated and distributed local LLM + replicas. Distributed launches delegate to ``_launch_omni_core_engines`` so + registration/address ownership stays centralized in ``OmniMasterServer``. + Colocated launches use an IPC handshake without master registration but + keep the same returned resource bundle. + """ + if omni_master_server is not None: + with _launch_omni_core_engines( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=log_stats, + omni_master_server=omni_master_server, + stage_id=stage_id, + stage_config=stage_config, + replica_id=replica_id, + omni_coordinator_address=omni_coordinator_address, + stage_visible_devices=stage_visible_devices, + spawn_device_lock=spawn_device_lock, + ) as resources: + engine_manager, coordinator, addresses = resources + yield StageReplicaResources( + manager=engine_manager, + coordinator=coordinator, + addresses=addresses, + ) + return + + from vllm.utils.network_utils import get_open_zmq_ipc_path + + from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager + + addresses = get_engine_zmq_addresses(vllm_config) + handshake_address = get_open_zmq_ipc_path() + engines_to_handshake = [CoreEngine(index=0, local=True)] + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + engine_manager = StageEngineCoreProcManager( + local_engine_count=1, + start_index=0, + local_start_index=0, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + omni_stage_id=stage_id, + omni_coordinator_address=omni_coordinator_address, + omni_replica_base_id=replica_id, + ) + + with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: + yield StageReplicaResources( + manager=engine_manager, + addresses=addresses, + ) + wait_for_engine_startup( + handshake_socket, + addresses, + engines_to_handshake, + vllm_config.parallel_config, + False, # coordinated_dp + vllm_config.cache_config, + engine_manager, + None, # coordinator_proc + ) + + +def launch_headless_llm_replica( + *, + vllm_config: VllmConfig, + executor_class: type[Executor], + log_stats: bool, + omni_master_address: str, + omni_master_port: int, + stage_id: int, + stage_config: Any, + coordinator: DPCoordinator | None = None, + replica_bind_address: str | None = None, +) -> CoreEngineProcManager: + """Register and launch one headless LLM replica. + + Headless LLM workers are connectors: the head binds the handshake/input/ + output sockets, while this process connects to the addresses allocated by + ``OmniMasterServer``. Keep that registration/manager wiring in the startup + layer so the CLI only handles argument/config normalization. + """ + from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager + + parallel_config = vllm_config.parallel_config + local_engine_count = parallel_config.data_parallel_size_local + if local_engine_count <= 0: + raise ValueError("data_parallel_size_local must be > 0 in headless mode") + dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 + + response = register_stage_with_omni_master( + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_stage_id=stage_id, + omni_stage_config=stage_config, + coordinator=coordinator, + replica_id=None, + replica_bind_address=replica_bind_address, + ) + + manager = StageEngineCoreProcManager( + local_engine_count=local_engine_count, + start_index=dp_rank, + local_start_index=0, + vllm_config=vllm_config, + local_client=False, + handshake_address=response.handshake_address, + executor_class=executor_class, + log_stats=log_stats, + omni_stage_id=stage_id, + omni_coordinator_address=response.coordinator_router_address, + omni_replica_base_id=response.replica_id, + ) + logger.info( + "[Headless] Stage %d replica id=%d up (coord=%s)", + stage_id, + response.replica_id, + response.coordinator_router_address, + ) + return manager + + +def launch_headless_llm_replicas( + *, + vllm_config: VllmConfig, + executor_class: type[Executor], + log_stats: bool, + omni_master_address: str, + omni_master_port: int, + stage_id: int, + stage_config: Any, + omni_dp_size_local: int, + per_replica_devices: list[str | None], + replica_bind_address: str | None = None, +) -> None: + """Launch, monitor, and clean up all local headless LLM replicas.""" + parallel_config = vllm_config.parallel_config + local_engine_count = parallel_config.data_parallel_size_local + if local_engine_count <= 0: + raise ValueError("data_parallel_size_local must be > 0 in headless mode") + + dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 + coordinator = None + if vllm_config.needs_dp_coordinator and dp_rank == 0: + coordinator = DPCoordinator( + parallel_config, + enable_wave_coordination=vllm_config.model_config.is_moe, + ) + logger.info( + "[Headless] Started DP Coordinator process for stage %d (PID: %d)", + stage_id, + coordinator.proc.pid, + ) + + logger.info( + "[Headless] Launching %d omni replica(s) (vLLM dp_size_local=%d each) for stage %d " + "via OmniMasterServer at %s:%d", + omni_dp_size_local, + local_engine_count, + stage_id, + omni_master_address, + omni_master_port, + ) + + try: + + def _launch_one(rep_idx: int) -> Any: + return launch_headless_llm_replica( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=log_stats, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + stage_id=stage_id, + stage_config=stage_config, + coordinator=coordinator, + replica_bind_address=replica_bind_address, + ) + + launch_headless_replica_group( + stage_id=stage_id, + omni_dp_size_local=omni_dp_size_local, + per_replica_devices=per_replica_devices, + launch_one=_launch_one, + ) + finally: + if coordinator is not None: + coordinator.shutdown() + + +def launch_headless_diffusion_replica( + *, + model: str, + od_config: Any, + stage_config: Any, + stage_id: int, + omni_master_address: str, + omni_master_port: int, + replica_bind_address: str | None = None, +) -> Any: + """Register and launch one headless diffusion replica. + + Headless diffusion follows the LLM remote-attach model: the head binds + handshake/input/output sockets, while this backend process only connects. + """ + response = register_stage_with_omni_master( + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_stage_id=stage_id, + omni_stage_config=stage_config, + replica_id=None, + replica_bind_address=replica_bind_address, + ) + from vllm_omni.diffusion import stage_diffusion_proc + + manager = stage_diffusion_proc.StageDiffusionProcManager.launch_headless( + model=model, + od_config=od_config, + handshake_address=response.handshake_address, + addresses=EngineZmqAddresses( + inputs=[response.input_address], + outputs=[response.output_address], + ), + omni_coordinator_address=response.coordinator_router_address, + omni_stage_id=stage_id, + omni_replica_id=response.replica_id, + ) + logger.info( + "[Headless] Diffusion replica id=%d for stage %d is up (coord=%s)", + response.replica_id, + stage_id, + response.coordinator_router_address, + ) + return manager + + +@contextlib.contextmanager +def replica_device_env(stage_id: int, devices: str | None): + """Temporarily scope device visibility for one replica launch.""" + device_control_env = current_omni_platform.device_control_env_var + previous_visible_devices = os.environ.get(device_control_env) + try: + if devices is not None: + stage_init_utils.setup_stage_devices(stage_id, {"devices": devices}) + yield + finally: + if previous_visible_devices is None: + current_omni_platform.unset_device_control_env_var() + else: + current_omni_platform.set_device_control_env_var(previous_visible_devices) + + +def get_headless_replica_devices( + stage_cfg: Any, + stage_id: int, + omni_dp_size_local: int, +) -> list[str | None]: + """Return per-replica device slices for a headless stage.""" + runtime_cfg = getattr(stage_cfg, "runtime", None) + devices_str: str | None = None + if runtime_cfg is not None: + devices_str = ( + runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + ) + if not devices_str: + return [None] * omni_dp_size_local + + devices_per_replica = stage_init_utils.get_stage_devices_per_replica(stage_cfg) + per_replica_devices = stage_init_utils.split_devices_for_replicas( + devices_str, omni_dp_size_local, devices_per_replica, stage_id + ) + logger.info( + "[Headless] Stage %d: %d local replicas, devices_per_replica=%d, per-replica devices: %s", + stage_id, + omni_dp_size_local, + devices_per_replica, + per_replica_devices, + ) + return per_replica_devices + + +def wait_for_manager_liveness(engine_managers: list[Any]) -> None: + """Block until one or more engine managers report process exit.""" + if len(engine_managers) == 1: + engine_managers[0].monitor_engine_liveness() + return + + def _monitor_target(mgr: Any) -> None: + try: + mgr.monitor_engine_liveness() + except Exception: + logger.exception("[Headless] monitor_engine_liveness raised") + + monitor_threads: list[threading.Thread] = [] + for mgr in engine_managers: + t = threading.Thread( + target=_monitor_target, + args=(mgr,), + name=f"omni-replica-monitor-{id(mgr):x}", + ) + t.start() + monitor_threads.append(t) + for t in monitor_threads: + t.join() + + +def wait_for_diffusion_manager_liveness(managers: list[Any]) -> None: + """Block until one diffusion manager exits and surface non-zero exits.""" + sentinel_to_proc = {manager.proc.sentinel: manager.proc for manager in managers} + died = connection.wait(list(sentinel_to_proc.keys())) + first = sentinel_to_proc[died[0]] + logger.info( + "[Headless] Diffusion replica %s exited (code=%s).", + first.name, + first.exitcode, + ) + if first.exitcode not in (None, 0): + raise RuntimeError(f"Diffusion replica {first.name!r} exited with code {first.exitcode}") + + +def launch_headless_replica_group( + *, + stage_id: int, + omni_dp_size_local: int, + per_replica_devices: list[str | None], + launch_one: Callable[[int], Any], + wait_for_replicas: Callable[[list[Any]], None] = wait_for_manager_liveness, +) -> None: + """Launch, monitor, and clean up a group of local headless replicas.""" + managers: list[Any] = [] + try: + for rep_idx in range(omni_dp_size_local): + with replica_device_env(stage_id, per_replica_devices[rep_idx]): + managers.append(launch_one(rep_idx)) + wait_for_replicas(managers) + finally: + logger.info("[Headless] Shutting down stage %d (%d manager(s)).", stage_id, len(managers)) + for manager in managers: + try: + manager.shutdown() + except Exception: + logger.exception("[Headless] manager shutdown failed") + + +def launch_headless_diffusion_replicas( + *, + model: str, + stage_cfg: Any, + stage_configs: list[Any], + stage_id: int, + omni_master_address: str, + omni_master_port: int, + omni_dp_size_local: int, + per_replica_devices: list[str | None], + config_path: str, + replica_bind_address: str | None = None, +) -> None: + """Prepare diffusion config, launch replicas, monitor, and clean up.""" + omni_transfer_config = stage_init_utils.load_omni_transfer_config_for_model(model, config_path) + omni_conn_cfg, omni_from, omni_to = initialization.resolve_omni_kv_config_for_stage( + omni_transfer_config, + stage_id, + ) + + metadata = stage_init_utils.extract_stage_metadata(stage_cfg) + if omni_conn_cfg: + inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) + # Headless single-stage launch must still infer cross-stage TP topology + # from the loaded deploy config so heterogeneous KV routing keys match the + # head process (e.g. from_tp=2, to_tp=1). + stage_init_utils.inject_kv_stage_info(stage_cfg, stage_id, stage_configs) + od_config = stage_init_utils.build_diffusion_config(model, stage_cfg, metadata) + + logger.info( + "[Headless] Launching %d diffusion replica(s) for stage %d via OmniMasterServer at %s:%d", + omni_dp_size_local, + stage_id, + omni_master_address, + omni_master_port, + ) + + def _launch_one(rep_idx: int) -> Any: + # Keep torch.distributed ports away from the ZMQ ephemeral range that + # OmniMasterServer pre-allocates for sibling headless replicas. + if omni_dp_size_local > 1: + od_config.master_port = od_config.settle_port( + 61000 + rep_idx * 100, + port_inc=37, + ) + return launch_headless_diffusion_replica( + model=model, + od_config=od_config, + stage_config=stage_cfg, + stage_id=stage_id, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + replica_bind_address=replica_bind_address, + ) + + launch_headless_replica_group( + stage_id=stage_id, + omni_dp_size_local=omni_dp_size_local, + per_replica_devices=per_replica_devices, + launch_one=_launch_one, + wait_for_replicas=wait_for_diffusion_manager_liveness, + ) + + +def launch_diffusion_stage_replica( + *, + model: str, + stage_config: Any, + metadata: Any, + stage_init_timeout: int, + batch_size: int, + use_inline: bool, + replica_id: int = 0, + omni_master_server: OmniMasterServer | None = None, + omni_coordinator_address: str | None = None, +) -> tuple[Any, StageReplicaResources]: + """Launch a local diffusion stage replica. + + Colocated mode delegates to ``initialize_diffusion_stage``. Distributed + local mode registers with ``OmniMasterServer`` and spawns a + ``StageDiffusionProc`` that heartbeats to ``OmniCoordinator``. + """ + if omni_master_server is None: + client = initialize_diffusion_stage( + metadata.stage_id, + model, + stage_config, + metadata, + stage_init_timeout=stage_init_timeout, + batch_size=batch_size, + use_inline=use_inline, + ) + return client, StageReplicaResources() + + from vllm_omni.diffusion import stage_diffusion_proc + from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient + + od_config = build_diffusion_config(model, stage_config, metadata) + parallel_config = getattr(od_config, "parallel_config", None) + world_size = getattr(parallel_config, "world_size", 1) + try: + world_size = max(1, int(world_size)) + except (TypeError, ValueError): + world_size = 1 + lock_fds = acquire_device_locks( + metadata.stage_id, + {"tensor_parallel_size": world_size}, + stage_init_timeout, + ) + proc_manager = None + try: + registration = register_stage_with_omni_master( + omni_master_address=omni_master_server.address, + omni_master_port=omni_master_server.port, + omni_stage_id=metadata.stage_id, + omni_stage_config=stage_config, + replica_id=replica_id, + ) + proc_manager = stage_diffusion_proc.StageDiffusionProcManager( + model=model, + od_config=od_config, + stage_init_timeout=stage_init_timeout, + handshake_address=registration.handshake_address, + addresses=EngineZmqAddresses( + inputs=[registration.input_address], + outputs=[registration.output_address], + ), + omni_coordinator_address=omni_coordinator_address, + omni_stage_id=metadata.stage_id, + omni_replica_id=replica_id, + ) + client = StageDiffusionClient.from_addresses( + metadata, + request_address=proc_manager.addresses.inputs[0], + response_address=proc_manager.addresses.outputs[0], + proc_manager=proc_manager, + batch_size=batch_size, + ) + return client, StageReplicaResources( + manager=proc_manager, + addresses=proc_manager.addresses, + lock_fds=lock_fds, + ) + except Exception: + if proc_manager is not None: + proc_manager.shutdown() + if lock_fds: + release_device_locks(lock_fds) + raise diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 53267e23b..5c8f91169 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -51,6 +51,7 @@ class ReplicaInitPlan: omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None] stage_vllm_config: Any | None = None executor_class: type | None = None + engine_args_dict: dict[str, Any] | None = None @dataclass @@ -58,65 +59,8 @@ class LogicalStageInitPlan: """Startup plan for one logical stage.""" stage_idx: int - configured_stage_id: int - replicas: list[ReplicaInitPlan] - - -@dataclass -class StageRemoteFactoryContext: - """Per-stage context cached by AsyncOmniEngine for dynamic replica attach. - - Populated once during ``_bootstrap_orchestrator`` from the per-stage - init plans. ``_build_remote_replica`` consumes it to construct the - right head-side stage client when a headless replica registers. - """ - stage_id: int - stage_type: str - stage_cfg: Any - base_metadata: Any - # LLM-only fields: - vllm_config: Any | None = None - executor_class: type | None = None - # Diffusion-only fields: - diffusion_batch_size: int = 1 - - -def capture_stage_factory_contexts( - stage_plans: Sequence[LogicalStageInitPlan], - diffusion_batch_size: int, -) -> dict[int, StageRemoteFactoryContext]: - """Snapshot per-stage construction context for dynamic replica attach. - - Called once after ``_initialize_stages`` finishes. The captured - context holds everything ``_build_remote_replica`` needs to build a - fresh head-side client when a new headless replica registers - (vllm_config / executor_class for LLM, batch_size for diffusion, - plus the base stage metadata). - - Per-replica fields like ``replica_id`` are filled in at build time, - not at capture time. - """ - contexts: dict[int, StageRemoteFactoryContext] = {} - for plan in stage_plans: - if not plan.replicas: - # Stage was declared but has zero replicas locally; we still - # want to be able to attach incoming headless ones, so use - # the stage_cfg-derived context if any replica plan exists. - continue - template = plan.replicas[0] - stage_id = int(plan.configured_stage_id) - stage_type = template.metadata.stage_type or "llm" - contexts[stage_id] = StageRemoteFactoryContext( - stage_id=stage_id, - stage_type=stage_type, - stage_cfg=template.stage_cfg, - base_metadata=template.metadata, - vllm_config=template.stage_vllm_config, - executor_class=template.executor_class, - diffusion_batch_size=diffusion_batch_size, - ) - return contexts + replicas: list[ReplicaInitPlan] def _resolve_model_to_local_path(model: str) -> str: @@ -359,6 +303,26 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[ logger.debug("Failed to inject stage info into omni_kv_config: %s", e) +def inject_omni_kv_connector_config( + engine_args_dict: dict[str, Any], + omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None], + stage_id: int, +) -> None: + """Inject resolved connector config into a stage engine-args dict.""" + omni_conn_cfg, omni_from, omni_to = omni_kv_connector + if not omni_conn_cfg: + return + + omni_kv = engine_args_dict.get("omni_kv_config") or {} + if not isinstance(omni_kv, dict): + omni_kv = dict(omni_kv) + omni_kv["connector_config"] = omni_conn_cfg + omni_kv["omni_from_stage"] = omni_from + omni_kv["omni_to_stage"] = omni_to + omni_kv.setdefault("stage_id", stage_id) + engine_args_dict["omni_kv_config"] = omni_kv + + @dataclass class StageMetadata: """Lightweight stage attributes extracted from stage_config.""" @@ -1061,32 +1025,6 @@ def release_device_locks(lock_fds: list[int]) -> None: pass -def acquire_diffusion_device_locks( - stage_id: int, - od_config: Any, - stage_init_timeout: int, -) -> list[int]: - """Acquire init locks for the GPU set used by a diffusion stage. - - Diffusion stages express their device count via ``OmniDiffusionConfig``'s - ``parallel_config.world_size`` rather than the LLM-style - ``tensor_parallel_size`` knob, so adapt to the shape that - ``acquire_device_locks`` understands. - """ - parallel_config = getattr(od_config, "parallel_config", None) - world_size = getattr(parallel_config, "world_size", 1) - try: - world_size = max(1, int(world_size)) - except (TypeError, ValueError): - world_size = 1 - - return acquire_device_locks( - stage_id, - {"tensor_parallel_size": world_size}, - stage_init_timeout, - ) - - def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> Any: """Load omni transfer config from an explicit path or resolved model config. diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py new file mode 100644 index 000000000..1819e7032 --- /dev/null +++ b/vllm_omni/engine/stage_runtime.py @@ -0,0 +1,1106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Stage Runtime implementations for single-node and distributed omni stages.""" + +from __future__ import annotations + +import concurrent.futures +import copy +import os +import threading +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any + +import janus +from omegaconf import OmegaConf +from vllm.logger import init_logger + +from vllm_omni.distributed.omni_connectors.utils.initialization import ( + resolve_omni_kv_config_for_stage, +) +from vllm_omni.distributed.omni_coordinator import ( + LeastQueueLengthBalancer, + LoadBalancer, + LoadBalancingPolicy, + RandomBalancer, + RoundRobinBalancer, +) +from vllm_omni.engine.messages import ( + EngineQueueMessage, + RegisterRemoteReplicaMessage, +) +from vllm_omni.engine.output_modality import FinalOutputModalityType +from vllm_omni.engine.stage_client import StageClient, StagePoolClient +from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase +from vllm_omni.engine.stage_engine_startup import ( + OmniMasterServer, + StageReplicaResources, + connect_remote_diffusion_proc, + connect_remote_engine_cores, + launch_diffusion_stage_replica, + launch_stage_replica, +) +from vllm_omni.engine.stage_init_utils import ( + LogicalStageInitPlan, + ReplicaInitPlan, + _inject_inferred_kv_tp_topology, + acquire_device_locks, + build_engine_args_dict, + build_llm_stage_output_processor, + build_vllm_config, + compute_replica_layout, + extract_stage_metadata, + get_stage_connector_spec, + inject_kv_stage_info, + inject_omni_kv_connector_config, + load_omni_transfer_config_for_model, + prepare_engine_environment, + release_device_locks, +) +from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.entrypoints.stage_utils import resolve_stage_physical_devices +from vllm_omni.entrypoints.utils import inject_omni_kv_config +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + + +@dataclass(frozen=True, slots=True) +class StageRuntimeInfo: + final_output: bool + final_output_type: FinalOutputModalityType | None + stage_type: str + model_stage: str | None = None + + +@dataclass +class StageRemoteFactoryContext: + """Per-stage context for creating head-side clients for remote replicas.""" + + stage_id: int + stage_type: str + stage_cfg: Any + base_metadata: Any + vllm_config: Any | None = None + executor_class: type | None = None + diffusion_batch_size: int = 1 + + +def _build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]: + try: + normalized = LoadBalancingPolicy(policy) + except ValueError as exc: + valid = ", ".join(p.value for p in LoadBalancingPolicy) + raise ValueError(f"unknown --omni-lb-policy {policy!r} (valid: {valid})") from exc + if normalized is LoadBalancingPolicy.RANDOM: + return RandomBalancer + if normalized is LoadBalancingPolicy.ROUND_ROBIN: + return RoundRobinBalancer + if normalized is LoadBalancingPolicy.LEAST_QUEUE_LENGTH: + return LeastQueueLengthBalancer + raise ValueError(f"unhandled load balancing policy {normalized!r}") + + +# =========================================================================== +# StageRuntime +# =========================================================================== + + +class StageRuntime: + """Stage runtime for single-node (non-distributed) mode. + + No coordinator, no master server, no hub. Launches stage processes + directly and creates StagePool with static clients. + """ + + def __init__( + self, + stage_configs: list[Any], + model: str, + config_path: str, + *, + stage_init_timeout: int, + diffusion_batch_size: int, + async_chunk: bool, + tokenizer: str | None = None, + ) -> None: + self._stage_configs = stage_configs + self._model = model + self._config_path = config_path + self._stage_init_timeout = stage_init_timeout + self._diffusion_batch_size = diffusion_batch_size + self._async_chunk = async_chunk + self._tokenizer = tokenizer + self._num_stages = len(stage_configs) + + # Populated by initialize() + self.stage_pools: list[StagePool] = [] + self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None + self._spawn_device_lock = threading.Lock() + self._init_visible_devices_baseline: str | None = None + + @staticmethod + def _client_addresses_from_zmq(addresses: Any) -> dict[str, str]: + client_addresses = { + "input_address": addresses.inputs[0], + "output_address": addresses.outputs[0], + } + if addresses.frontend_stats_publish_address is not None: + client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address + return client_addresses + + @staticmethod + def _cleanup_launched_resources( + *, + stage_id: int, + resources: StageReplicaResources | None = None, + ) -> None: + """Release launch-only resources when client creation never completed.""" + if resources is None: + return + + for resource, resource_name in ( + (resources.manager, "manager"), + (resources.coordinator, "coordinator"), + ): + if resource is None: + continue + try: + resource.shutdown() + except Exception as cleanup_error: + logger.warning( + "[StageRuntime] Failed to cleanup launched %s for stage %s: %s", + resource_name, + stage_id, + cleanup_error, + ) + + @staticmethod + def _collect_initialized_clients_for_cleanup( + 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[StageClient] = [] + seen: set[int] = set() + + def _add_client(client: StageClient | None) -> None: + if client is None: + return + client_id = id(client) + if client_id in seen: + return + seen.add(client_id) + collected.append(client) + + for pool in stage_pools: + for client in getattr(pool, "clients", ()): + _add_client(client) + + for clients in initialized_clients_by_stage.values(): + for client in clients: + _add_client(client) + + return collected + + @staticmethod + def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> None: + """Best-effort shutdown for attached clients after init failure.""" + for client in reversed(list(clients)): + if client is None: + continue + try: + client.shutdown() + except Exception as cleanup_error: + logger.warning( + "[StageRuntime] Failed to shutdown initialized client after init failure: %s", + cleanup_error, + ) + + def initialize(self) -> None: + """Run the full stage initialization sequence.""" + stage_plans = self._prepare_stage_plans() + initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = { + plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans + } + try: + self._before_initialize_stage_replicas(stage_plans) + initialized_clients = self._initialize_stage_replicas(stage_plans, self._stage_init_timeout) + initialized_clients_by_stage = initialized_clients + self._finalize_initialized_stages(stage_plans, initialized_clients) + except Exception as exc: + initialized_clients_by_stage = getattr( + exc, + "_initialized_clients_by_stage", + initialized_clients_by_stage, + ) + cleanup_clients = self._collect_initialized_clients_for_cleanup( + self.stage_pools, + initialized_clients_by_stage, + ) + logger.exception( + "[StageRuntime] Stage initialization failed; shutting down %s initialized client(s)", + len(cleanup_clients), + ) + self._shutdown_initialized_clients(cleanup_clients) + self._cleanup_after_initialize_failure() + raise exc + + def shutdown(self) -> None: + for pool in self.stage_pools: + for client in pool.clients: + if client is not None and hasattr(client, "shutdown"): + try: + client.shutdown() + except Exception: + logger.warning("[StageRuntime] client shutdown failed", exc_info=True) + if self._stage_init_executor is not None: + self._stage_init_executor.shutdown(wait=True, cancel_futures=True) + self._stage_init_executor = None + + def create_membership_controller(self) -> Any | None: + """Return a distributed membership controller, if this runtime needs one.""" + return None + + def _prepare_stage_plans(self) -> list[LogicalStageInitPlan]: + """Build logical stage plans and cache prompt expansion metadata.""" + replicas_per_stage, replica_devices_map = compute_replica_layout(self._stage_configs) + prepare_engine_environment() + omni_transfer_config = load_omni_transfer_config_for_model(self._model, self._config_path) + stage_plans = self._build_logical_stage_init_plans( + omni_transfer_config, + replicas_per_stage, + replica_devices_map, + ) + return stage_plans + + def _finalize_initialized_stages( + self, + stage_plans: Sequence[LogicalStageInitPlan], + initialized_clients: Mapping[int, Sequence[StagePoolClient | None]], + ) -> None: + """Populate runtime fields after replica initialization succeeds.""" + self.stage_pools = self._assemble_stage_pools(stage_plans, initialized_clients) + + def _before_initialize_stage_replicas(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + """Hook for runtimes that need infrastructure before replica init.""" + return None + + def _cleanup_after_initialize_failure(self) -> None: + """Hook for runtimes that own extra infrastructure during init.""" + return None + + @contextmanager + def _scoped_spawn_device_env(self, physical_devices: str | None) -> Iterator[None]: + """Briefly scope device visibility for spawn-sensitive setup steps.""" + from vllm_omni.engine.stage_engine_startup import scoped_spawn_device_env + + with scoped_spawn_device_env( + physical_devices, + self._spawn_device_lock, + ): + yield + + def _resolve_replica_physical_devices(self, stage_id: int, runtime_cfg: Any) -> str | None: + if runtime_cfg is None: + runtime_cfg = {} + devices = runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + return resolve_stage_physical_devices( + stage_id, + devices, + visible_baseline=self._init_visible_devices_baseline, + ) + + @contextmanager + def _stage_device_scope(self, stage_id: int, runtime_cfg: Any) -> Iterator[None]: + """Temporarily apply the stage device env while launching a replica.""" + physical_devices = self._resolve_replica_physical_devices(stage_id, runtime_cfg) + with self._scoped_spawn_device_env(physical_devices): + yield + + # ---- Internal methods ---- + + def _build_logical_stage_init_plans( + self, + omni_transfer_config: Any, + replicas_per_stage: Sequence[int], + replica_devices_map: Mapping[int, Sequence[str]], + ) -> list[LogicalStageInitPlan]: + """Build startup plans for every logical stage and replica.""" + stage_plans: list[LogicalStageInitPlan] = [] + + for stage_idx, stage_cfg in enumerate(self._stage_configs): + base_metadata = extract_stage_metadata(stage_cfg) + stage_id = int(base_metadata.stage_id) + if stage_id != stage_idx: + raise ValueError( + "stage_id must match its position in stage_configs: " + f"stage_configs[{stage_idx}] has stage_id={stage_id}. " + "Stage ids must be contiguous and zero-based because the " + "orchestrator indexes stage pools by stage_id." + ) + + stage_connector_spec = get_stage_connector_spec( + omni_transfer_config=omni_transfer_config, + stage_id=stage_id, + async_chunk=self._async_chunk, + ) + omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) + num_replicas = replicas_per_stage[stage_idx] + launch_mode = self._get_launch_mode(stage_id) + + replicas: list[ReplicaInitPlan] = [] + stage_vllm_config = None + executor_class = None + engine_args_dict = None + if base_metadata.stage_type != "diffusion": + engine_args_dict = build_engine_args_dict( + stage_cfg, + self._model, + stage_connector_spec=stage_connector_spec, + cli_tokenizer=self._tokenizer, + ) + inject_omni_kv_connector_config( + engine_args_dict, + omni_kv_connector, + stage_id, + ) + _inject_inferred_kv_tp_topology( + engine_args_dict.get("omni_kv_config"), + stage_id, + self._stage_configs, + ) + stage_vllm_config, executor_class = build_vllm_config( + stage_cfg, + self._model, + stage_connector_spec=stage_connector_spec, + engine_args_dict=engine_args_dict, + ) + + for replica_id in range(num_replicas): + replica_cfg = copy.deepcopy(stage_cfg) if replica_id > 0 else stage_cfg + if stage_idx in replica_devices_map: + replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id] + + replica_metadata = extract_stage_metadata(replica_cfg) + replica_metadata.replica_id = replica_id + if launch_mode == "remote" and replica_metadata.stage_type != "diffusion": + replica_metadata.runtime_cfg = None + + replicas.append( + ReplicaInitPlan( + replica_id=replica_id, + num_replicas=num_replicas, + launch_mode=launch_mode, + stage_cfg=replica_cfg, + metadata=replica_metadata, + stage_connector_spec=stage_connector_spec, + omni_kv_connector=omni_kv_connector, + stage_vllm_config=stage_vllm_config, + executor_class=executor_class, + engine_args_dict=copy.deepcopy(engine_args_dict) if engine_args_dict is not None else None, + ) + ) + + stage_plans.append( + LogicalStageInitPlan( + stage_idx=stage_idx, + stage_id=stage_id, + replicas=replicas, + ) + ) + + return stage_plans + + def _get_launch_mode(self, stage_id: int) -> str: + """Determine launch mode for a stage. Overridden by DistStageRuntime.""" + return "local" + + def _initialize_stage_replicas( + self, + stage_plans: Sequence[LogicalStageInitPlan], + stage_init_timeout: int, + ) -> dict[int, list[StagePoolClient | None]]: + """Initialize all stage replicas. + + Stages sharing the same GPU are initialized sequentially to avoid + memory profiling interference. Stages on different GPUs are + initialized in parallel. + """ + 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 + init_state_lock = threading.Lock() + self._init_visible_devices_baseline = os.environ.get(current_omni_platform.device_control_env_var) + + init_groups: dict[str, list[tuple[int, ReplicaInitPlan]]] = {} + for plan in stage_plans: + for replica in plan.replicas: + init_groups.setdefault(self._replica_init_group_key(replica), []).append((plan.stage_idx, replica)) + + def _init_group(group: list[tuple[int, ReplicaInitPlan]]) -> None: + """Initialize replicas in one scheduling group sequentially.""" + nonlocal primary_exc + for stage_idx, replica in group: + with init_state_lock: + if primary_exc is not None: + return + try: + client = self._initialize_replica( + replica, + stage_init_timeout, + ) + except Exception as exc: + with init_state_lock: + if primary_exc is None: + primary_exc = exc + return + with init_state_lock: + initialized_clients_by_stage[stage_idx][replica.replica_id] = client + + inline_keys = [key for key in init_groups if key.startswith("inline:")] + for key in inline_keys: + _init_group(init_groups.pop(key)) + + if primary_exc is None and init_groups: + if len(init_groups) == 1: + _init_group(next(iter(init_groups.values()))) + else: + future_to_group: dict[concurrent.futures.Future[None], str] = {} + if self._stage_init_executor is None: + self._stage_init_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=len(init_groups), + thread_name_prefix="stage-init", + ) + for group_key, group in init_groups.items(): + future_to_group[self._stage_init_executor.submit(_init_group, group)] = group_key + + for future in concurrent.futures.as_completed(future_to_group): + try: + future.result() + except Exception as exc: + with init_state_lock: + if primary_exc is None: + primary_exc = exc + + if primary_exc is not None: + setattr(primary_exc, "_initialized_clients_by_stage", initialized_clients_by_stage) + raise primary_exc + + return initialized_clients_by_stage + + def _replica_init_group_key(self, replica: ReplicaInitPlan) -> str: + """Return the scheduling group used during replica initialization.""" + if replica.launch_mode == "local" and replica.metadata.stage_type == "diffusion": + # Local diffusion process spawning must stay on the orchestrator + # thread. Keep all local diffusion replicas in one sequential group. + return "inline:diffusion" + if replica.launch_mode == "remote": + return f"remote:{replica.metadata.stage_id}:{replica.replica_id}" + + runtime_cfg = replica.metadata.runtime_cfg or {} + devices = runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + return f"device:{devices}" + + def _initialize_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> StagePoolClient: + if plan.launch_mode == "remote": + return self._initialize_remote_replica(plan, stage_init_timeout) + if plan.metadata.stage_type == "diffusion": + return self._initialize_local_diffusion_replica(plan, stage_init_timeout) + return self._initialize_local_llm_replica(plan, stage_init_timeout) + + def _initialize_remote_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> StagePoolClient: + """Initialize a remote replica. Only distributed runtime implements this.""" + raise NotImplementedError("Remote replicas require DistStageRuntime") + + def _initialize_local_llm_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> StageEngineCoreClientBase: + """Initialize one local LLM replica using vLLM's launch/attach pattern.""" + resources: StageReplicaResources | None = None + stage_client = None + lock_fds: list[int] = [] + try: + physical_devices = self._resolve_replica_physical_devices( + plan.metadata.stage_id, + plan.metadata.runtime_cfg, + ) + if physical_devices: + logger.info( + "[stage_init] Stage-%s set runtime devices: %s", + plan.metadata.stage_id, + physical_devices, + ) + vllm_config = plan.stage_vllm_config + executor_class = plan.executor_class + if vllm_config is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing vllm_config") + if executor_class is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing executor_class") + if plan.engine_args_dict is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing engine args") + with self._scoped_spawn_device_env(physical_devices): + lock_fds = acquire_device_locks( + plan.metadata.stage_id, + plan.engine_args_dict, + stage_init_timeout, + ) + with launch_stage_replica( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=False, + stage_id=plan.metadata.stage_id, + replica_id=plan.replica_id, + stage_config=plan.stage_cfg, + omni_master_server=self._get_omni_master_server(), + omni_coordinator_address=self._get_coordinator_address(), + stage_visible_devices=physical_devices, + spawn_device_lock=self._spawn_device_lock, + ) as resources: + pass + + logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id) + if resources is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no resources") + if resources.addresses is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no addresses") + stage_client = StageEngineCoreClientBase.make_async_mp_client( + vllm_config=vllm_config, + executor_class=executor_class, + metadata=plan.metadata, + client_addresses=self._client_addresses_from_zmq(resources.addresses), + engine_manager=resources.manager, + coordinator=resources.coordinator, + ) + + logger.info("[StageRuntime] Stage %s initialized", plan.metadata.stage_id) + return stage_client + except Exception: + if stage_client is not None: + try: + stage_client.shutdown() + except Exception as cleanup_error: + logger.warning( + "[StageRuntime] Failed to cleanup stage %s after init failure: %s", + plan.metadata.stage_id, + cleanup_error, + ) + else: + self._cleanup_launched_resources( + stage_id=plan.metadata.stage_id, + resources=resources, + ) + raise + finally: + if lock_fds: + release_device_locks(lock_fds) + + def _get_coordinator_address(self) -> str | None: + """Return coordinator router address. Overridden by DistStageRuntime.""" + return None + + def _get_omni_master_server(self) -> OmniMasterServer | None: + """Return the master server for local distributed launches, if any.""" + return None + + def _initialize_local_diffusion_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> Any: + """Initialize one local diffusion replica end-to-end.""" + client = None + resources = None + try: + with self._stage_device_scope(plan.metadata.stage_id, plan.metadata.runtime_cfg): + omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector + if omni_conn_cfg: + inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to) + inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self._stage_configs) + client, resources = launch_diffusion_stage_replica( + model=self._model, + stage_config=plan.stage_cfg, + metadata=plan.metadata, + stage_init_timeout=stage_init_timeout, + batch_size=self._diffusion_batch_size, + use_inline=self._num_stages == 1 and plan.num_replicas == 1, + replica_id=plan.replica_id, + omni_master_server=self._get_omni_master_server(), + omni_coordinator_address=self._get_coordinator_address(), + ) + + logger.info( + "[StageRuntime] Stage %s replica %s initialized (diffusion, batch_size=%d)", + plan.metadata.stage_id, + plan.replica_id, + self._diffusion_batch_size, + ) + return client + except Exception: + if client is not None: + try: + client.shutdown() + except Exception as cleanup_error: + logger.warning( + "[StageRuntime] Failed to cleanup stage %s after init failure: %s", + plan.metadata.stage_id, + cleanup_error, + ) + else: + self._cleanup_launched_resources( + stage_id=plan.metadata.stage_id, + resources=resources, + ) + raise + finally: + if resources is not None and resources.lock_fds: + release_device_locks(resources.lock_fds) + + def _assemble_stage_pools( + self, + stage_plans: Sequence[LogicalStageInitPlan], + initialized_clients_by_stage: Mapping[int, Sequence[StagePoolClient | None]], + ) -> list[StagePool]: + """Assemble logical stage pools.""" + stage_pools: list[StagePool] = [] + + for plan in stage_plans: + replica_clients = initialized_clients_by_stage[plan.stage_idx] + first_client = replica_clients[0] if replica_clients else None + if first_client is None: + raise RuntimeError(f"Stage {plan.stage_idx} initialization completed with a missing client") + + 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": + stage_vllm_config = plan.replicas[0].stage_vllm_config + if stage_vllm_config is None: + raise RuntimeError(f"Stage {plan.stage_id} is missing vllm_config") + output_processor = build_llm_stage_output_processor(plan, stage_vllm_config) + + stage_pools.append( + StagePool( + plan.stage_idx, + clients, + output_processor=output_processor, + stage_vllm_config=stage_vllm_config, + ) + ) + + return stage_pools + + +# =========================================================================== +# DistStageRuntime +# =========================================================================== + + +class DistStageRuntime(StageRuntime): + """Stage runtime for distributed (single_stage_mode) deployment. + + Extends StageRuntime with: + - OmniCoordinatorRuntime (independent process) + - OmniMasterServer for replica registration + - Remote replica support + - Dynamic membership via MembershipController + """ + + def __init__( + self, + stage_configs: list[Any], + model: str, + config_path: str, + *, + stage_init_timeout: int, + diffusion_batch_size: int, + async_chunk: bool, + tokenizer: str | None = None, + single_stage_id_filter: int | None, + omni_master_address: str, + omni_master_port: int, + omni_dp_size_local: int = 1, + omni_heartbeat_timeout: float = 30.0, + omni_lb_policy: str = "random", + request_queue: janus.Queue[EngineQueueMessage] | None = None, + ) -> None: + super().__init__( + stage_configs=stage_configs, + model=model, + config_path=config_path, + stage_init_timeout=stage_init_timeout, + diffusion_batch_size=diffusion_batch_size, + async_chunk=async_chunk, + tokenizer=tokenizer, + ) + self._single_stage_id_filter = single_stage_id_filter + self._omni_master_address = omni_master_address + self._omni_master_port = omni_master_port + self._omni_heartbeat_timeout = omni_heartbeat_timeout + self._omni_lb_policy = omni_lb_policy + self._request_queue = request_queue + + self._omni_master_server: OmniMasterServer | None = None + self._coordinator_runtime: Any | None = None + self._omni_dp_size_local = omni_dp_size_local + self._stage_remote_factory_contexts: dict[int, StageRemoteFactoryContext] = {} + + def create_membership_controller(self) -> Any | None: + if self._coordinator_runtime is None: + return None + + from vllm_omni.engine.membership_controller import MembershipController + + return MembershipController( + stage_pools=self.stage_pools, + coordinator_pub_address=self._coordinator_runtime.pub_address, + load_balancer_factory=_build_load_balancer_factory(self._omni_lb_policy), + remote_replica_factory=self._build_remote_replica, + ) + + def _prepare_stage_plans(self) -> list[LogicalStageInitPlan]: + self._validate_single_stage_mode_replica_constraints() + return super()._prepare_stage_plans() + + def _validate_single_stage_mode_replica_constraints(self) -> None: + """Apply --omni-dp-size-local to the local stage's runtime.num_replicas.""" + target_stage_id = self._single_stage_id_filter + if target_stage_id is None: + return + + for idx, stage_cfg in enumerate(self._stage_configs): + stage_id = int(getattr(stage_cfg, "stage_id", idx)) + runtime_cfg = getattr(stage_cfg, "runtime", None) + if runtime_cfg is None: + continue + if stage_id == target_stage_id: + try: + runtime_cfg.num_replicas = self._omni_dp_size_local + except (AttributeError, TypeError): + if hasattr(runtime_cfg, "__setitem__"): + runtime_cfg["num_replicas"] = self._omni_dp_size_local + continue + logger.warning( + "[DistStageRuntime] Failed to apply omni_dp_size_local=%s to stage %s runtime config", + self._omni_dp_size_local, + stage_id, + ) + + def _before_initialize_stage_replicas(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + self._start_omni_master_server(stage_plans) + self._stage_remote_factory_contexts = self._capture_stage_factory_contexts(stage_plans) + + def _cleanup_after_initialize_failure(self) -> None: + self._cleanup_distributed_infra() + + def shutdown(self) -> None: + super().shutdown() + self._cleanup_distributed_infra() + + def _cleanup_distributed_infra(self) -> None: + if self._omni_master_server is not None: + try: + self._omni_master_server.stop() + except Exception: + logger.warning("[DistStageRuntime] master server stop failed", exc_info=True) + self._omni_master_server = None + if self._coordinator_runtime is not None: + try: + self._coordinator_runtime.close() + except Exception: + logger.warning("[DistStageRuntime] coordinator close failed", exc_info=True) + self._coordinator_runtime = None + + # ---- Distributed overrides ---- + + def _get_launch_mode(self, stage_id: int) -> str: + if self._single_stage_id_filter is not None and stage_id != self._single_stage_id_filter: + return "remote" + return "local" + + def _get_coordinator_address(self) -> str | None: + if self._coordinator_runtime is not None: + return self._coordinator_runtime.router_address + return None + + def _get_omni_master_server(self) -> OmniMasterServer | None: + return self._omni_master_server + + def _initialize_remote_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> StagePoolClient: + """Wait for a configured remote replica and create its head-side client.""" + if self._omni_master_server is None: + raise RuntimeError("OmniMasterServer is not running; cannot initialize remote replica") + registered_stage_cfg = self._omni_master_server.get_stage_config( + plan.metadata.stage_id, + timeout_s=stage_init_timeout, + replica_id=plan.replica_id, + ) + if registered_stage_cfg is None: + raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config") + + metadata = ( + extract_stage_metadata(OmegaConf.create(registered_stage_cfg)) + if plan.metadata.stage_type == "diffusion" + else copy.deepcopy(plan.metadata) + ) + metadata.replica_id = plan.replica_id + ctx = StageRemoteFactoryContext( + stage_id=plan.metadata.stage_id, + stage_type=plan.metadata.stage_type, + stage_cfg=plan.stage_cfg, + base_metadata=metadata, + vllm_config=plan.stage_vllm_config, + executor_class=plan.executor_class, + diffusion_batch_size=self._diffusion_batch_size, + ) + return self._create_remote_replica_client(ctx, plan.replica_id) + + # ---- Distributed infrastructure ---- + + def _start_omni_master_server(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + """Start OmniMasterServer and OmniCoordinatorRuntime.""" + if not self._omni_master_address or not self._omni_master_port: + raise ValueError( + "AsyncOmniEngine single_stage_mode requires both omni_master_address and omni_master_port to be set." + ) + + from vllm_omni.distributed.omni_coordinator import OmniCoordinatorRuntime + + all_stage_ids: list[int] = [] + stage_replica_counts: dict[int, int] = {} + head_local_replicas: dict[int, list[int]] = {} + seen_stage_ids: set[int] = set() + for plan in stage_plans: + stage_id = plan.stage_id + if stage_id in seen_stage_ids: + raise ValueError(f"Duplicate stage_id {stage_id!r} detected") + seen_stage_ids.add(stage_id) + all_stage_ids.append(stage_id) + stage_replica_counts[stage_id] = len(plan.replicas) + local_rids = [rep.replica_id for rep in plan.replicas if rep.launch_mode == "local"] + if local_rids: + head_local_replicas[stage_id] = local_rids + + self._coordinator_runtime = OmniCoordinatorRuntime( + host=self._omni_master_address, + heartbeat_timeout=self._omni_heartbeat_timeout, + ) + + self._omni_master_server = OmniMasterServer( + master_address=self._omni_master_address, + master_port=self._omni_master_port, + stage_ids=all_stage_ids, + stage_replica_counts=stage_replica_counts, + coordinator_router_address=self._coordinator_runtime.router_address, + on_register=self._dispatch_master_register, + head_local_replicas=head_local_replicas, + ) + self._omni_master_server.start() + logger.info("[DistStageRuntime] OmniMasterServer started for stages %s", all_stage_ids) + + def _capture_stage_factory_contexts( + self, stage_plans: Sequence[LogicalStageInitPlan] + ) -> dict[int, StageRemoteFactoryContext]: + contexts: dict[int, StageRemoteFactoryContext] = {} + for plan in stage_plans: + if not plan.replicas: + continue + template = plan.replicas[0] + stage_id = int(plan.stage_id) + contexts[stage_id] = StageRemoteFactoryContext( + stage_id=stage_id, + stage_type=template.metadata.stage_type or "llm", + stage_cfg=template.stage_cfg, + base_metadata=template.metadata, + vllm_config=template.stage_vllm_config, + executor_class=template.executor_class, + diffusion_batch_size=self._diffusion_batch_size, + ) + return contexts + + def _dispatch_master_register(self, stage_id: int, replica_id: int, alloc: Any) -> None: + """Callback from OmniMasterServer when a headless replica registers.""" + if self._request_queue is None: + logger.warning("[DistStageRuntime] on_register fired but no request_queue wired") + return + try: + self._request_queue.sync_q.put_nowait( + RegisterRemoteReplicaMessage(stage_id=stage_id, replica_id=replica_id) + ) + except Exception: + logger.exception("[DistStageRuntime] Failed to enqueue register message") + + def _build_remote_replica(self, stage_id: int, replica_id: int) -> StagePoolClient: + ctx = self._stage_remote_factory_contexts.get(stage_id) + if ctx is None: + raise ValueError(f"No factory context for stage {stage_id}") + return self._create_remote_replica_client(ctx, replica_id) + + def _create_remote_replica_client( + self, + ctx: StageRemoteFactoryContext, + replica_id: int, + ) -> StagePoolClient: + """Create the head-side client for a remote replica. + + Used by both initial remote slots and dynamic headless registrations. + """ + if self._omni_master_server is None: + raise RuntimeError("OmniMasterServer is not running; cannot create remote replica") + stage_id = ctx.stage_id + metadata = copy.deepcopy(ctx.base_metadata) + metadata.replica_id = replica_id + + if ctx.stage_type == "diffusion": + from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient + + resources = None + try: + logger.info( + "[DistStageRuntime] Remote diffusion handshake started stage=%d replica=%d", + stage_id, + replica_id, + ) + with connect_remote_diffusion_proc( + omni_master_server=self._omni_master_server, + stage_id=stage_id, + replica_id=replica_id, + ) as remote_resources: + resources = remote_resources + if resources is None: + raise RuntimeError(f"Remote diffusion stage {stage_id} returned no resources") + if resources.addresses is None: + raise RuntimeError(f"Remote diffusion stage {stage_id} returned no addresses") + except Exception: + self._cleanup_launched_resources( + stage_id=stage_id, + resources=resources, + ) + raise + + client = StageDiffusionClient.from_addresses( + metadata, + request_address=resources.addresses.inputs[0], + response_address=resources.addresses.outputs[0], + batch_size=ctx.diffusion_batch_size, + ) + logger.info( + "[DistStageRuntime] Remote diffusion replica attached stage=%d replica=%d", + stage_id, + replica_id, + ) + return client + + if ctx.vllm_config is None: + raise RuntimeError(f"Remote LLM stage {stage_id} is missing vllm_config") + if ctx.executor_class is None: + raise RuntimeError(f"Remote LLM stage {stage_id} is missing executor_class") + vllm_config = copy.deepcopy(ctx.vllm_config) + vllm_config.parallel_config.data_parallel_size_local = 0 + resources = None + try: + logger.info("[DistStageRuntime] Remote LLM handshake started stage=%d replica=%d", stage_id, replica_id) + with connect_remote_engine_cores( + vllm_config=vllm_config, + omni_master_server=self._omni_master_server, + stage_id=stage_id, + replica_id=replica_id, + ) as remote_resources: + resources = remote_resources + if resources is None: + raise RuntimeError(f"Remote LLM stage {stage_id} returned no resources") + if resources.addresses is None: + raise RuntimeError(f"Remote LLM stage {stage_id} returned no addresses") + client_addresses = self._client_addresses_from_zmq(resources.addresses) + replica_host = self._omni_master_server.get_replica_host(stage_id, replica_id) + if replica_host: + client_addresses["replica_host"] = replica_host + client = StageEngineCoreClientBase.make_async_mp_client( + vllm_config=vllm_config, + executor_class=ctx.executor_class, + metadata=metadata, + client_addresses=client_addresses, + engine_manager=resources.manager, + coordinator=resources.coordinator, + ) + logger.info("[DistStageRuntime] Remote LLM replica attached stage=%d replica=%d", stage_id, replica_id) + return client + except Exception: + self._cleanup_launched_resources( + stage_id=stage_id, + resources=resources, + ) + raise + + +# =========================================================================== +# Factory +# =========================================================================== + + +def create_stage_runtime( + stage_configs: list[Any], + model: str, + config_path: str, + *, + single_stage_mode: bool, + stage_init_timeout: int, + diffusion_batch_size: int, + async_chunk: bool, + tokenizer: str | None = None, + # Distributed-only params: + single_stage_id_filter: int | None = None, + omni_master_address: str | None = None, + omni_master_port: int | None = None, + omni_dp_size_local: int = 1, + omni_heartbeat_timeout: float = 30.0, + omni_lb_policy: str = "random", + request_queue: janus.Queue[EngineQueueMessage] | None = None, +) -> StageRuntime: + """Factory: select StageRuntime or DistStageRuntime.""" + if single_stage_mode: + if not omni_master_address or not omni_master_port: + raise ValueError("Distributed mode requires omni_master_address and omni_master_port") + return DistStageRuntime( + stage_configs=stage_configs, + model=model, + config_path=config_path, + stage_init_timeout=stage_init_timeout, + diffusion_batch_size=diffusion_batch_size, + async_chunk=async_chunk, + tokenizer=tokenizer, + single_stage_id_filter=single_stage_id_filter, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_dp_size_local=omni_dp_size_local, + omni_heartbeat_timeout=omni_heartbeat_timeout, + omni_lb_policy=omni_lb_policy, + request_queue=request_queue, + ) + return StageRuntime( + stage_configs=stage_configs, + model=model, + config_path=config_path, + stage_init_timeout=stage_init_timeout, + diffusion_batch_size=diffusion_batch_size, + async_chunk=async_chunk, + tokenizer=tokenizer, + ) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 9468a632f..905ff4ede 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -9,10 +9,7 @@ import argparse import json import os import signal -import threading -from multiprocessing import connection from types import FrameType -from typing import Any import uvloop from vllm.entrypoints.cli.types import CLISubcommand @@ -697,33 +694,24 @@ def run_headless(args: TrackingNamespace) -> None: headless invocations can coexist) and reports heartbeats to the head's OmniCoordinator. """ - from vllm.v1.engine.coordinator import DPCoordinator from vllm.v1.executor.multiproc_executor import MultiprocExecutor from vllm.version import __version__ as VLLM_VERSION - from vllm_omni.diffusion.stage_diffusion_proc import ( - complete_diffusion_handshake, - spawn_diffusion_proc, - ) from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_omni_kv_config_for_stage - from vllm_omni.engine.omni_core_engine_proc_manager import OmniCoreEngineProcManager - from vllm_omni.engine.stage_engine_startup import register_stage_with_omni_master + from vllm_omni.engine.stage_engine_startup import ( + get_headless_replica_devices, + launch_headless_diffusion_replicas, + launch_headless_llm_replicas, + ) from vllm_omni.engine.stage_init_utils import ( - build_diffusion_config, build_engine_args_dict, build_vllm_config, - extract_stage_metadata, get_stage_connector_spec, - get_stage_devices_per_replica, - inject_kv_stage_info, + inject_omni_kv_connector_config, load_omni_transfer_config_for_model, prepare_engine_environment, - setup_stage_devices, - split_devices_for_replicas, - terminate_alive_proc, ) - from vllm_omni.entrypoints.utils import inject_omni_kv_config, load_and_resolve_stage_configs - from vllm_omni.platforms import current_omni_platform + from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs model = args.model stage_id: int | None = args.stage_id @@ -778,155 +766,25 @@ def run_headless(args: TrackingNamespace) -> None: ) prepare_engine_environment() - omni_transfer_config = load_omni_transfer_config_for_model(model, config_path) - omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) - - # When ``--omni-dp-size-local > 1``, slice the YAML's ``devices:`` field - # into per-replica subsets so each subprocess we spawn below sees a - # narrowed ``CUDA_VISIBLE_DEVICES`` and doesn't stack on cuda:0. Mirrors - # the head-side per-replica device application at - # ``async_omni_engine.py`` (setup_stage_devices around each launch). - runtime_cfg = getattr(stage_cfg, "runtime", None) - devices_str: str | None = None - if runtime_cfg is not None: - devices_str = ( - runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) - ) - devices_per_replica = get_stage_devices_per_replica(stage_cfg) - if devices_str: - # Always remap YAML's logical devices through setup_stage_devices, - # even for omni_dp_size_local==1. The launcher's CUDA_VISIBLE_DEVICES - # is dropped from the engine-subprocess env between vllm-serve and - # OmniCoreEngineProcManager.Process, so the worker would otherwise - # default cuda:0 to physical GPU 0 and collide with a co-located - # head on the same host (see hyi3_multi_host_1 reproducer). - per_replica_devices: list[str | None] = split_devices_for_replicas( - devices_str, omni_dp_size_local, devices_per_replica, stage_id - ) - logger.info( - "[Headless] Stage %d: %d local replicas, devices_per_replica=%d, per-replica devices: %s", - stage_id, - omni_dp_size_local, - devices_per_replica, - per_replica_devices, - ) - else: - per_replica_devices = [None] * omni_dp_size_local - device_control_env = current_omni_platform.device_control_env_var + per_replica_devices = get_headless_replica_devices(stage_cfg, stage_id, omni_dp_size_local) if stage_cfg.stage_type == "diffusion": - metadata = extract_stage_metadata(stage_cfg) - if omni_conn_cfg: - inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) - # Headless single-stage launch must still infer cross-stage TP topology - # from the loaded deploy config so heterogeneous KV routing keys match - # the head process (e.g. from_tp=2, to_tp=1). - inject_kv_stage_info(stage_cfg, stage_id, stage_configs) - od_config = build_diffusion_config(model, stage_cfg, metadata) - - logger.info( - "[Headless] Launching %d diffusion replica(s) for stage %d via OmniMasterServer at %s:%d", - omni_dp_size_local, - stage_id, - omni_master_address, - omni_master_port, + launch_headless_diffusion_replicas( + model=model, + stage_cfg=stage_cfg, + stage_configs=stage_configs, + stage_id=stage_id, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_dp_size_local=omni_dp_size_local, + per_replica_devices=per_replica_devices, + config_path=config_path, + replica_bind_address=omni_replica_address, ) + return - procs: list[Any] = [] - try: - for _rep_idx in range(omni_dp_size_local): - # Always auto-assign: headless processes carry no knowledge - # of their per-replica id and the master server is the sole - # authority on the per-stage id namespace. - response = register_stage_with_omni_master( - omni_master_address=omni_master_address, - omni_master_port=omni_master_port, - omni_stage_id=stage_id, - omni_stage_config=stage_cfg, - replica_id=None, - return_full_response=True, - replica_bind_address=omni_replica_address, - ) - # Apply this replica's CUDA_VISIBLE_DEVICES (only when - # ``--omni-dp-size-local > 1`` and the YAML's stage devices - # field is set). The spawned subprocess inherits the env at - # spawn time; we restore the parent env afterwards so the - # next replica's setup sees the same baseline. - previous_visible_devices = os.environ.get(device_control_env) - try: - if per_replica_devices[_rep_idx] is not None: - setup_stage_devices(stage_id, {"devices": per_replica_devices[_rep_idx]}) - # Each StageDiffusionProc starts its own - # torch.distributed group bound to - # ``od_config.master_port``. Without an explicit - # per-replica override all spawned subprocesses - # share the value ``OmniDiffusionConfig.__post_init__`` - # picked once (and the second binder hits EADDRINUSE - # on ``init_process_group``). We can't use - # kernel-ephemeral allocation either, because the - # master server's pre-allocated ZMQ ports (returned - # by ``register_stage_with_omni_master``) also live - # in the ephemeral range and are not actually bound - # until the headless ``_perform_diffusion_handshake`` - # runs — so picking an ephemeral port here can steal - # a port the master server already promised to a - # sibling headless. Use ``settle_port`` from a base - # above the Linux default ephemeral range - # (32768-60999) so torch.distributed master ports - # never overlap with ZMQ allocations. - if omni_dp_size_local > 1: - od_config.master_port = od_config.settle_port( - 61000 + _rep_idx * 100, - port_inc=37, - ) - proc, _, _, _ = spawn_diffusion_proc( - model, - od_config, - handshake_address=response.handshake_address, - request_address=response.input_address, - response_address=response.output_address, - omni_coordinator_address=response.coordinator_router_address, - omni_stage_id=stage_id, - omni_replica_id=response.replica_id, - ) - finally: - if previous_visible_devices is None: - current_omni_platform.unset_device_control_env_var() - else: - current_omni_platform.set_device_control_env_var(previous_visible_devices) - complete_diffusion_handshake(proc, response.handshake_address, args.stage_init_timeout) - procs.append(proc) - logger.info( - "[Headless] Diffusion replica id=%d for stage %d is up (coord=%s)", - response.replica_id, - stage_id, - response.coordinator_router_address, - ) - - # Block on the sentinel set so any replica crash is detected - # immediately (the previous per-proc join loop only noticed - # crashes in registration order). Any exit triggers fleet - # shutdown via the finally block; non-zero exits propagate. - sentinel_to_proc = {p.sentinel: p for p in procs} - died = connection.wait(list(sentinel_to_proc.keys())) - first = sentinel_to_proc[died[0]] - logger.info( - "[Headless] Diffusion replica %s exited (code=%s); shutting down stage %d.", - first.name, - first.exitcode, - stage_id, - ) - if first.exitcode not in (None, 0): - raise RuntimeError( - f"Diffusion stage {stage_id} replica {first.name!r} exited with code {first.exitcode}" - ) - return - finally: - logger.info("[Headless] Shutting down %d diffusion replica(s) for stage %d.", len(procs), stage_id) - for p in procs: - if p.is_alive(): - terminate_alive_proc(p) - + omni_transfer_config = load_omni_transfer_config_for_model(model, config_path) + omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) stage_connector_spec = get_stage_connector_spec( omni_transfer_config=omni_transfer_config, stage_id=stage_id, @@ -944,17 +802,7 @@ def run_headless(args: TrackingNamespace) -> None: cli_tokenizer=getattr(args, "tokenizer", None), ) - # Inject omni KV connector config so the engine runner can initialize the - # correct connector (sender/receiver role, type, addresses, etc.). - if omni_conn_cfg: - omni_kv = engine_args_dict.get("omni_kv_config") or {} - if not isinstance(omni_kv, dict): - omni_kv = dict(omni_kv) - omni_kv["connector_config"] = omni_conn_cfg - omni_kv["omni_from_stage"] = omni_from - omni_kv["omni_to_stage"] = omni_to - omni_kv.setdefault("stage_id", stage_id) - engine_args_dict["omni_kv_config"] = omni_kv + inject_omni_kv_connector_config(engine_args_dict, omni_kv_connector, stage_id) vllm_config, executor_class = build_vllm_config( stage_cfg, @@ -964,10 +812,6 @@ def run_headless(args: TrackingNamespace) -> None: headless=True, ) parallel_config = vllm_config.parallel_config - local_engine_count = parallel_config.data_parallel_size_local - - if local_engine_count <= 0: - raise ValueError("data_parallel_size_local must be > 0 in headless mode") shutdown_requested = False @@ -994,126 +838,22 @@ def run_headless(args: TrackingNamespace) -> None: executor.start_worker_monitor(inline=True) return - dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 - coordinator = None - if vllm_config.needs_dp_coordinator and dp_rank == 0: - coordinator = DPCoordinator( - parallel_config, - enable_wave_coordination=vllm_config.model_config.is_moe, - ) - logger.info( - "[Headless] Started DP Coordinator process for stage %d (PID: %d)", - stage_id, - coordinator.proc.pid, - ) - - logger.info( - "[Headless] Launching %d omni replica(s) (vLLM dp_size_local=%d each) for stage %d " - "via OmniMasterServer at %s:%d", - omni_dp_size_local, - local_engine_count, - stage_id, - omni_master_address, - omni_master_port, - ) - - # One OmniMasterServer registration per omni replica; each registration - # yields its own (handshake, input, output) allocation and the head's - # OmniCoordinator ROUTER address. We then spawn one - # OmniCoreEngineProcManager per replica so its subprocess gets the - # right replica id wired into its OmniCoordClientForStage. log_stats = bool(args.log_stats) if args.disable_log_stats: log_stats = False - engine_managers: list[Any] = [] - monitor_threads: list[threading.Thread] = [] - - def _monitor_target(mgr: Any) -> None: - try: - mgr.monitor_engine_liveness() - except Exception: - logger.exception("[Headless] monitor_engine_liveness raised") - - try: - for _rep_idx in range(omni_dp_size_local): - # Always auto-assign: see the diffusion branch comment above - # for the rationale (headless owns no replica-id namespace). - response = register_stage_with_omni_master( - omni_master_address=omni_master_address, - omni_master_port=omni_master_port, - omni_stage_id=stage_id, - omni_stage_config=stage_cfg, - coordinator=coordinator, - replica_id=None, - return_full_response=True, - replica_bind_address=omni_replica_address, - # LLM headless: the head binds *all* three sockets — - # handshake ROUTER (``connect_remote_engine_cores``), - # input ROUTER and output PULL (``CoreClient`` — - # ``make_zmq_socket`` defaults bind=True for PULL). - # The remote LLM worker is purely a connector: it - # opens 3 outbound TCP connections to the master's - # host. So the master must keep every address on - # its own host; rewriting any of them to this - # replica's NIC makes the head's ``bind`` go - # EADDRNOTAVAIL on a cross-host launch. - replica_binds_sockets=False, - ) - # Per-replica CUDA_VISIBLE_DEVICES, same pattern as the diffusion - # branch above. OmniCoreEngineProcManager.__init__ spawns its - # subprocesses via context.Process inside the constructor, so we - # must set the env *before* instantiation and restore after. - previous_visible_devices = os.environ.get(device_control_env) - try: - if per_replica_devices[_rep_idx] is not None: - setup_stage_devices(stage_id, {"devices": per_replica_devices[_rep_idx]}) - mgr = OmniCoreEngineProcManager( - local_engine_count=local_engine_count, - start_index=dp_rank, - local_start_index=0, - vllm_config=vllm_config, - local_client=False, - handshake_address=response.handshake_address, - executor_class=executor_class, - log_stats=log_stats, - omni_stage_id=stage_id, - omni_coordinator_address=response.coordinator_router_address, - omni_replica_base_id=response.replica_id, - ) - finally: - if previous_visible_devices is None: - current_omni_platform.unset_device_control_env_var() - else: - current_omni_platform.set_device_control_env_var(previous_visible_devices) - engine_managers.append(mgr) - logger.info( - "[Headless] Stage %d replica id=%d up (coord=%s)", - stage_id, - response.replica_id, - response.coordinator_router_address, - ) - - # Run all managers' liveness monitors in parallel. Each blocks - # until its own subprocesses exit (or fail). - if len(engine_managers) == 1: - engine_managers[0].monitor_engine_liveness() - else: - for mgr in engine_managers: - t = threading.Thread(target=_monitor_target, args=(mgr,), name=f"omni-replica-monitor-{id(mgr):x}") - t.start() - monitor_threads.append(t) - for t in monitor_threads: - t.join() - finally: - logger.info("[Headless] Shutting down stage %d (%d managers).", stage_id, len(engine_managers)) - for mgr in engine_managers: - try: - mgr.shutdown() - except Exception: - logger.exception("[Headless] engine manager shutdown failed") - if coordinator is not None: - coordinator.shutdown() + launch_headless_llm_replicas( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=log_stats, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + stage_id=stage_id, + stage_config=stage_cfg, + omni_dp_size_local=omni_dp_size_local, + per_replica_devices=per_replica_devices, + replica_bind_address=omni_replica_address, + ) def cmd_init() -> list[CLISubcommand]: diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 74b533b51..c899b2697 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -11,9 +11,39 @@ from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) +def resolve_stage_physical_devices( + stage_id: int, + devices: str | int | None, + *, + visible_baseline: str | None = None, +) -> str | None: + """Map logical stage devices to physical IDs without mutating process env. + + When ``visible_baseline`` is provided it is used for logical-to-physical + mapping instead of the current device-control env var. This keeps parallel + stage initialization correct even if another thread has temporarily + narrowed ``CUDA_VISIBLE_DEVICES`` for a different stage. + """ + if devices in (None, "cpu"): + return None + + if not isinstance(devices, (int, str)): + raise TypeError(f"Expected str or int device IDs for stage initialization, got type {type(devices)}") + + device_list = _parse_device_list(devices) + env_var = current_omni_platform.device_control_env_var + visible_devices = visible_baseline if visible_baseline is not None else os.environ.get(env_var) + if visible_devices is not None: + visible_device_list = _parse_device_list(visible_devices) + device_list = _map_device_list(stage_id, device_list, visible_device_list) + return ",".join(device_list) + + def set_stage_devices( stage_id: int, devices: str | int | None, + *, + visible_baseline: str | None = None, ) -> str | None: """Configure per-stage device visibility and current device (CUDA or NPU). @@ -42,19 +72,18 @@ def set_stage_devices( The list of physical devices that were set for the given stage or None if we have no passed devices / are using cpu. """ - env_var = current_omni_platform.device_control_env_var - vis = os.environ.get(env_var) - if devices in (None, "cpu"): logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices) return None elif isinstance(devices, (int, str)): - device_list = _parse_device_list(devices) - if vis is not None: - visible_device_list = _parse_device_list(vis) - device_list = _map_device_list(stage_id, device_list, visible_device_list) - device_str = ",".join(device_list) + device_str = resolve_stage_physical_devices( + stage_id, + devices, + visible_baseline=visible_baseline, + ) + if device_str is None: + return None current_omni_platform.set_device_control_env_var(device_str) return device_str