From cb91cbe61e87c4a2fec5c11d1597e4c2bd922ad2 Mon Sep 17 00:00:00 2001 From: wuhang Date: Fri, 10 Apr 2026 14:40:37 +0800 Subject: [PATCH] [Entrypoint][Refactor]Stage CLI Refactor (#2020) Signed-off-by: wuhang Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com> --- .../bagel/run_server_stage_cli.sh | 178 +- tests/conftest.py | 245 ++- tests/dfx/perf/scripts/run_benchmark.py | 2 +- .../test_flux_2_dev_expansion.py | 16 +- .../test_qwen3_omni_expansion.py | 10 +- tests/engine/test_arg_utils.py | 12 + .../test_async_omni_engine_stage_init.py | 23 +- tests/engine/test_single_stage_mode.py | 1645 +++++++++++++++++ tests/entrypoints/test_serve.py | 195 ++ tests/entrypoints/test_utils.py | 69 + vllm_omni/diffusion/stage_diffusion_client.py | 55 +- vllm_omni/diffusion/stage_diffusion_proc.py | 9 +- vllm_omni/engine/arg_utils.py | 23 +- vllm_omni/engine/async_omni_engine.py | 526 ++++-- vllm_omni/engine/stage_engine_core_client.py | 85 +- vllm_omni/engine/stage_engine_startup.py | 599 ++++++ vllm_omni/engine/stage_init_utils.py | 115 +- vllm_omni/entrypoints/cli/serve.py | 236 ++- vllm_omni/entrypoints/openai/api_server.py | 3 +- 19 files changed, 3773 insertions(+), 273 deletions(-) create mode 100644 tests/engine/test_single_stage_mode.py create mode 100644 tests/entrypoints/test_serve.py create mode 100644 vllm_omni/engine/stage_engine_startup.py diff --git a/examples/online_serving/bagel/run_server_stage_cli.sh b/examples/online_serving/bagel/run_server_stage_cli.sh index 51639153f..2d0b4bc36 100644 --- a/examples/online_serving/bagel/run_server_stage_cli.sh +++ b/examples/online_serving/bagel/run_server_stage_cli.sh @@ -1,34 +1,164 @@ #!/bin/bash -# Bagel multi-stage online serving startup script -# Starts stage 0 as master with API server, and stage 1 in headless mode +# Bagel multi-stage online serving startup script. +# +# Usage: +# ./run_server_stage_cli.sh --stage 0 +# ./run_server_stage_cli.sh --stage 1 +# ./run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2 +# ./run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9 +# +# By default, `--stage all` keeps the old behavior and launches both stages in +# one session. Use `--stage 0` / `--stage 1` to launch each stage separately in +# different terminal sessions, with stage-specific extra CLI arguments passed +# after `--`. + +set -euo pipefail MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}" PORT="${PORT:-8091}" MASTER_ADDRESS="${MASTER_ADDRESS:-127.0.0.1}" MASTER_PORT="${MASTER_PORT:-8092}" -STAGE_CONFIGS_PATH="$(dirname "$0")/../../../vllm_omni/model_executor/stage_configs/bagel.yaml" +STAGE="all" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +STAGE_CONFIGS_PATH="${STAGE_CONFIGS_PATH:-$SCRIPT_DIR/../../../vllm_omni/model_executor/stage_configs/bagel.yaml}" +EXTRA_ARGS=() + +usage() { + cat <&2 + usage + exit 1 + ;; + esac +done + +if [[ "$STAGE" != "0" && "$STAGE" != "1" && "$STAGE" != "all" ]]; then + echo "Invalid --stage value: $STAGE" >&2 + usage + exit 1 +fi + +print_config() { + echo "Model: $MODEL" + echo "API Port: $PORT" + echo "Master Address: $MASTER_ADDRESS" + echo "Master Port: $MASTER_PORT" + echo "Stage Configs: $STAGE_CONFIGS_PATH" + echo "Selected Stage: $STAGE" + if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then + echo "Extra Args: ${EXTRA_ARGS[*]}" + fi +} + +run_stage_0() { + echo "Starting Stage 0 (Thinker) as master..." + vllm serve "$MODEL" --omni \ + --port "$PORT" \ + --stage-configs-path "$STAGE_CONFIGS_PATH" \ + --stage-id 0 \ + -oma "$MASTER_ADDRESS" \ + -omp "$MASTER_PORT" \ + "${EXTRA_ARGS[@]}" +} + +run_stage_1() { + echo "Starting Stage 1 (DiT) in headless mode..." + vllm serve "$MODEL" --omni \ + --stage-configs-path "$STAGE_CONFIGS_PATH" \ + --stage-id 1 \ + --headless \ + -oma "$MASTER_ADDRESS" \ + -omp "$MASTER_PORT" \ + "${EXTRA_ARGS[@]}" +} echo "Starting Bagel multi-stage server..." -echo "Model: $MODEL" -echo "API Port: $PORT" -echo "Master Address: $MASTER_ADDRESS" -echo "Master Port: $MASTER_PORT" -echo "Stage Configs: $STAGE_CONFIGS_PATH" +print_config -# Start stage 1 (DiT) in headless mode first -echo "Starting Stage 1 (DiT) in headless mode..." -vllm serve "$MODEL" --omni \ - --stage-configs-path "$STAGE_CONFIGS_PATH" \ - --stage-id 1 \ - --headless \ - -oma "$MASTER_ADDRESS" \ - -omp "$MASTER_PORT" & +case "$STAGE" in + 0) + run_stage_0 + ;; + 1) + run_stage_1 + ;; + all) + echo "Launching both stages in one session (legacy mode)..." + echo "Starting Stage 0 (Thinker) in background first..." + run_stage_0 & + STAGE_0_PID=$! -# Start stage 0 (Thinker) as master with API server -echo "Starting Stage 0 (Thinker) as master..." -vllm serve "$MODEL" --omni \ - --port "$PORT" \ - --stage-configs-path "$STAGE_CONFIGS_PATH" \ - --stage-id 0 \ - -oma "$MASTER_ADDRESS" \ - -omp "$MASTER_PORT" + cleanup() { + if [[ -n "${STAGE_0_PID:-}" ]]; then + kill "$STAGE_0_PID" 2>/dev/null || true + wait "$STAGE_0_PID" 2>/dev/null || true + fi + } + + trap cleanup EXIT INT TERM + + echo "Waiting briefly for Stage 0 to initialize..." + sleep 2 + run_stage_1 + ;; +esac diff --git a/tests/conftest.py b/tests/conftest.py index 8ac790f13..27833fe28 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -73,6 +73,8 @@ class OmniServerParams(NamedTuple): server_args: list[str] | None = None env_dict: dict[str, str] | None = None use_omni: bool = True + use_stage_cli: bool = False + init_timeout: int | None = None def assert_image_diffusion_response( @@ -1546,6 +1548,183 @@ class OmniServer: cleanup_dist_env_and_memory() +class OmniServerStageCli(OmniServer): + """Omni server harness that exercises the stage CLI flow.""" + + def __init__( + self, + model: str, + stage_config_path: str, + serve_args: list[str] | None = None, + *, + stage_ids: list[int] | None = None, + port: int | None = None, + env_dict: dict[str, str] | None = None, + ) -> None: + super().__init__(model, serve_args or [], port=port, env_dict=env_dict, use_omni=True) + self.stage_config_path = stage_config_path + self.master_port = get_open_port() + self.visible_device_list = self._load_visible_device_list(env_dict) + self.stage_runtime_devices = self._load_stage_runtime_devices(stage_config_path) + self.stage_ids = stage_ids or self._load_stage_ids(stage_config_path) + if 0 not in self.stage_ids: + raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}") + self.stage_procs: dict[int, subprocess.Popen] = {} + self.proc = None + + @staticmethod + def _load_stage_ids(stage_config_path: str) -> list[int]: + with open(stage_config_path, encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + + stage_ids = [stage["stage_id"] for stage in cfg.get("stage_args", []) if "stage_id" in stage] + if not stage_ids: + raise ValueError(f"No stage IDs found in config: {stage_config_path}") + return stage_ids + + @staticmethod + def _load_stage_runtime_devices(stage_config_path: str) -> dict[int, str]: + with open(stage_config_path, encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + + runtime_devices: dict[int, str] = {} + for stage in cfg.get("stage_args", []): + stage_id = stage.get("stage_id") + devices = stage.get("runtime", {}).get("devices") + if stage_id is not None and devices: + runtime_devices[int(stage_id)] = str(devices) + return runtime_devices + + @classmethod + def _parse_device_list(cls, devices: str | int) -> list[str]: + if isinstance(devices, int): + if devices < 0: + raise ValueError("Device IDs must be non-negative integers") + return [str(devices)] + return [token.strip() for token in str(devices).split(",") if token.strip()] + + @classmethod + def _load_visible_device_list(cls, env_dict: dict[str, str] | None) -> list[str] | None: + env = os.environ.copy() + if env_dict is not None: + env.update(env_dict) + + env_var = getattr(current_omni_platform, "device_control_env_var", None) + if env_var and env_var in env: + return [token.strip() for token in env[env_var].split(",") if token.strip()] + return None + + @classmethod + def _map_stage_devices(cls, stage_id: int, visible_device_list: list[str] | None, devices: str) -> str: + device_list = cls._parse_device_list(devices) + + if visible_device_list is None: + return ",".join(device_list) + + if not all(device.isdigit() for device in device_list): + raise ValueError("Logical devices must be non-negative integers") + + logical_ids = [int(device) for device in device_list] + if logical_ids and max(logical_ids) >= len(visible_device_list): + raise ValueError( + f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices" + ) + + return ",".join(visible_device_list[idx] for idx in logical_ids) + + def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str) -> None: + mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, devices) + env_var = getattr(current_omni_platform, "device_control_env_var", None) + if env_var: + env[env_var] = mapped_devices + + def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]: + cmd = [ + sys.executable, + "-m", + "vllm_omni.entrypoints.cli.main", + "serve", + self.model, + "--omni", + "--stage-configs-path", + self.stage_config_path, + "--stage-id", + str(stage_id), + "--omni-master-address", + self.host, + "--omni-master-port", + str(self.master_port), + ] + + if headless: + cmd.append("--headless") + else: + cmd += ["--host", self.host, "--port", str(self.port)] + + cmd += self.serve_args + return cmd + + def _launch_stage(self, stage_id: int, *, headless: bool) -> None: + env = os.environ.copy() + env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + if self.env_dict is not None: + env.update(self.env_dict) + + devices = self.stage_runtime_devices.get(stage_id) + if devices: + self._set_stage_device_env(stage_id, env, devices) + + cmd = self._build_stage_cmd(stage_id, headless=headless) + print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}") + proc = subprocess.Popen( + cmd, + env=env, + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + ) + self.stage_procs[stage_id] = proc + if stage_id == 0: + self.proc = proc + + def _ensure_stage_processes_alive(self) -> None: + for stage_id, proc in self.stage_procs.items(): + ret = proc.poll() + if ret is not None: + raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.") + + def _start_server(self) -> None: + ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]] + + self._launch_stage(0, headless=False) + time.sleep(2) + self._ensure_stage_processes_alive() + + for stage_id in ordered_stage_ids[1:]: + self._launch_stage(stage_id, headless=True) + + max_wait = 1200 + start_time = time.time() + while time.time() - start_time < max_wait: + self._ensure_stage_processes_alive() + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(1) + result = sock.connect_ex((self.host, self.port)) + if result == 0: + print(f"OmniServerStageCli ready on {self.host}:{self.port}") + return + time.sleep(2) + + raise RuntimeError(f"OmniServerStageCli failed to start within {max_wait} seconds") + + def __exit__(self, exc_type, exc_val, exc_tb): + for stage_id in sorted(self.stage_procs, reverse=True): + proc = self.stage_procs[stage_id] + if proc.poll() is None: + self._kill_process_tree(proc.pid) + _run_pre_test_cleanup(enable_force=True) + _run_post_test_cleanup(enable_force=True) + cleanup_dist_env_and_memory() + + def pytest_addoption(parser): parser.addoption( "--run-level", @@ -1568,9 +1747,11 @@ _omni_server_lock = threading.Lock() @pytest.fixture(scope="module") def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]: - """Start vLLM-Omni server as a subprocess with actual model weights. - Uses session scope so the server starts only once for the entire test session. - Multi-stage initialization can take 10-20+ minutes. + """Start vLLM-Omni through the standard or stage-CLI launcher. + + The fixture stays module-scoped because multi-stage initialization is costly. + The ``use_stage_cli`` flag on ``OmniServerParams`` routes the setup through the + stage-CLI harness while still reusing the same fixture grouping semantics. """ with _omni_server_lock: params: OmniServerParams = request.param @@ -1589,28 +1770,47 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st server_args = params.server_args or [] if params.use_omni: server_args = ["--stage-init-timeout", "120", *server_args] - if stage_config_path is not None: - server_args += ["--stage-configs-path", stage_config_path] + if params.init_timeout is not None: + server_args = [*server_args, "--init-timeout", str(params.init_timeout)] + if params.use_stage_cli: + if not params.use_omni: + raise ValueError("omni_server with use_stage_cli=True requires use_omni=True") + if stage_config_path is None: + raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path") - with ( - OmniServer( + with OmniServerStageCli( model, + stage_config_path, server_args, port=port, env_dict=params.env_dict, - use_omni=params.use_omni, - ) - if port - else OmniServer( - model, - server_args, - env_dict=params.env_dict, - use_omni=params.use_omni, - ) - ) as server: - print("OmniServer started successfully") - yield server - print("OmniServer stopping...") + ) as server: + print("OmniServer started successfully") + yield server + print("OmniServer stopping...") + else: + if stage_config_path is not None: + server_args += ["--stage-configs-path", stage_config_path] + + with ( + OmniServer( + model, + server_args, + port=port, + env_dict=params.env_dict, + use_omni=params.use_omni, + ) + if port + else OmniServer( + model, + server_args, + env_dict=params.env_dict, + use_omni=params.use_omni, + ) + ) as server: + print("OmniServer started successfully") + yield server + print("OmniServer stopping...") print("OmniServer stopped") @@ -2653,10 +2853,11 @@ class OpenAIClientHandler: @pytest.fixture -def openai_client(omni_server: OmniServer, run_level: str): +def openai_client(request: pytest.FixtureRequest, run_level: str): """Create OpenAIClientHandler fixture to facilitate communication with OmniServer with encapsulated request sending, concurrent requests, response handling, and validation.""" - return OpenAIClientHandler(host=omni_server.host, port=omni_server.port, api_key="EMPTY", run_level=run_level) + server = request.getfixturevalue("omni_server") + return OpenAIClientHandler(host=server.host, port=server.port, api_key="EMPTY", run_level=run_level) class OmniRunner: diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 9e375fa9f..c625239e5 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -43,7 +43,7 @@ def omni_server(request): print(f"Starting OmniServer with test: {test_name}, model: {model}") - server_args = ["--stage-init-timeout", "120"] + server_args = ["--stage-init-timeout", "120", "--init-timeout", "900"] if stage_config_path: server_args = ["--stage-configs-path", stage_config_path] + server_args with OmniServer(model, server_args) as server: diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py index c7140769b..9d96a48c0 100644 --- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py +++ b/tests/e2e/online_serving/test_flux_2_dev_expansion.py @@ -2,13 +2,11 @@ End-to-end diffusion coverage for FLUX.2-dev in online serving mode. Coverage: -- Cache-DiT cache acceleration backend - CPU offload -This test verifies that FLUX.2-dev can be launched with the Cache-DiT backend -and CPU offload enabled, accepts text-to-image requests through the -OpenAI-compatible API, and returns valid generated images with the requested -resolution. +This test verifies that FLUX.2-dev can be launched with CPU offload enabled, +accepts text-to-image requests through the OpenAI-compatible API, and returns +valid generated images with the requested resolution. assert_diffusion_response validates successful generation and the expected image resolution. @@ -33,19 +31,17 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2) def _get_flux_2_dev_feature_cases(model: str): - """Return FLUX.2-dev diffusion feature cases for Cache-DiT + CPU offload.""" + """Return FLUX.2-dev diffusion feature cases for CPU offload.""" return [ pytest.param( OmniServerParams( model=model, server_args=[ - "--cache-backend", - "cache_dit", "--enable-cpu-offload", ], ), - id="cache_dit_cpu_offload", + id="cpu_offload", marks=SINGLE_CARD_FEATURE_MARKS, ), pytest.param( @@ -76,7 +72,7 @@ def test_flux_2_dev( omni_server: OmniServer, openai_client: OpenAIClientHandler, ): - """Validate FLUX.2-dev online serving with Cache-DiT and CPU offload.""" + """Validate FLUX.2-dev online serving with CPU offload.""" messages = dummy_messages_from_mix_data(content_text=PROMPT) diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 0bcc86840..163762769 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -67,13 +67,17 @@ if current_omni_platform.is_xpu(): # Create parameter combinations for model and stage config test_params = [ - pytest.param(OmniServerParams(model=model, stage_config_path=default_path), id="default"), - pytest.param(OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path)), id="async_chunk"), + pytest.param(OmniServerParams(model=model, stage_config_path=default_path, use_stage_cli=True), id="default"), + pytest.param( + OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path), use_stage_cli=True), + id="async_chunk", + ), ] test_token_params = [ pytest.param( - OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path)), id="batch_token_64" + OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path), use_stage_cli=True), + id="batch_token_64", ) ] diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 7ba1cebec..5584b15d9 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -24,6 +24,18 @@ def test_sync_config_is_omni(): assert isinstance(cfg, OmniModelConfig) +def test_default_stage_id_is_concrete_int(): + """Ensure `stage_id` stays safe for downstream arithmetic/indexing.""" + engine_args = OmniEngineArgs() + + assert engine_args.stage_id == 0 + assert isinstance(engine_args.stage_id, int) + assert engine_args.log_stats is False + + cfg = engine_args.create_model_config() + assert cfg.stage_id == 0 + + def test_multimodal_kwarg_overrides(): """Ensure that overrides in the multimodal config are preserved.""" # Get a different value than the default for a multimodal field diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 9f47fd449..31d3ed775 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1,3 +1,4 @@ +import importlib import os import types @@ -8,6 +9,17 @@ from vllm_omni.engine.async_omni_engine import AsyncOmniEngine pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): + """Regression test for forward references in make_async_mp_client.""" + import vllm_omni.engine.stage_engine_core_client as client_mod + + importlib.reload(client_mod) + + assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == ( + "StageEngineCoreClient | DPLBStageEngineCoreClient" + ) + + def test_initialize_stages_restores_device_visibility_after_diffusion_init(monkeypatch): """Regression test for stage device env leakage across stage init. @@ -23,6 +35,9 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke 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.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")] env_var = current_omni_platform.device_control_env_var @@ -49,7 +64,7 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke current_omni_platform.set_device_control_env_var("1") monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices) - monkeypatch.setattr(engine_mod, "_inject_kv_stage_info", lambda *_: None) + monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None) monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: diffusion_client) monkeypatch.setattr( engine_mod, @@ -101,7 +116,11 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch): self.vllm_config = vllm_config self.renderer = renderer - monkeypatch.setattr(engine_mod, "StageEngineCoreClient", DummyStageEngineCoreClient) + monkeypatch.setattr( + engine_mod.StageEngineCoreClientBase, + "make_async_mp_client", + staticmethod(lambda **kwargs: DummyStageEngineCoreClient(**kwargs)), + ) monkeypatch.setattr(engine_mod, "MultimodalOutputProcessor", DummyOutputProcessor) monkeypatch.setattr(engine_mod, "InputProcessor", DummyInputProcessor) monkeypatch.setattr(engine_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor) diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py new file mode 100644 index 000000000..627a98395 --- /dev/null +++ b/tests/engine/test_single_stage_mode.py @@ -0,0 +1,1645 @@ +"""Unit tests for AsyncOmniEngine single-stage mode and OmniMasterServer. + +These tests cover: +- OmniMasterServer address pre-allocation & ZMQ registration handshake +- AsyncOmniEngine single_stage_mode detection / _single_stage_id_filter setup +- _initialize_stages stage routing (local launch vs. remote-wait) in + single_stage_mode +- _create_remote_llm_stage delegation to connect_remote_engine_cores +- _launch_llm_stage delegation to launch_omni_core_engines in + single_stage_mode + +All tests run without real hardware by mocking ZMQ, vllm_config, and the +heavy initialization helpers. +""" + +from __future__ import annotations + +import threading +from contextlib import contextmanager +from typing import Any +from unittest.mock import MagicMock, Mock, patch + +import pytest +from vllm.v1.engine.utils import EngineZmqAddresses + +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.stage_engine_startup import ( + OmniMasterServer, + StageAllocation, + StageCoordinatorAddresses, + connect_remote_engine_cores, + launch_omni_core_engines, +) +from vllm_omni.engine.stage_init_utils import StartedLlmStage + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_stage_cfg(stage_id: int, stage_type: str = "llm") -> Mock: + """Return a lightweight stage config mock.""" + cfg = Mock() + cfg.stage_id = stage_id + cfg.stage_type = stage_type + cfg.engine_args = MagicMock() + cfg.engine_args.async_chunk = False + cfg.engine_args.model_stage = None + cfg.engine_args.engine_output_type = None + return cfg + + +def _make_started_llm_stage(stage_id: int) -> StartedLlmStage: + """Return a minimal StartedLlmStage for mocking.""" + addresses = Mock() + addresses.inputs = ["tcp://127.0.0.1:5000"] + addresses.outputs = ["tcp://127.0.0.1:5001"] + addresses.frontend_stats_publish_address = None + return StartedLlmStage( + stage_id=stage_id, + metadata=Mock(stage_id=stage_id), + vllm_config=Mock(), + executor_class=Mock(), + engine_manager=Mock(), + coordinator=Mock(), + addresses=addresses, + ) + + +# --------------------------------------------------------------------------- +# OmniMasterServer – address pre-allocation +# --------------------------------------------------------------------------- + + +class TestOmniMasterServerAllocation: + """Test address pre-allocation in OmniMasterServer.__init__.""" + + def test_public_address_and_port_properties_expose_registration_endpoint(self): + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15000, + stage_ids=[0], + ) + assert server.address == "127.0.0.1" + assert server.port == 15000 + + def test_allocations_created_for_each_stage_id(self): + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15000, + stage_ids=[0, 1, 2], + ) + assert set(server._allocations.keys()) == {0, 1, 2} + + def test_each_allocation_is_stage_allocation(self): + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15000, + stage_ids=[0, 1], + ) + for sid in (0, 1): + alloc = server._allocations[sid] + assert isinstance(alloc, StageAllocation) + + def test_allocation_addresses_reference_master_address(self): + server = OmniMasterServer( + master_address="192.168.1.10", + master_port=20000, + stage_ids=[0], + ) + alloc = server._allocations[0] + for addr in ( + alloc.handshake_bind_address, + alloc.handshake_connect_address, + alloc.input_bind_address, + alloc.input_connect_address, + alloc.output_bind_address, + alloc.output_connect_address, + ): + assert "192.168.1.10" in addr, f"Expected master address in {addr}" + + def test_port_uniqueness_within_single_allocation(self): + """Each allocation uses three distinct ports.""" + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15001, + stage_ids=[0], + ) + alloc = server._allocations[0] + hs_port = int(alloc.handshake_bind_address.split(":")[-1]) + inp_port = int(alloc.input_bind_address.split(":")[-1]) + out_port = int(alloc.output_bind_address.split(":")[-1]) + assert len({hs_port, inp_port, out_port}) == 3, "Expected three distinct ports per stage allocation" + + def test_get_zmq_addresses_returns_bind_addresses(self): + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15002, + stage_ids=[0], + ) + alloc = server._allocations[0] + zmq_addrs = server.get_zmq_addresses(0) + assert zmq_addrs.inputs == [alloc.input_bind_address] + assert zmq_addrs.outputs == [alloc.output_bind_address] + + def test_get_engine_zmq_addresses_returns_connect_addresses(self): + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15003, + stage_ids=[0], + ) + alloc = server._allocations[0] + engine_addrs = server.get_engine_zmq_addresses(0) + assert engine_addrs.inputs == [alloc.input_connect_address] + assert engine_addrs.outputs == [alloc.output_connect_address] + + def test_get_allocation_returns_correct_object(self): + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15004, + stage_ids=[3], + ) + assert server.get_allocation(3) is server._allocations[3] + + +# --------------------------------------------------------------------------- +# OmniMasterServer – ZMQ registration flow +# --------------------------------------------------------------------------- + + +class TestOmniMasterServerRegistration: + """Test that the server correctly handles a stage registration.""" + + def test_registration_reply_contains_handshake_address(self): + """A DEALER client that sends a registration msg gets the handshake + address back from the ROUTER registration socket.""" + import msgspec + import zmq + from vllm.utils.network_utils import get_open_port + + master_port = get_open_port() + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=master_port, + stage_ids=[0], + ) + server.start() + expected_hs = server._allocations[0].handshake_connect_address + + ctx = zmq.Context() + try: + sock = ctx.socket(zmq.DEALER) + sock.connect(f"tcp://127.0.0.1:{master_port}") + sock.send(msgspec.msgpack.encode({"stage_id": 0})) + if not sock.poll(timeout=5_000): + pytest.fail("No reply received from OmniMasterServer within 5 s") + reply = msgspec.msgpack.decode(sock.recv()) + assert reply["handshake_address"] == expected_hs + finally: + sock.close(linger=0) + ctx.term() + server.stop() + + def test_server_handles_unknown_stage_id_gracefully(self): + """A registration for an unrecognised stage_id must not crash the server.""" + import msgspec + import zmq + from vllm.utils.network_utils import get_open_port + + master_port = get_open_port() + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=master_port, + stage_ids=[0], + ) + server.start() + + ctx = zmq.Context() + try: + bad_sock = ctx.socket(zmq.DEALER) + bad_sock.connect(f"tcp://127.0.0.1:{master_port}") + # Send unknown stage_id=99 + bad_sock.send(msgspec.msgpack.encode({"stage_id": 99})) + # Server should NOT reply for an unknown id; wait briefly + has_reply = bad_sock.poll(timeout=500) + assert not has_reply, "Server should not reply to unknown stage_id" + # Then register the valid stage so the server thread can exit + good_sock = ctx.socket(zmq.DEALER) + good_sock.connect(f"tcp://127.0.0.1:{master_port}") + good_sock.send(msgspec.msgpack.encode({"stage_id": 0})) + good_sock.poll(timeout=2_000) + finally: + for s in (bad_sock, good_sock): + try: + s.close(linger=0) + except Exception: + pass + ctx.term() + server.stop() + + def test_registration_stores_stage_config(self): + """Stage registration should persist the sender's stage config.""" + import msgspec + import zmq + from vllm.utils.network_utils import get_open_port + + master_port = get_open_port() + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=master_port, + stage_ids=[0], + ) + server.start() + + payload = { + "stage_id": 0, + "stage_config": { + "stage_id": 0, + "stage_type": "llm", + "engine_args": {"model": "fake-model"}, + }, + } + + ctx = zmq.Context() + try: + sock = ctx.socket(zmq.DEALER) + sock.connect(f"tcp://127.0.0.1:{master_port}") + sock.send(msgspec.msgpack.encode(payload)) + assert sock.poll(timeout=5_000) + sock.recv() + + stored = server.get_stage_config(0, timeout_s=0.1) + assert stored == payload["stage_config"] + finally: + sock.close(linger=0) + ctx.term() + server.stop() + + def test_registration_stores_coordinator_addresses(self): + """Stage registration should persist optional coordinator addresses.""" + import msgspec + import zmq + from vllm.utils.network_utils import get_open_port + + master_port = get_open_port() + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=master_port, + stage_ids=[0], + ) + server.start() + + payload = { + "stage_id": 0, + "stage_config": {"stage_id": 0}, + "coordinator_input": "tcp://127.0.0.1:31001", + "coordinator_output": "tcp://127.0.0.1:31002", + "frontend_stats_publish_address": "tcp://127.0.0.1:31003", + } + + ctx = zmq.Context() + try: + sock = ctx.socket(zmq.DEALER) + sock.connect(f"tcp://127.0.0.1:{master_port}") + sock.send(msgspec.msgpack.encode(payload)) + assert sock.poll(timeout=5_000) + sock.recv() + + stored = server.get_stage_coordinator_addresses(0, timeout_s=0.1) + assert stored == StageCoordinatorAddresses( + coordinator_input=payload["coordinator_input"], + coordinator_output=payload["coordinator_output"], + frontend_stats_publish_address=payload["frontend_stats_publish_address"], + ) + finally: + sock.close(linger=0) + ctx.term() + server.stop() + + def test_stop_joins_server_thread(self): + from vllm.utils.network_utils import get_open_port + + master_port = get_open_port() + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=master_port, + stage_ids=[], # no stages → thread exits immediately + ) + server.start() + assert server._thread is not None + server.stop() + # Thread should have exited (joined with timeout=10 inside stop()) + assert not server._thread.is_alive() + + +# --------------------------------------------------------------------------- +# AsyncOmniEngine – single_stage_mode detection in __init__ +# --------------------------------------------------------------------------- + + +class TestSingleStageModeDetection: + """Test __init__ single_stage_mode / _single_stage_id_filter setup. + + We bypass the real __init__ by patching _resolve_stage_configs and + the orchestrator thread, so no actual engines are started. + """ + + def _make_engine_no_thread(self, **kwargs: Any) -> AsyncOmniEngine: + """Create an AsyncOmniEngine without starting the orchestrator thread.""" + stage_cfg = _make_stage_cfg(0) + mock_stage_configs = [stage_cfg] + + with ( + patch.object( + AsyncOmniEngine, + "_resolve_stage_configs", + return_value=("/fake/path", mock_stage_configs), + ), + patch.object( + AsyncOmniEngine, + "_bootstrap_orchestrator", + ), + patch("threading.Thread") as mock_thread_cls, + patch("concurrent.futures.Future") as mock_future_cls, + ): + mock_future = Mock() + mock_future.result.return_value = Mock() # simulates a loop + mock_future_cls.return_value = mock_future + + mock_thread = Mock() + mock_thread.is_alive.return_value = False + mock_thread_cls.return_value = mock_thread + + engine = AsyncOmniEngine(model="fake-model", **kwargs) + return engine + + def test_explicit_single_stage_mode_true(self): + engine = self._make_engine_no_thread( + single_stage_mode=True, + omni_master_address="127.0.0.1", + omni_master_port=20000, + ) + assert engine.single_stage_mode is True + + def test_stage_id_kwarg_promotes_to_single_stage_mode(self): + engine = self._make_engine_no_thread( + stage_id=0, + omni_master_address="127.0.0.1", + omni_master_port=20001, + ) + assert engine.single_stage_mode is True + + def test_stage_id_kwarg_sets_filter(self): + engine = self._make_engine_no_thread( + stage_id=1, + omni_master_address="127.0.0.1", + omni_master_port=20002, + ) + assert engine._single_stage_id_filter == 1 + + def test_no_stage_id_no_single_stage_mode(self): + engine = self._make_engine_no_thread() + assert engine.single_stage_mode is False + assert engine._single_stage_id_filter is None + + def test_single_stage_mode_without_stage_id_has_no_filter(self): + engine = self._make_engine_no_thread( + single_stage_mode=True, + omni_master_address="127.0.0.1", + omni_master_port=20003, + ) + assert engine._single_stage_id_filter is None + + def test_master_address_and_port_stored(self): + engine = self._make_engine_no_thread( + stage_id=0, + omni_master_address="10.0.0.1", + omni_master_port=12345, + ) + assert engine._omni_master_address == "10.0.0.1" + assert engine._omni_master_port == 12345 + + def test_omni_master_server_starts_as_none(self): + engine = self._make_engine_no_thread() + assert engine._omni_master_server is None + + +# --------------------------------------------------------------------------- +# AsyncOmniEngine – _initialize_stages stage routing +# --------------------------------------------------------------------------- + + +class TestInitializeStagesRouting: + """Verify that _initialize_stages routes each stage to the correct launch + function depending on single_stage_mode and _single_stage_id_filter.""" + + _COMMON_PATCHES = [ + "vllm_omni.engine.async_omni_engine.prepare_engine_environment", + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + ] + + def _build_engine_skeleton( + self, + stage_cfgs: list[Mock], + single_stage_mode: bool, + stage_id_filter: int | None, + omni_master_address: str = "127.0.0.1", + omni_master_port: int = 25000, + ) -> AsyncOmniEngine: + """Build a bare AsyncOmniEngine without launching any threads.""" + engine = object.__new__(AsyncOmniEngine) + engine.model = "fake-model" + engine.config_path = "/fake" + engine.stage_configs = stage_cfgs + engine.num_stages = len(stage_cfgs) + engine.async_chunk = False + engine.single_stage_mode = single_stage_mode + engine._single_stage_id_filter = stage_id_filter + engine._omni_master_address = omni_master_address + engine._omni_master_port = omni_master_port + engine._omni_master_server = None + engine._llm_stage_launch_lock = __import__("threading").Lock() + engine.diffusion_batch_size = 1 + engine.stage_clients = [] + engine.stage_vllm_configs = [] + engine.output_processors = [] + engine.input_processor = None + engine.supported_tasks = ("generate",) + engine.default_sampling_params_list = [] + engine.stage_metadata = [] + engine.prompt_expand_func = None + return engine + + def _fake_metadata(self, stage_id: int, stage_type: str = "llm") -> Mock: + meta = Mock() + meta.stage_id = stage_id + meta.stage_type = stage_type + meta.runtime_cfg = {} + meta.prompt_expand_func = None + meta.engine_output_type = None + meta.is_comprehension = False + meta.final_output = True if stage_id == 0 else False + meta.final_output_type = None + return meta + + def _run_initialize_stages_mocked( + self, + engine: AsyncOmniEngine, + stage_cfgs: list[Mock], + *, + launch_side_effect: Any = None, + remote_side_effect: Any = None, + attach_result: Any = None, + ) -> tuple[Mock, Mock]: + """Execute _initialize_stages with all heavy helpers mocked. + + Returns (mock_launch_llm_stage, mock_create_remote_llm_stage). + """ + started_by_stage: dict[int, StartedLlmStage] = { + cfg.stage_id: _make_started_llm_stage(cfg.stage_id) + for cfg in stage_cfgs + if getattr(cfg, "stage_type", "llm") != "diffusion" + } + + default_attach = (Mock(), Mock(), Mock(), Mock()) + + mock_launch = Mock( + side_effect=launch_side_effect + or (lambda cfg, meta, spec, timeout, llm_stage_launch_lock, kv: started_by_stage[meta.stage_id]) + ) + mock_remote = Mock( + side_effect=remote_side_effect or (lambda cfg, meta, spec, timeout, srv: started_by_stage[meta.stage_id]) + ) + mock_attach = Mock(return_value=attach_result or default_attach) + + mock_oms = Mock(spec=OmniMasterServer) + mock_oms.get_zmq_addresses.side_effect = lambda sid: Mock() + + finalized = ( + [Mock() for _ in stage_cfgs], + [Mock() for _ in stage_cfgs], + [{"final_output": True, "final_output_type": None, "stage_type": "llm"} for _ in stage_cfgs], + ) + + with ( + patch.object(engine, "_launch_llm_stage", mock_launch), + patch.object(engine, "_create_remote_llm_stage", mock_remote), + patch.object(engine, "_attach_llm_stage", mock_attach), + patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ), + patch( + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + return_value={}, + ), + patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ), + patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(cfg.stage_id, getattr(cfg, "stage_type", "llm")), + ), + patch( + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + return_value=finalized, + ), + ): + engine._initialize_stages(stage_init_timeout=60) + + return mock_launch, mock_remote + + # -- single-stage mode: stage matches filter → local launch --------------- + + def test_matching_stage_uses_launch_llm_stage(self): + """stage_id == _single_stage_id_filter → _launch_llm_stage is called.""" + stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) + mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + + launched_ids = [c.args[1].stage_id for c in mock_launch.call_args_list] + assert 0 in launched_ids, "_launch_llm_stage should be called for stage 0" + + def test_non_matching_stage_uses_create_remote_llm_stage(self): + """stage_id != _single_stage_id_filter → _create_remote_llm_stage is called.""" + stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) + mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + + remote_ids = [c.args[1].stage_id for c in mock_remote.call_args_list] + assert 1 in remote_ids, "_create_remote_llm_stage should be called for stage 1" + + def test_filter_1_routes_correctly(self): + """With filter=1, stage 0 is remote and stage 1 is local.""" + stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=1) + mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + + launched_ids = [c.args[1].stage_id for c in mock_launch.call_args_list] + remote_ids = [c.args[1].stage_id for c in mock_remote.call_args_list] + assert 1 in launched_ids, "stage 1 should be launched locally with filter=1" + assert 0 in remote_ids, "stage 0 should use remote path with filter=1" + + def test_no_filter_all_stages_use_launch_path(self): + """single_stage_mode=True but no filter → all stages use _launch_llm_stage.""" + stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=None) + mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + + assert mock_remote.call_count == 0, "No remote launches without a filter" + launched_ids = [c.args[1].stage_id for c in mock_launch.call_args_list] + assert set(launched_ids) == {0, 1} + + def test_non_single_stage_mode_never_calls_create_remote(self): + """Outside single_stage_mode, _create_remote_llm_stage must not be called.""" + stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=False, stage_id_filter=None) + mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + + assert mock_remote.call_count == 0 + + def test_omni_master_server_started_in_single_stage_mode(self): + """OmniMasterServer.start() must be called when single_stage_mode=True.""" + stage_cfgs = [_make_stage_cfg(0)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) + mock_oms = Mock(spec=OmniMasterServer) + mock_oms.get_zmq_addresses.return_value = Mock() + finalized = ([Mock()], [Mock()], [{"final_output": True, "final_output_type": None, "stage_type": "llm"}]) + + with ( + patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)), + patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(0)), + patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), + patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), + patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + ), + patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(cfg.stage_id), + ), + patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), + ): + engine._initialize_stages(stage_init_timeout=60) + + mock_oms.start.assert_called_once() + + def test_omni_master_server_uses_configured_stage_ids(self): + """Configured stage IDs, not list indexes, should drive pre-allocation.""" + stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=7) + mock_oms = Mock(spec=OmniMasterServer) + mock_oms.get_zmq_addresses.return_value = Mock() + finalized = ( + [Mock(), Mock()], + [Mock(), Mock()], + [{"final_output": False, "final_output_type": None, "stage_type": "llm"} for _ in stage_cfgs], + ) + + with ( + patch.object( + engine, "_launch_llm_stage", side_effect=[_make_started_llm_stage(7), _make_started_llm_stage(11)] + ), + patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(11)), + patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms) as mock_oms_cls, + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), + patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), + patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + ), + patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(cfg.stage_id), + ), + patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), + ): + engine._initialize_stages(stage_init_timeout=60) + + mock_oms_cls.assert_called_once_with( + master_address=engine._omni_master_address, + master_port=engine._omni_master_port, + stage_ids=[7, 11], + ) + + def test_single_stage_filter_uses_configured_stage_ids(self): + """Local/remote dispatch should compare against configured stage IDs.""" + stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=7) + mock_oms = Mock(spec=OmniMasterServer) + finalized = ( + [Mock(), Mock()], + [Mock(), Mock()], + [{"final_output": False, "final_output_type": None, "stage_type": "llm"} for _ in stage_cfgs], + ) + + with ( + patch.object(engine, "_launch_llm_stage", side_effect=[_make_started_llm_stage(7)]) as mock_launch, + patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(11)) as mock_remote, + patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), + patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), + patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + ), + patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(cfg.stage_id), + ), + patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), + ): + engine._initialize_stages(stage_init_timeout=60) + + assert [call.args[1].stage_id for call in mock_launch.call_args_list] == [7] + assert [call.args[1].stage_id for call in mock_remote.call_args_list] == [11] + + def test_omni_master_server_preallocates_diffusion_stage_ids(self): + """Diffusion stages should also receive OmniMasterServer allocations.""" + stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11, stage_type="diffusion")] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=7) + mock_oms = Mock(spec=OmniMasterServer) + finalized = ( + [Mock(), Mock()], + [Mock(), Mock()], + [ + {"final_output": False, "final_output_type": None, "stage_type": "llm"}, + {"final_output": False, "final_output_type": None, "stage_type": "diffusion"}, + ], + ) + + with ( + patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(7)), + patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(7)), + patch.object(engine, "_launch_diffusion_stage", return_value=Mock()), + patch.object(engine, "_create_remote_diffusion_stage", return_value=Mock()), + patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms) as mock_oms_cls, + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), + patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), + patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + ), + patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(cfg.stage_id, getattr(cfg, "stage_type", "llm")), + ), + patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), + ): + engine._initialize_stages(stage_init_timeout=60) + + mock_oms_cls.assert_called_once_with( + master_address=engine._omni_master_address, + master_port=engine._omni_master_port, + stage_ids=[7, 11], + ) + + def test_duplicate_llm_stage_ids_raise(self): + """Duplicate configured LLM stage IDs should fail fast.""" + stage_cfgs = [_make_stage_cfg(3), _make_stage_cfg(3)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=3) + + with ( + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), + pytest.raises(ValueError, match="Duplicate stage_id"), + ): + engine._initialize_stages(stage_init_timeout=60) + + def test_omni_master_server_not_started_in_normal_mode(self): + """OmniMasterServer must NOT be instantiated outside single_stage_mode.""" + stage_cfgs = [_make_stage_cfg(0)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=False, stage_id_filter=None) + finalized = ([Mock()], [Mock()], [{"final_output": True, "final_output_type": None, "stage_type": "llm"}]) + + with ( + patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)), + patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.OmniMasterServer") as mock_oms_cls, + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), + patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), + patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + ), + patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(cfg.stage_id), + ), + patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), + ): + engine._initialize_stages(stage_init_timeout=60) + + mock_oms_cls.assert_not_called() + + def test_single_stage_mode_missing_master_address_raises(self): + """single_stage_mode without master address/port raises ValueError.""" + stage_cfgs = [_make_stage_cfg(0)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) + engine._omni_master_address = None # missing + engine._omni_master_port = None + + with ( + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), + pytest.raises(ValueError, match="omni_master_address"), + ): + engine._initialize_stages(stage_init_timeout=60) + + def test_matching_diffusion_stage_uses_local_registered_launch(self): + """A local diffusion stage should use the registered single-stage launch path.""" + stage_cfgs = [_make_stage_cfg(0, stage_type="diffusion"), _make_stage_cfg(1)] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) + mock_oms = Mock(spec=OmniMasterServer) + diffusion_client = Mock(stage_type="diffusion") + finalized = ( + [diffusion_client, Mock()], + [Mock(), Mock()], + [ + {"final_output": False, "final_output_type": None, "stage_type": "diffusion"}, + {"final_output": False, "final_output_type": None, "stage_type": "llm"}, + ], + ) + + with ( + patch.object(engine, "_launch_diffusion_stage", return_value=diffusion_client) as mock_local_diff, + patch.object(engine, "_create_remote_diffusion_stage") as mock_remote_diff, + patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(1)), + patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(1)), + patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), + patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), + patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + ), + patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(cfg.stage_id, getattr(cfg, "stage_type", "llm")), + ), + patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), + ): + engine._initialize_stages(stage_init_timeout=60) + + assert mock_local_diff.call_count == 1 + assert mock_local_diff.call_args.args[1].stage_id == 0 + mock_remote_diff.assert_not_called() + + def test_non_matching_diffusion_stage_uses_remote_diffusion_client(self): + """A non-local diffusion stage should attach via the remote diffusion path.""" + stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1, stage_type="diffusion")] + engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) + mock_oms = Mock(spec=OmniMasterServer) + remote_diffusion_client = Mock(stage_type="diffusion") + finalized = ( + [Mock(), remote_diffusion_client], + [Mock(), Mock()], + [ + {"final_output": False, "final_output_type": None, "stage_type": "llm"}, + {"final_output": False, "final_output_type": None, "stage_type": "diffusion"}, + ], + ) + + with ( + patch.object(engine, "_launch_diffusion_stage") as mock_local_diff, + patch.object( + engine, "_create_remote_diffusion_stage", return_value=remote_diffusion_client + ) as mock_remote_diff, + patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)), + patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(0)), + patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), + patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), + patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), + patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), + patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + ), + patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(cfg.stage_id, getattr(cfg, "stage_type", "llm")), + ), + patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), + ): + engine._initialize_stages(stage_init_timeout=60) + + mock_local_diff.assert_not_called() + assert mock_remote_diff.call_count == 1 + assert mock_remote_diff.call_args.args[0].stage_id == 1 + + +# --------------------------------------------------------------------------- +# AsyncOmniEngine – _launch_diffusion_stage +# --------------------------------------------------------------------------- + + +class TestLaunchDiffusionStage: + """Test local diffusion stage launch wiring.""" + + def test_registers_stage_with_public_master_properties(self): + engine = object.__new__(AsyncOmniEngine) + engine.model = "fake-model" + engine.diffusion_batch_size = 4 + + stage_cfg = _make_stage_cfg(5, stage_type="diffusion") + metadata = Mock(stage_id=5) + omni_master_server = Mock(spec=OmniMasterServer) + omni_master_server.address = "127.0.0.1" + omni_master_server.port = 25000 + + proc = Mock() + diffusion_client = Mock() + + with ( + patch("vllm_omni.engine.async_omni_engine.build_diffusion_config", return_value="diffusion-config"), + patch( + "vllm_omni.engine.async_omni_engine.register_stage_with_omni_master", + return_value=( + "tcp://127.0.0.1:25001", + "tcp://127.0.0.1:25002", + "tcp://127.0.0.1:25003", + ), + ) as mock_register, + patch( + "vllm_omni.engine.async_omni_engine.spawn_diffusion_proc", + return_value=(proc, None, None, None), + ) as mock_spawn, + patch("vllm_omni.engine.async_omni_engine.complete_diffusion_handshake") as mock_handshake, + patch( + "vllm_omni.engine.async_omni_engine.StageDiffusionClient.from_addresses", + return_value=diffusion_client, + ) as mock_from_addresses, + ): + result = engine._launch_diffusion_stage( + stage_cfg=stage_cfg, + metadata=metadata, + omni_master_server=omni_master_server, + ) + + mock_register.assert_called_once_with( + omni_master_address="127.0.0.1", + omni_master_port=25000, + omni_stage_id=5, + omni_stage_config=stage_cfg, + return_addresses=True, + ) + mock_spawn.assert_called_once_with( + "fake-model", + "diffusion-config", + handshake_address="tcp://127.0.0.1:25001", + request_address="tcp://127.0.0.1:25002", + response_address="tcp://127.0.0.1:25003", + ) + mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:25001") + mock_from_addresses.assert_called_once_with( + metadata, + request_address="tcp://127.0.0.1:25002", + response_address="tcp://127.0.0.1:25003", + proc=proc, + batch_size=4, + ) + assert result is diffusion_client + + +# --------------------------------------------------------------------------- +# AsyncOmniEngine – _create_remote_llm_stage +# --------------------------------------------------------------------------- + + +class TestCreateRemoteLlmStage: + """Test _create_remote_llm_stage delegates correctly.""" + + def _engine(self) -> AsyncOmniEngine: + engine = object.__new__(AsyncOmniEngine) + engine.model = "fake-model" + engine.single_stage_mode = True + engine._single_stage_id_filter = 0 + engine._omni_master_server = Mock(spec=OmniMasterServer) + engine._omni_master_server.get_zmq_addresses.return_value = Mock() + engine._omni_master_server.get_allocation.return_value = Mock() + engine._omni_master_server.get_stage_config.return_value = { + "stage_id": 0, + "stage_type": "llm", + "engine_args": {}, + } + return engine + + @contextmanager + def _patch_build_and_connect(self, stage_id: int): + fake_vllm_config = Mock() + fake_executor_cls = Mock() + fake_addresses = Mock() + fake_addresses.inputs = ["tcp://127.0.0.1:5000"] + fake_addresses.outputs = ["tcp://127.0.0.1:5001"] + fake_addresses.frontend_stats_publish_address = None + + eng_mgr = Mock() + coordinator = Mock() + + @contextmanager + def fake_connect_cm(*args, **kwargs): + yield eng_mgr, coordinator, fake_addresses + + with ( + patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": stage_id}, + ), + patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(fake_vllm_config, fake_executor_cls), + ), + patch( + "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", + return_value=fake_connect_cm(), + ) as mock_connect, + ): + yield mock_connect, fake_vllm_config, fake_executor_cls, fake_addresses + + def test_returns_started_llm_stage_with_correct_stage_id(self): + engine = self._engine() + stage_cfg = _make_stage_cfg(1) + metadata = Mock(stage_id=1) + omni_ms = engine._omni_master_server + omni_ms.get_stage_config.return_value = { + "stage_id": 1, + "stage_type": "llm", + "engine_args": {}, + } + + with self._patch_build_and_connect(1): + result = engine._create_remote_llm_stage( + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + omni_master_server=omni_ms, + ) + assert isinstance(result, StartedLlmStage) + assert result.stage_id == 1 + + def test_connect_remote_engine_cores_called_with_stage_id(self): + engine = self._engine() + stage_cfg = _make_stage_cfg(2) + metadata = Mock(stage_id=2) + omni_ms = engine._omni_master_server + omni_ms.get_zmq_addresses.return_value = Mock(inputs=["x"], outputs=["y"]) + omni_ms.get_stage_config.return_value = { + "stage_id": 2, + "stage_type": "llm", + "engine_args": {}, + } + + fake_vllm_config = Mock() + fake_executor_cls = Mock() + fake_addresses = Mock() + fake_addresses.inputs = ["tcp://127.0.0.1:5000"] + fake_addresses.outputs = ["tcp://127.0.0.1:5001"] + fake_addresses.frontend_stats_publish_address = None + + @contextmanager + def fake_connect_cm(*args, **kwargs): + yield Mock(), Mock(), fake_addresses + + with ( + patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 2}, + ), + patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(fake_vllm_config, fake_executor_cls), + ), + patch( + "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", return_value=fake_connect_cm() + ) as mock_connect, + ): + engine._create_remote_llm_stage( + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + omni_master_server=omni_ms, + ) + + mock_connect.assert_called_once() + _, kwargs = mock_connect.call_args + assert kwargs.get("stage_id") == 2 or mock_connect.call_args.args[-1] == 2 + omni_ms.get_stage_config.assert_called_once_with(2, timeout_s=60) + + def test_missing_registered_stage_config_raises_value_error(self): + engine = self._engine() + stage_cfg = _make_stage_cfg(3) + metadata = Mock(stage_id=3) + omni_ms = engine._omni_master_server + omni_ms.get_stage_config.return_value = None + + with patch("vllm_omni.engine.async_omni_engine.build_engine_args_dict") as mock_build_args: + with pytest.raises( + ValueError, + match="Remote stage 3 registered without stage config", + ): + engine._create_remote_llm_stage( + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + omni_master_server=omni_ms, + ) + + mock_build_args.assert_not_called() + + def test_exception_during_connect_closes_started_stage(self): + """If an error occurs after StartedLlmStage creation, close_started_llm_stage is called.""" + engine = self._engine() + stage_cfg = _make_stage_cfg(1) + metadata = Mock(stage_id=1) + omni_ms = engine._omni_master_server + omni_ms.get_stage_config.return_value = { + "stage_id": 1, + "stage_type": "llm", + "engine_args": {}, + } + + @contextmanager + def boom(*args, **kwargs): + yield Mock(), Mock(), Mock() + raise RuntimeError("handshake failed") + + with ( + patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 1}, + ), + patch("vllm_omni.engine.async_omni_engine.build_vllm_config", return_value=(Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", return_value=boom()), + patch("vllm_omni.engine.async_omni_engine.close_started_llm_stage") as mock_close, + ): + with pytest.raises(RuntimeError, match="handshake failed"): + engine._create_remote_llm_stage( + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + omni_master_server=omni_ms, + ) + mock_close.assert_called_once() + + +class TestConnectRemoteEngineCoresCoordinator: + """Test coordinator launch parity with launch_core_engines.""" + + @staticmethod + def _build_vllm_config(*, dp_rank: int = 0, offline_mode: bool = False, needs_dp_coordinator: bool = True) -> Mock: + parallel_config = Mock() + parallel_config.data_parallel_size_local = 1 + parallel_config.data_parallel_size = 2 + parallel_config.data_parallel_rank = dp_rank + parallel_config.data_parallel_rank_local = 0 if offline_mode else None + + vllm_config = Mock() + vllm_config.parallel_config = parallel_config + vllm_config.needs_dp_coordinator = needs_dp_coordinator + vllm_config.model_config = Mock(is_moe=False) + return vllm_config + + def test_uses_registered_coordinator_addresses(self): + vllm_config = self._build_vllm_config(dp_rank=0, offline_mode=False, needs_dp_coordinator=True) + + omni_master_server = 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 = Mock(handshake_bind_address="tcp://127.0.0.1:26001") + omni_master_server.get_stage_coordinator_addresses.return_value = StageCoordinatorAddresses( + coordinator_input="tcp://coord-in", + coordinator_output="tcp://coord-out", + frontend_stats_publish_address="tcp://stats", + ) + + @contextmanager + def fake_socket_ctx(*args, **kwargs): + yield Mock() + + with ( + patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()), + patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") as mock_wait, + ): + with connect_remote_engine_cores( + vllm_config=vllm_config, + omni_master_server=omni_master_server, + stage_id=7, + ) as (_, yielded_coordinator, yielded_addresses): + assert yielded_coordinator is 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" + + omni_master_server.get_stage_coordinator_addresses.assert_called_once_with(7) + mock_wait.assert_called_once() + + def test_defaults_to_no_coordinator_addresses_when_none_registered(self): + vllm_config = self._build_vllm_config( + dp_rank=0, + offline_mode=False, + needs_dp_coordinator=True, + ) + + omni_master_server = 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 = Mock(handshake_bind_address="tcp://127.0.0.1:26001") + omni_master_server.get_stage_coordinator_addresses.return_value = StageCoordinatorAddresses() + + @contextmanager + def fake_socket_ctx(*args, **kwargs): + yield Mock() + + with ( + patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()), + patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_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): + assert yielded_coordinator is None + assert yielded_addresses.coordinator_input is None + assert yielded_addresses.coordinator_output is None + assert yielded_addresses.frontend_stats_publish_address is None + + +class TestLaunchOmniCoreEngines: + """Tests for local omni engine launch wiring.""" + + def test_registers_stage_once_and_reuses_handshake_for_all_local_engines(self): + parallel_config = Mock( + data_parallel_size_local=2, + data_parallel_size=4, + data_parallel_rank=3, + ) + vllm_config = Mock(parallel_config=parallel_config) + + omni_master_server = Mock(spec=OmniMasterServer) + omni_master_server.address = "127.0.0.1" + omni_master_server.port = 26000 + omni_master_server.get_allocation.return_value = Mock(handshake_bind_address="tcp://127.0.0.1:26001") + + stage_config = {"stage_id": 7, "stage_type": "llm"} + local_engine_manager = Mock() + + @contextmanager + def fake_socket_ctx(*args, **kwargs): + yield Mock() + + with ( + patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value="tcp://127.0.0.1:26001", + ) as mock_register, + patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()), + patch( + "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", + return_value=local_engine_manager, + ) as mock_manager_cls, + patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup"), + ): + with launch_omni_core_engines( + vllm_config=vllm_config, + executor_class=Mock(), + log_stats=False, + omni_master_server=omni_master_server, + stage_id=7, + stage_config=stage_config, + ) as (yielded_manager, yielded_coordinator, yielded_addresses): + assert yielded_manager is local_engine_manager + assert yielded_coordinator is None + + mock_register.assert_called_once_with( + omni_master_address="127.0.0.1", + omni_master_port=26000, + omni_stage_id=7, + omni_stage_config=stage_config, + coordinator=None, + ) + mock_manager_cls.assert_called_once() + manager_kwargs = mock_manager_cls.call_args.kwargs + assert manager_kwargs["local_engine_count"] == 2 + assert manager_kwargs["start_index"] == 3 + assert manager_kwargs["local_start_index"] == 0 + assert manager_kwargs["vllm_config"] is vllm_config + assert manager_kwargs["local_client"] is True + assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" + assert manager_kwargs["executor_class"] is not None + + def test_registers_stage_with_coordinator_when_started(self): + parallel_config = Mock( + data_parallel_size_local=1, + data_parallel_size=2, + data_parallel_rank=0, + ) + vllm_config = Mock(parallel_config=parallel_config) + vllm_config.needs_dp_coordinator = True + vllm_config.model_config = Mock(is_moe=False) + + omni_master_server = Mock(spec=OmniMasterServer) + omni_master_server.address = "127.0.0.1" + omni_master_server.port = 26000 + omni_master_server.get_zmq_addresses.return_value = EngineZmqAddresses( + inputs=["tcp://client-in"], outputs=["tcp://client-out"] + ) + omni_master_server.get_allocation.return_value = Mock(handshake_bind_address="tcp://127.0.0.1:26001") + + coordinator = Mock() + coordinator.proc.pid = 1234 + coordinator.get_engine_socket_addresses.return_value = ("tcp://coord-in", "tcp://coord-out") + coordinator.get_stats_publish_address.return_value = "tcp://stats" + + @contextmanager + def fake_socket_ctx(*args, **kwargs): + yield Mock() + + with ( + patch("vllm_omni.engine.stage_engine_startup.DPCoordinator", return_value=coordinator), + patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value="tcp://127.0.0.1:26001", + ) as mock_register, + patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()), + patch( + "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", + return_value=Mock(), + ) as mock_manager_cls, + patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") as mock_wait, + ): + with launch_omni_core_engines( + vllm_config=vllm_config, + executor_class=Mock(), + log_stats=False, + omni_master_server=omni_master_server, + stage_id=7, + stage_config={"stage_id": 7}, + ): + pass + + mock_register.assert_called_once_with( + omni_master_address="127.0.0.1", + omni_master_port=26000, + omni_stage_id=7, + omni_stage_config={"stage_id": 7}, + coordinator=coordinator, + ) + manager_kwargs = mock_manager_cls.call_args.kwargs + assert manager_kwargs["log_stats"] is False + mock_wait.assert_called_once() + + +# --------------------------------------------------------------------------- +# AsyncOmniEngine – _launch_llm_stage single_stage_mode codepath +# --------------------------------------------------------------------------- + + +class TestLaunchLlmStageSingleStageMode: + """Test that _launch_llm_stage selects launch_omni_core_engines when + single_stage_mode=True and _omni_master_server is set.""" + + def _build_engine_with_oms(self) -> AsyncOmniEngine: + engine = object.__new__(AsyncOmniEngine) + engine.model = "fake-model" + engine.single_stage_mode = True + engine._single_stage_id_filter = 0 + engine._llm_stage_launch_lock = threading.Lock() + mock_oms = Mock(spec=OmniMasterServer) + mock_oms.address = "127.0.0.1" + mock_oms.port = 25000 + alloc = Mock() + alloc.handshake_bind_address = "tcp://127.0.0.1:25001" + mock_oms.get_allocation.return_value = alloc + fake_addresses = Mock() + fake_addresses.inputs = ["tcp://127.0.0.1:5000"] + fake_addresses.outputs = ["tcp://127.0.0.1:5001"] + fake_addresses.frontend_stats_publish_address = None + mock_oms.get_zmq_addresses.return_value = fake_addresses + engine._omni_master_server = mock_oms + return engine + + @contextmanager + def _patch_launch_omni_cm(self, stage_id: int): + fake_vllm_config = Mock() + fake_executor_cls = Mock() + fake_addresses = Mock() + fake_addresses.inputs = ["tcp://127.0.0.1:5000"] + fake_addresses.outputs = ["tcp://127.0.0.1:5001"] + fake_addresses.frontend_stats_publish_address = None + + eng_mgr = Mock() + + @contextmanager + def fake_launch_omni(*args, **kwargs): + yield eng_mgr, None, fake_addresses + + with ( + patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), + patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": stage_id}, + ), + patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(fake_vllm_config, fake_executor_cls), + ), + patch( + "vllm_omni.engine.async_omni_engine.acquire_device_locks", + return_value=[], + ), + patch( + "vllm_omni.engine.async_omni_engine.release_device_locks", + ), + patch( + "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", + return_value=fake_launch_omni(), + ) as mock_launch_omni, + ): + yield mock_launch_omni + + def test_launch_omni_core_engines_used_in_single_stage_mode(self): + """single_stage_mode + _omni_master_server → launch_omni_core_engines.""" + engine = self._build_engine_with_oms() + metadata = Mock(stage_id=0, runtime_cfg={}) + stage_cfg = _make_stage_cfg(0) + + with self._patch_launch_omni_cm(0) as mock_launch_omni: + result = engine._launch_llm_stage( + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + llm_stage_launch_lock=threading.Lock(), + ) + + mock_launch_omni.assert_called_once() + assert mock_launch_omni.call_args.kwargs["stage_config"] is stage_cfg + assert isinstance(result, StartedLlmStage) + assert result.stage_id == 0 + + def test_spawn_stage_core_used_in_normal_mode(self): + """~single_stage_mode → spawn_stage_core + complete_stage_handshake.""" + engine = object.__new__(AsyncOmniEngine) + engine.model = "fake-model" + engine.single_stage_mode = False + engine._omni_master_server = None + engine._llm_stage_launch_lock = threading.Lock() + + fake_vllm_config = Mock() + fake_executor_cls = Mock() + fake_addresses = Mock() + fake_addresses.inputs = ["tcp://127.0.0.1:5000"] + fake_addresses.outputs = ["tcp://127.0.0.1:5001"] + fake_addresses.frontend_stats_publish_address = None + + fake_proc = Mock() + fake_handshake_address = "ipc:///tmp/fake-handshake" + + with ( + patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), + patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 0}, + ), + patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(fake_vllm_config, fake_executor_cls), + ), + patch("vllm_omni.engine.async_omni_engine.acquire_device_locks", return_value=[]), + patch("vllm_omni.engine.async_omni_engine.release_device_locks"), + patch( + "vllm_omni.engine.async_omni_engine.spawn_stage_core", + return_value=(fake_addresses, fake_proc, fake_handshake_address), + ) as mock_spawn, + patch("vllm_omni.engine.async_omni_engine.complete_stage_handshake") as mock_handshake, + patch("vllm_omni.engine.async_omni_engine.launch_omni_core_engines") as mock_omni, + ): + metadata = Mock(stage_id=0, runtime_cfg={}) + result = engine._launch_llm_stage( + stage_cfg=_make_stage_cfg(0), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + llm_stage_launch_lock=threading.Lock(), + ) + + mock_spawn.assert_called_once_with( + vllm_config=fake_vllm_config, + executor_class=fake_executor_cls, + log_stats=False, + ) + mock_handshake.assert_called_once_with( + fake_proc, + fake_handshake_address, + fake_addresses, + fake_vllm_config, + ) + mock_omni.assert_not_called() + assert isinstance(result, StartedLlmStage) + assert result.proc is fake_proc + + def test_launch_omni_passes_stage_id_and_master_server(self): + """launch_omni_core_engines receives the correct stage_id and omni_master_server.""" + engine = self._build_engine_with_oms() + metadata = Mock(stage_id=0, runtime_cfg={}) + + captured_kwargs: dict[str, Any] = {} + + @contextmanager + def capturing_launch(*args, **kwargs): + captured_kwargs.update(kwargs) + fake_addresses = Mock() + fake_addresses.inputs = ["tcp://127.0.0.1:5000"] + fake_addresses.outputs = ["tcp://127.0.0.1:5001"] + fake_addresses.frontend_stats_publish_address = None + yield Mock(), None, fake_addresses + + with ( + patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), + patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 0}, + ), + patch("vllm_omni.engine.async_omni_engine.build_vllm_config", return_value=(Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.acquire_device_locks", return_value=[]), + patch("vllm_omni.engine.async_omni_engine.release_device_locks"), + patch("vllm_omni.engine.async_omni_engine.launch_omni_core_engines", side_effect=capturing_launch), + ): + engine._launch_llm_stage( + stage_cfg=_make_stage_cfg(0), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + llm_stage_launch_lock=threading.Lock(), + ) + + assert captured_kwargs.get("stage_id") == 0 + assert captured_kwargs.get("omni_master_server") is engine._omni_master_server + + def test_launch_omni_context_exits_before_stage_cleanup_on_error(self): + """Errors after entering the omni launch context still unwind it first.""" + engine = self._build_engine_with_oms() + metadata = Mock(stage_id=0, runtime_cfg={}) + + fake_addresses = Mock() + fake_addresses.inputs = ["tcp://127.0.0.1:5000"] + fake_addresses.outputs = ["tcp://127.0.0.1:5001"] + fake_addresses.frontend_stats_publish_address = None + + events: list[str] = [] + + @contextmanager + def fake_launch_omni(*args, **kwargs): + try: + yield Mock(), None, fake_addresses + finally: + events.append("launch_exit") + + with ( + patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), + patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 0}, + ), + patch("vllm_omni.engine.async_omni_engine.build_vllm_config", return_value=(Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.acquire_device_locks", return_value=[]), + patch("vllm_omni.engine.async_omni_engine.release_device_locks"), + patch( + "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", + return_value=fake_launch_omni(), + ), + patch("vllm_omni.engine.async_omni_engine.logger.info", side_effect=RuntimeError("boom")), + patch( + "vllm_omni.engine.async_omni_engine.close_started_llm_stage", + side_effect=lambda _started: events.append("stage_close"), + ) as mock_close_stage, + ): + with pytest.raises(RuntimeError, match="boom"): + engine._launch_llm_stage( + stage_cfg=_make_stage_cfg(0), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + llm_stage_launch_lock=threading.Lock(), + ) + + mock_close_stage.assert_called_once() + assert events == ["launch_exit", "stage_close"] + + def test_base_exception_propagates_without_started_stage_cleanup(self): + """BaseException subclasses should bypass the Exception cleanup path.""" + engine = self._build_engine_with_oms() + metadata = Mock(stage_id=0, runtime_cfg={}) + + fake_addresses = Mock() + fake_addresses.inputs = ["tcp://127.0.0.1:5000"] + fake_addresses.outputs = ["tcp://127.0.0.1:5001"] + fake_addresses.frontend_stats_publish_address = None + + events: list[str] = [] + + class FatalLaunchInterrupt(BaseException): + pass + + @contextmanager + def fake_launch_omni(*args, **kwargs): + try: + yield Mock(), None, fake_addresses + finally: + events.append("launch_exit") + + with ( + patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), + patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 0}, + ), + patch("vllm_omni.engine.async_omni_engine.build_vllm_config", return_value=(Mock(), Mock())), + patch("vllm_omni.engine.async_omni_engine.acquire_device_locks", return_value=[]), + patch("vllm_omni.engine.async_omni_engine.release_device_locks"), + patch( + "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", + return_value=fake_launch_omni(), + ), + patch( + "vllm_omni.engine.async_omni_engine.logger.info", + side_effect=FatalLaunchInterrupt("stop"), + ), + patch("vllm_omni.engine.async_omni_engine.close_started_llm_stage") as mock_close_stage, + ): + with pytest.raises(FatalLaunchInterrupt, match="stop"): + engine._launch_llm_stage( + stage_cfg=_make_stage_cfg(0), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + llm_stage_launch_lock=threading.Lock(), + ) + + mock_close_stage.assert_not_called() + assert events == ["launch_exit"] diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py new file mode 100644 index 000000000..916db3cc2 --- /dev/null +++ b/tests/entrypoints/test_serve.py @@ -0,0 +1,195 @@ +"""Unit tests for the Omni serve CLI helpers.""" + +from __future__ import annotations + +import argparse +from unittest.mock import Mock, patch + +import pytest + +from vllm_omni.entrypoints.cli.serve import run_headless + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_headless_args() -> argparse.Namespace: + return argparse.Namespace( + model="fake-model", + stage_id=3, + omni_master_address="127.0.0.1", + omni_master_port=26000, + api_server_count=0, + worker_backend="multi_process", + stage_configs_path=None, + log_stats=False, + disable_log_stats=False, + ) + + +def test_run_headless_registers_stage_once_and_launches_all_local_engines() -> None: + args = _make_headless_args() + stage_cfg = Mock(stage_id=3) + stage_cfgs = [stage_cfg] + parallel_config = Mock( + data_parallel_size_local=2, + data_parallel_rank=4, + data_parallel_rank_local=1, + node_rank_within_dp=0, + ) + vllm_config = Mock(parallel_config=parallel_config) + executor_class = Mock() + engine_manager = Mock() + + with ( + patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", stage_cfgs), + ), + patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"), + patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()), + patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}), + patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}), + patch( + "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ), + patch( + "vllm_omni.engine.stage_init_utils.build_vllm_config", + return_value=(vllm_config, executor_class), + ) as mock_build_vllm_config, + patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value="tcp://127.0.0.1:26001", + ) as mock_register, + patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls, + patch("signal.signal"), + ): + run_headless(args) + + mock_build_vllm_config.assert_called_once_with( + stage_cfg, + "fake-model", + stage_connector_spec={}, + engine_args_dict={}, + headless=True, + ) + mock_register.assert_called_once_with( + omni_master_address="127.0.0.1", + omni_master_port=26000, + omni_stage_id=3, + omni_stage_config=stage_cfg, + coordinator=None, + ) + mock_manager_cls.assert_called_once() + manager_kwargs = mock_manager_cls.call_args.kwargs + assert manager_kwargs["local_engine_count"] == 2 + assert manager_kwargs["start_index"] == 4 + assert manager_kwargs["local_start_index"] == 0 + assert manager_kwargs["local_client"] is False + assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" + assert manager_kwargs["log_stats"] is False + engine_manager.join_first.assert_called_once_with() + engine_manager.shutdown.assert_called_once_with() + + +def test_run_headless_honors_explicit_log_stats_flag() -> None: + args = _make_headless_args() + args.log_stats = True + stage_cfg = Mock(stage_id=3) + stage_cfgs = [stage_cfg] + parallel_config = Mock( + data_parallel_size_local=2, + data_parallel_rank=4, + data_parallel_rank_local=1, + node_rank_within_dp=0, + ) + vllm_config = Mock(parallel_config=parallel_config) + executor_class = Mock() + engine_manager = Mock() + + with ( + patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", stage_cfgs), + ), + patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"), + patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()), + patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}), + patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}), + patch( + "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ), + patch( + "vllm_omni.engine.stage_init_utils.build_vllm_config", + return_value=(vllm_config, executor_class), + ), + patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value="tcp://127.0.0.1:26001", + ), + patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls, + patch("signal.signal"), + ): + run_headless(args) + + manager_kwargs = mock_manager_cls.call_args.kwargs + assert manager_kwargs["log_stats"] is True + + +def test_run_headless_launches_diffusion_stage_via_omni_master() -> None: + args = _make_headless_args() + stage_cfg = Mock(stage_id=3, stage_type="diffusion") + stage_cfg.engine_args = Mock() + stage_cfg.engine_input_source = [] + stage_cfgs = [stage_cfg] + metadata = Mock(stage_id=3) + od_config = Mock() + proc = Mock() + proc.exitcode = 0 + proc.is_alive.return_value = False + + with ( + patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", stage_cfgs), + ), + patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"), + patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()), + patch( + "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ), + patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata), + patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info") as mock_inject_stage_info, + patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config), + patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value=("tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"), + ) as mock_register, + patch( + "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc", + return_value=(proc, "tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"), + ) as mock_spawn, + patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") as mock_handshake, + patch("signal.signal"), + ): + run_headless(args) + + mock_inject_stage_info.assert_called_once_with(stage_cfg, 3) + mock_register.assert_called_once_with( + omni_master_address="127.0.0.1", + omni_master_port=26000, + omni_stage_id=3, + omni_stage_config=stage_cfg, + return_addresses=True, + ) + mock_spawn.assert_called_once_with( + "fake-model", + od_config, + handshake_address="tcp://127.0.0.1:26001", + request_address="tcp://127.0.0.1:26002", + response_address="tcp://127.0.0.1:26003", + ) + mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001") + proc.join.assert_called_once_with() diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 6e44fe533..94e254c25 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -8,6 +8,7 @@ import pytest import torch from pytest_mock import MockerFixture +from vllm_omni.config.yaml_util import create_config from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.engine.async_omni_engine import AsyncOmniEngine @@ -16,6 +17,7 @@ from vllm_omni.entrypoints.utils import ( _filter_dict_like_object, filter_dataclass_kwargs, load_and_resolve_stage_configs, + load_stage_configs_from_yaml, resolve_model_config_path, ) @@ -322,3 +324,70 @@ class TestLoadAndResolveStageConfigs: assert config_path is None assert len(stage_configs) == 1 assert "dtype" in stage_configs[0]["engine_args"] + + +class TestLoadStageConfigsFromYaml: + """Regression tests for stage-config loading and merging.""" + + def test_deep_merges_stage_engine_args(self, mocker: MockerFixture): + yaml_config = create_config( + { + "async_chunk": True, + "stage_args": [ + { + "stage_id": 0, + "runtime": {"device": 0}, + "engine_args": { + "parallel_config": {"tensor_parallel_size": 4}, + }, + } + ], + } + ) + mocker.patch( + "vllm_omni.entrypoints.utils.load_yaml_config", + return_value=yaml_config, + ) + + stages = load_stage_configs_from_yaml( + "fake.yaml", + base_engine_args={ + "parallel_config": { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 2, + }, + "model": "base-model", + }, + ) + + merged_engine_args = stages[0]["engine_args"] + assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4 + assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2 + assert merged_engine_args["model"] == "base-model" + assert merged_engine_args["async_chunk"] is True + + def test_merges_nested_stage_engine_args(self, mocker: MockerFixture): + yaml_config = create_config( + { + "stage_args": [ + { + "stage_id": 0, + "engine_args": { + "nested": {"override": 2}, + }, + } + ], + } + ) + mocker.patch( + "vllm_omni.entrypoints.utils.load_yaml_config", + return_value=yaml_config, + ) + + stages = load_stage_configs_from_yaml( + "fake.yaml", + base_engine_args={"nested": {"base": 1}}, + ) + + assert stages[0]["engine_args"]["nested"]["base"] == 1 + assert stages[0]["engine_args"]["nested"]["override"] == 2 diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index a1a4766de..cd7159b68 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -51,6 +51,41 @@ class StageDiffusionClient: od_config: OmniDiffusionConfig, metadata: StageMetadata, 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) + self._initialize_client(metadata, request_address, response_address, proc=proc, batch_size=batch_size) + + @classmethod + def from_addresses( + cls, + metadata: StageMetadata, + request_address: str, + response_address: str, + *, + proc: Any = None, + batch_size: int = 1, + ) -> StageDiffusionClient: + """Create a client for an already-running diffusion subprocess.""" + client = cls.__new__(cls) + client._initialize_client( + metadata, + request_address, + response_address, + proc=proc, + batch_size=batch_size, + ) + return client + + def _initialize_client( + self, + metadata: StageMetadata, + request_address: str, + response_address: str, + *, + proc: Any, + batch_size: int, ) -> None: self.stage_id = metadata.stage_id self.final_output = metadata.final_output @@ -58,13 +93,9 @@ class StageDiffusionClient: self.default_sampling_params = metadata.default_sampling_params self.custom_process_input_func = metadata.custom_process_input_func self.engine_input_source = metadata.engine_input_source - - # 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) self._proc = proc + self._owns_process = proc is not None - # ZMQ sockets (sync) for communicating with the subprocess. self._zmq_ctx = zmq.Context() self._request_socket = self._zmq_ctx.socket(zmq.PUSH) self._request_socket.connect(request_address) @@ -74,14 +105,18 @@ class StageDiffusionClient: self._encoder = OmniMsgpackEncoder() self._decoder = OmniMsgpackDecoder() - # Buffers for demultiplexing response messages. self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue() self._rpc_results: dict[str, Any] = {} self._pending_rpcs: set[str] = set() self._tasks: dict[str, asyncio.Task] = {} self._shutting_down = False - logger.info("[StageDiffusionClient] Stage-%s initialized (batch_size=%d)", self.stage_id, batch_size) + logger.info( + "[StageDiffusionClient] Stage-%s initialized (owns_process=%s, batch_size=%d)", + self.stage_id, + self._owns_process, + batch_size, + ) # ------------------------------------------------------------------ # Internal helpers @@ -253,7 +288,7 @@ class StageDiffusionClient: try: return self._output_queue.get_nowait() except asyncio.QueueEmpty: - if not self._shutting_down and self._proc is not None and not self._proc.is_alive(): + if not self._shutting_down and self._owns_process and self._proc is not None and not self._proc.is_alive(): exitcode = self._proc.exitcode # One final drain – the last ZMQ frame may have arrived # between the first drain and the is_alive() check. @@ -325,7 +360,7 @@ class StageDiffusionClient: self._drain_responses() if rpc_id in self._rpc_results: return self._rpc_results.pop(rpc_id) - if self._proc is not None and not self._proc.is_alive(): + if self._owns_process and self._proc is not None and not self._proc.is_alive(): raise RuntimeError( f"StageDiffusionProc died while waiting for " f"collective_rpc '{method}' (exit code {self._proc.exitcode})" @@ -343,7 +378,7 @@ class StageDiffusionClient: except Exception: pass - if self._proc is not None and self._proc.is_alive(): + 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) diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 9d8c06cce..2bba41925 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -580,14 +580,17 @@ class StageDiffusionProc: def spawn_diffusion_proc( model: str, od_config: OmniDiffusionConfig, + handshake_address: str | None = None, + request_address: str | None = None, + response_address: str | None = None, ) -> tuple[BaseProcess, str, str, str]: """Spawn a StageDiffusionProc subprocess. Returns ``(proc, handshake_address, request_address, response_address)``. """ - handshake_address = get_open_zmq_ipc_path() - request_address = get_open_zmq_ipc_path() - response_address = get_open_zmq_ipc_path() + 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( diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index b66378926..d43f1b8fd 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -86,7 +86,11 @@ class OmniEngineArgs(EngineArgs): Adds omni-specific configuration fields for multi-stage pipeline processing and output type specification. Args: - stage_id: Identifier for the stage in a multi-stage pipeline (default: 0) + stage_id: Identifier for the stage in a multi-stage pipeline. + Defaults to 0 for per-stage engine construction. The CLI-level + single-stage selector remains optional on the parsed argparse + namespace and should not be forwarded as a nullable per-stage + engine argument. model_stage: Stage type identifier, e.g., "thinker" or "talker" (default: "thinker") model_arch: Model architecture name @@ -105,6 +109,18 @@ class OmniEngineArgs(EngineArgs): worker_type: Model Type, e.g., "ar" or "generation" task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base). If not specified, will be inferred from model path. + omni_master_address: TCP address that the OmniMasterServer (running + inside AsyncOmniEngine) listens on for engine core registrations. + Required when single-stage mode is active. + omni_master_port: TCP port for the OmniMasterServer registration + socket. Required when single-stage mode is active. + stage_configs_path: Optional path to a JSON/YAML file containing + stage configurations for the multi-stage pipeline. If None, + stage configs are resolved from the model's default configuration. + output_modalities: Optional list of output modality names to enable + (e.g. ["text", "audio"]). If None, all modalities supported by + the model are used. + log_stats: Whether to log engine statistics. Defaults to False. """ stage_id: int = 0 @@ -119,6 +135,11 @@ class OmniEngineArgs(EngineArgs): quantization_config: Any | None = None worker_type: str | None = None task_type: str | None = None + omni_master_address: str | None = None + omni_master_port: int | None = None + stage_configs_path: str | None = None + output_modalities: list[str] | None = None + log_stats: bool = False def __post_init__(self) -> None: load_omni_general_plugins() diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index f7e7d53d5..7dc5db0ac 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -18,12 +18,10 @@ import time import uuid import weakref from collections.abc import Mapping, Sequence +from contextlib import ExitStack from dataclasses import asdict from typing import TYPE_CHECKING, Any -if TYPE_CHECKING: - from vllm_omni.engine.arg_utils import OmniEngineArgs - import janus import torch from omegaconf import OmegaConf @@ -34,26 +32,36 @@ from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.diffusion.data import DiffusionParallelConfig +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.engine import ( - OmniEngineCoreRequest, -) +from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.orchestrator import Orchestrator from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.engine.serialization import ( deserialize_additional_information, serialize_additional_information, ) -from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +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_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 ( StartedLlmStage, acquire_device_locks, + build_diffusion_config, build_engine_args_dict, build_vllm_config, cleanup_failed_stage_initialization, @@ -62,17 +70,20 @@ from vllm_omni.engine.stage_init_utils import ( finalize_initialized_stages, get_stage_connector_spec, initialize_diffusion_stage, + inject_kv_stage_info, load_omni_transfer_config_for_model, prepare_engine_environment, release_device_locks, setup_stage_devices, + terminate_alive_proc, ) -from vllm_omni.entrypoints.utils import ( - load_and_resolve_stage_configs, -) +from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs from vllm_omni.inputs.preprocess import OmniInputPreprocessor from vllm_omni.platforms import current_omni_platform +if TYPE_CHECKING: + from vllm_omni.engine.arg_utils import OmniEngineArgs + logger = init_logger(__name__) @@ -86,39 +97,6 @@ def _patch_generation_config_if_needed(model_config: Any) -> None: model_config.try_get_generation_config = lambda: {} -def _inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None: - """Inject stage_id and engine_input_source into omni_kv_config. - - OmniKVTransferManager needs stage_id to compute recv_stages for the - receiving side. In the old Omni architecture, OmniDiffusion.__init__ - performed this injection; replicate it here for AsyncOmniEngine. - """ - try: - engine_args = stage_cfg.engine_args - if hasattr(engine_args, "get"): - omni_kv = engine_args.get("omni_kv_config", None) - else: - omni_kv = getattr(engine_args, "omni_kv_config", None) - - if omni_kv is None: - return - - if hasattr(omni_kv, "setdefault"): - omni_kv.setdefault("stage_id", stage_id) - elif hasattr(omni_kv, "__setitem__"): - if "stage_id" not in omni_kv: - omni_kv["stage_id"] = stage_id - - engine_input_source = getattr(stage_cfg, "engine_input_source", None) - if engine_input_source is not None: - if hasattr(omni_kv, "setdefault"): - omni_kv.setdefault("engine_input_source", list(engine_input_source)) - elif hasattr(omni_kv, "__setitem__") and "engine_input_source" not in omni_kv: - omni_kv["engine_input_source"] = list(engine_input_source) - except Exception as e: - logger.debug("Failed to inject stage info into omni_kv_config: %s", e) - - def _inject_global_id(target: Any, request_id: str) -> None: """Inject global_request_id into a prompt dict's additional_information.""" if isinstance(target, dict): @@ -255,6 +233,7 @@ class AsyncOmniEngine: stage_init_timeout: int = 300, init_timeout: int = 600, diffusion_batch_size: int = 1, + single_stage_mode: bool = False, **kwargs: Any, ) -> None: self.model = model @@ -274,6 +253,31 @@ class AsyncOmniEngine: ea_dict.pop("model", None) kwargs = {**ea_dict, **kwargs} + # ------------------------------------------------------------------ # + # Single-stage mode detection # + # ------------------------------------------------------------------ # + # Single-stage mode is enabled when the caller explicitly passes # + # single_stage_mode=True, or when a stage_id is provided in the args. # + _stage_id_kwarg = kwargs.get("stage_id") + if isinstance(_stage_id_kwarg, int) and not single_stage_mode: + single_stage_mode = True + + self.single_stage_mode: bool = single_stage_mode + self._single_stage_id_filter: int | None = ( + int(_stage_id_kwarg) if single_stage_mode and isinstance(_stage_id_kwarg, int) else None + ) + 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 + + if single_stage_mode: + logger.info( + "[AsyncOmniEngine] Single-stage mode enabled (stage_id_filter=%s, master=%s:%s)", + self._single_stage_id_filter, + self._omni_master_address, + self._omni_master_port, + ) + self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs) self.num_stages = len(self.stage_configs) @@ -350,61 +354,89 @@ class AsyncOmniEngine: started_stage: StartedLlmStage | None = None lock_fds: list[int] = [] device_control_env = current_omni_platform.device_control_env_var - try: - with llm_stage_launch_lock: - previous_visible_devices = os.environ.get(device_control_env) - try: - setup_stage_devices(metadata.stage_id, metadata.runtime_cfg) - engine_args_dict = build_engine_args_dict( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - ) - 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", metadata.stage_id) - engine_args_dict["omni_kv_config"] = omni_kv - vllm_config, executor_class = build_vllm_config( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - engine_args_dict=engine_args_dict, - ) - lock_fds = acquire_device_locks( - metadata.stage_id, - engine_args_dict, - stage_init_timeout, - ) - addresses, proc, handshake_address = spawn_stage_core( - vllm_config=vllm_config, - executor_class=executor_class, - log_stats=False, - ) - started_stage = StartedLlmStage( - stage_id=metadata.stage_id, - metadata=metadata, - vllm_config=vllm_config, - executor_class=executor_class, - proc=proc, - addresses=addresses, - ) - logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id) - # Keep the stage-specific device visibility until vLLM - # finishes starting all child processes. - complete_stage_handshake(proc, handshake_address, addresses, vllm_config) - logger.info("[AsyncOmniEngine] Stage %s engine startup completed", 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) + proc = None + handshake_address = None + with ExitStack() as launch_stack: + with llm_stage_launch_lock: + previous_visible_devices = os.environ.get(device_control_env) + try: + setup_stage_devices(metadata.stage_id, metadata.runtime_cfg) + engine_args_dict = build_engine_args_dict( + stage_cfg, + self.model, + stage_connector_spec=stage_connector_spec, + ) + 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", metadata.stage_id) + engine_args_dict["omni_kv_config"] = omni_kv + vllm_config, executor_class = build_vllm_config( + stage_cfg, + self.model, + stage_connector_spec=stage_connector_spec, + engine_args_dict=engine_args_dict, + ) + lock_fds = acquire_device_locks( + metadata.stage_id, + engine_args_dict, + stage_init_timeout, + ) + if self.single_stage_mode and self._omni_master_server is not None: + engine_manager, coordinator, addresses = launch_stack.enter_context( + launch_omni_core_engines( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=False, + omni_master_server=self._omni_master_server, + stage_id=metadata.stage_id, + stage_config=stage_cfg, + ) + ) + started_stage = StartedLlmStage( + stage_id=metadata.stage_id, + metadata=metadata, + vllm_config=vllm_config, + executor_class=executor_class, + addresses=addresses, + engine_manager=engine_manager, + coordinator=coordinator, + ) + else: + addresses, proc, handshake_address = spawn_stage_core( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=False, + ) + started_stage = StartedLlmStage( + stage_id=metadata.stage_id, + metadata=metadata, + vllm_config=vllm_config, + executor_class=executor_class, + addresses=addresses, + proc=proc, + ) + logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id) + # Keep the stage-specific device visibility until vLLM + # finishes starting all child processes. + 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) + logger.info("[AsyncOmniEngine] Stage %s engine startup completed", 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) assert started_stage is not None return started_stage @@ -416,13 +448,138 @@ class AsyncOmniEngine: if lock_fds: release_device_locks(lock_fds) + def _create_remote_llm_stage( + self, + stage_cfg: Any, + metadata: Any, + stage_connector_spec: dict[str, Any], + stage_init_timeout: int, + omni_master_server: OmniMasterServer, + ) -> StartedLlmStage: + """Attach to a remote engine core and wait for its startup handshake.""" + started_stage: StartedLlmStage | None = None + try: + raw_stage_cfg = omni_master_server.get_stage_config( + metadata.stage_id, + timeout_s=stage_init_timeout, + ) + if raw_stage_cfg is None: + raise ValueError(f"Remote stage {metadata.stage_id} registered without stage config") + stage_cfg = OmegaConf.create(raw_stage_cfg) + engine_args_dict = build_engine_args_dict( + stage_cfg, + self.model, + stage_connector_spec=stage_connector_spec, + ) + vllm_config, executor_class = build_vllm_config( + stage_cfg, + self.model, + stage_connector_spec=stage_connector_spec, + engine_args_dict=engine_args_dict, + ) + vllm_config.parallel_config.data_parallel_size_local = 0 + launch_cm = connect_remote_engine_cores( + vllm_config=vllm_config, + omni_master_server=omni_master_server, + stage_id=metadata.stage_id, + ) + logger.info("[AsyncOmniEngine] Stage %s remote engine handshake started", metadata.stage_id) + with launch_cm as (engine_manager, coordinator, addresses): + started_stage = StartedLlmStage( + stage_id=metadata.stage_id, + metadata=metadata, + vllm_config=vllm_config, + executor_class=executor_class, + engine_manager=engine_manager, + coordinator=coordinator, + addresses=addresses, + ) + logger.info("[AsyncOmniEngine] Stage %s remote engine startup completed", metadata.stage_id) + assert started_stage is not None + return started_stage + except Exception: + if started_stage is not None: + close_started_llm_stage(started_stage) + raise + + def _launch_diffusion_stage( + self, + stage_cfg: Any, + metadata: Any, + omni_master_server: OmniMasterServer, + ) -> StageDiffusionClient: + """Launch a local diffusion stage on OmniMasterServer-allocated sockets.""" + proc = None + try: + od_config = build_diffusion_config(self.model, stage_cfg, metadata) + handshake_address, request_address, response_address = 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_cfg, + return_addresses=True, + ) + logger.info( + "[AsyncOmniEngine] Stage %s diffusion registration completed", + metadata.stage_id, + ) + proc, _, _, _ = spawn_diffusion_proc( + self.model, + od_config, + handshake_address=handshake_address, + request_address=request_address, + response_address=response_address, + ) + complete_diffusion_handshake(proc, handshake_address) + logger.info( + "[AsyncOmniEngine] Stage %s diffusion startup completed", + metadata.stage_id, + ) + return StageDiffusionClient.from_addresses( + metadata, + request_address=request_address, + response_address=response_address, + proc=proc, + batch_size=self.diffusion_batch_size, + ) + except Exception: + if proc is not None: + terminate_alive_proc(proc) + raise + + def _create_remote_diffusion_stage( + self, + metadata: Any, + stage_init_timeout: int, + omni_master_server: OmniMasterServer, + ) -> StageDiffusionClient: + """Attach to a remote diffusion stage registered with OmniMasterServer.""" + remote_stage_cfg = OmegaConf.create( + omni_master_server.get_stage_config( + metadata.stage_id, + timeout_s=stage_init_timeout, + ) + ) + remote_metadata = extract_stage_metadata(remote_stage_cfg) + addresses = omni_master_server.get_zmq_addresses(metadata.stage_id) + logger.info( + "[AsyncOmniEngine] Stage %s remote diffusion startup completed", + metadata.stage_id, + ) + return StageDiffusionClient.from_addresses( + remote_metadata, + request_address=addresses.inputs[0], + response_address=addresses.outputs[0], + batch_size=self.diffusion_batch_size, + ) + def _attach_llm_stage( self, started: StartedLlmStage, ) -> tuple[Any, Any, Any, InputProcessor | None]: """Attach a READY LLM stage to the orchestrator event loop.""" - client_addresses = { + client_addresses: dict[str, str] = { "input_address": started.addresses.inputs[0], "output_address": started.addresses.outputs[0], } @@ -430,14 +587,18 @@ class AsyncOmniEngine: client_addresses["stats_update_address"] = started.addresses.frontend_stats_publish_address try: - stage_client = StageEngineCoreClient( + stage_client = StageEngineCoreClientBase.make_async_mp_client( vllm_config=started.vllm_config, executor_class=started.executor_class, metadata=started.metadata, client_addresses=client_addresses, proc=started.proc, + engine_manager=started.engine_manager, + coordinator=started.coordinator, ) started.proc = None + started.engine_manager = None + started.coordinator = None except Exception: close_started_llm_stage(started) raise @@ -493,7 +654,7 @@ class AsyncOmniEngine: output_processors: list[Any | None] = [None] * num_stages stage_vllm_configs: list[Any | None] = [None] * num_stages input_processor: InputProcessor | None = None - llm_stage_ids: list[int] = [] + llm_stage_positions: list[int] = [] llm_launch_futures: dict[int, concurrent.futures.Future[StartedLlmStage]] = {} started_llm_stages: dict[int, StartedLlmStage] = {} llm_stage_launch_lock = threading.Lock() @@ -507,45 +668,102 @@ class AsyncOmniEngine: prepare_engine_environment() omni_transfer_config = load_omni_transfer_config_for_model(self.model, self.config_path) + # ------------------------------------------------------------------ # + # Single-stage mode: start OmniMasterServer before launching stages. # + # ------------------------------------------------------------------ # + if self.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." + ) + # Collect all configured stage IDs for pre-allocation. + all_stage_ids: list[int] = [] + seen_stage_ids: set[int] = set() + for i, sc in enumerate(self.stage_configs): + stage_id = int(getattr(sc, "stage_id", i)) + 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) + self._omni_master_server = OmniMasterServer( + master_address=self._omni_master_address, + master_port=self._omni_master_port, + stage_ids=all_stage_ids, + ) + self._omni_master_server.start() + logger.info( + "[AsyncOmniEngine] OmniMasterServer started for stages %s", + all_stage_ids, + ) + try: with concurrent.futures.ThreadPoolExecutor( max_workers=max(1, llm_stage_count), thread_name_prefix="llm-stage-launch", ) as launch_executor: - for stage_id, stage_cfg in enumerate(self.stage_configs): - logger.info("[AsyncOmniEngine] Initializing stage %s", stage_id) + for stage_idx, stage_cfg in enumerate(self.stage_configs): metadata = extract_stage_metadata(stage_cfg) + configured_stage_id = metadata.stage_id + logger.info("[AsyncOmniEngine] Initializing stage %s", configured_stage_id) if metadata.prompt_expand_func is not None: prompt_expand_func = metadata.prompt_expand_func + if self.single_stage_mode: + metadata.runtime_cfg = None + stage_connector_spec = get_stage_connector_spec( omni_transfer_config=omni_transfer_config, - stage_id=stage_id, + stage_id=configured_stage_id, async_chunk=async_chunk, ) - omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) + omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, configured_stage_id) if metadata.stage_type == "diffusion": + is_remote_diffusion_stage = ( + self.single_stage_mode + and self._single_stage_id_filter is not None + and configured_stage_id != self._single_stage_id_filter + ) + if is_remote_diffusion_stage: + assert self._omni_master_server is not None + stage_clients[stage_idx] = self._create_remote_diffusion_stage( + metadata, + stage_init_timeout, + self._omni_master_server, + ) + continue + with llm_stage_launch_lock: previous_visible_devices = os.environ.get(device_control_env) try: - setup_stage_devices(stage_id, metadata.runtime_cfg) + setup_stage_devices(configured_stage_id, metadata.runtime_cfg) omni_conn_cfg, omni_from, omni_to = omni_kv_connector if omni_conn_cfg: from vllm_omni.entrypoints.utils import inject_omni_kv_config inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) - _inject_kv_stage_info(stage_cfg, stage_id) - stage_clients[stage_id] = initialize_diffusion_stage( - self.model, - stage_cfg, - metadata, - batch_size=self.diffusion_batch_size, - ) + inject_kv_stage_info(stage_cfg, configured_stage_id) + if self.single_stage_mode: + assert self._omni_master_server is not None + stage_clients[stage_idx] = self._launch_diffusion_stage( + stage_cfg, + metadata, + self._omni_master_server, + ) + else: + stage_clients[stage_idx] = initialize_diffusion_stage( + self.model, + stage_cfg, + metadata, + batch_size=self.diffusion_batch_size, + ) logger.info( "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)", - stage_id, + configured_stage_id, self.diffusion_batch_size, ) finally: @@ -555,30 +773,58 @@ class AsyncOmniEngine: current_omni_platform.set_device_control_env_var(previous_visible_devices) continue - llm_stage_ids.append(stage_id) - llm_launch_futures[stage_id] = launch_executor.submit( - self._launch_llm_stage, - stage_cfg, - metadata, - stage_connector_spec, - stage_init_timeout, - llm_stage_launch_lock, - omni_kv_connector, - ) + llm_stage_positions.append(stage_idx) + + # In single-stage mode, stages that don't match the local + # stage_id filter are skipped. + if ( + self.single_stage_mode + and self._single_stage_id_filter is not None + and configured_stage_id != self._single_stage_id_filter + ): + assert self._omni_master_server is not None + llm_launch_futures[stage_idx] = launch_executor.submit( + self._create_remote_llm_stage, + stage_cfg, + metadata, + stage_connector_spec, + stage_init_timeout, + self._omni_master_server, + ) + else: + llm_launch_futures[stage_idx] = launch_executor.submit( + self._launch_llm_stage, + stage_cfg, + metadata, + stage_connector_spec, + stage_init_timeout, + llm_stage_launch_lock, + omni_kv_connector, + ) concurrent.futures.wait(list(llm_launch_futures.values())) - for stage_id in llm_stage_ids: - started_llm_stages[stage_id] = llm_launch_futures[stage_id].result() + for stage_idx in llm_stage_positions: + started_llm_stages[stage_idx] = llm_launch_futures[stage_idx].result() - for stage_id in llm_stage_ids: - started = started_llm_stages[stage_id] - stage_client, output_processor, vllm_config, stage0_input_processor = self._attach_llm_stage(started) - stage_clients[stage_id] = stage_client - output_processors[stage_id] = output_processor - stage_vllm_configs[stage_id] = vllm_config - if stage0_input_processor is not None: - input_processor = stage0_input_processor + attach_futures: dict[concurrent.futures.Future[tuple[Any, Any, Any, InputProcessor | None]], int] = {} + with concurrent.futures.ThreadPoolExecutor( + max_workers=max(1, len(llm_stage_positions)), + thread_name_prefix="llm-stage-attach", + ) as attach_executor: + for stage_idx in llm_stage_positions: + attach_futures[attach_executor.submit(self._attach_llm_stage, started_llm_stages[stage_idx])] = ( + stage_idx + ) + + for future in concurrent.futures.as_completed(attach_futures): + stage_idx = attach_futures[future] + stage_client, output_processor, vllm_config, stage0_input_processor = future.result() + stage_clients[stage_idx] = stage_client + output_processors[stage_idx] = output_processor + stage_vllm_configs[stage_idx] = vllm_config + if stage0_input_processor is not None: + input_processor = stage0_input_processor initialized_stage_clients, default_sampling_params_list, stage_metadata = finalize_initialized_stages( stage_clients, @@ -595,8 +841,13 @@ class AsyncOmniEngine: ) cleanup_failed_stage_initialization( stage_clients, - [started_llm_stages[stage_id] for stage_id in llm_stage_ids if stage_id in started_llm_stages], + [started_llm_stages[stage_idx] for stage_idx in llm_stage_positions if stage_idx in started_llm_stages], ) + 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") raise self.stage_clients = initialized_stage_clients @@ -1310,3 +1561,10 @@ class AsyncOmniEngine: q.close() except Exception: pass + + if self._omni_master_server is not None: + try: + self._omni_master_server.stop() + except Exception: + logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during shutdown") + self._omni_master_server = None diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 71a0aee4a..52e674f47 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -12,7 +12,7 @@ from urllib.parse import urlparse from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest -from vllm.v1.engine.core_client import AsyncMPClient +from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient from vllm_omni.distributed.omni_connectors.utils.initialization import KV_TRANSFER_PORT_OFFSET from vllm_omni.engine.stage_init_utils import StageMetadata @@ -25,18 +25,54 @@ if TYPE_CHECKING: logger = init_logger(__name__) -class StageEngineCoreClient(AsyncMPClient): - """Stage async client that inherits from vLLM's AsyncMPClient. +class StageEngineCoreClientBase: + """Shared stage-aware behavior for async EngineCore clients. - Fully reuses AsyncMPClient for: + The concrete transport/load-balancing behavior is supplied by the + multiprocessing client subclass in the MRO. + + Fully reuses the underlying vLLM async MP client ``__init__`` for: - ZMQ setup, sockets - 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. """ + @staticmethod + def make_async_mp_client( + vllm_config: Any, + executor_class: type, + metadata: StageMetadata, + client_addresses: dict[str, str] | None = None, + proc: Any = None, + engine_manager: Any = None, + coordinator: Any = None, + client_count: int = 1, + client_index: int = 0, + ) -> StageEngineCoreClient | DPLBStageEngineCoreClient: + """Create the appropriate stage async client for the DP mode.""" + parallel_config = vllm_config.parallel_config + client_args = dict( + vllm_config=vllm_config, + executor_class=executor_class, + metadata=metadata, + client_addresses=client_addresses, + proc=proc, + engine_manager=engine_manager, + coordinator=coordinator, + client_count=client_count, + client_index=client_index, + ) + + if parallel_config.data_parallel_size > 1 and not parallel_config.data_parallel_external_lb: + return DPLBStageEngineCoreClient(**client_args) + + return StageEngineCoreClient(**client_args) + def __init__( self, vllm_config: Any, @@ -85,8 +121,10 @@ class StageEngineCoreClient(AsyncMPClient): self._kv_sender_info: dict[str, Any] | None = None self._kv_sender_initialized = False + client_name = self.__class__.__name__ logger.info( - "[StageEngineCoreClient] Stage-%s initializing EngineCore", + "[%s] Stage-%s initializing EngineCore", + client_name, self.stage_id, ) try: @@ -98,23 +136,30 @@ class StageEngineCoreClient(AsyncMPClient): client_count=client_count, client_index=client_index, ) + if engine_manager is not None: + self.resources.engine_manager = engine_manager + if coordinator is not None: + self.resources.coordinator = coordinator except Exception: logger.exception( - "[StageEngineCoreClient] Stage-%s EngineCore init failed", + "[%s] Stage-%s EngineCore init failed", + client_name, self.stage_id, ) try: self.shutdown() except Exception as shutdown_error: logger.warning( - "[StageEngineCoreClient] Stage-%s cleanup after init failure failed: %s", + "[%s] Stage-%s cleanup after init failure failed: %s", + client_name, self.stage_id, shutdown_error, ) raise self._initialize_kv_sender_endpoint() logger.info( - "[StageEngineCoreClient] Stage-%s EngineCore running", + "[%s] Stage-%s EngineCore running", + client_name, self.stage_id, ) @@ -122,7 +167,12 @@ class StageEngineCoreClient(AsyncMPClient): async def add_request_async(self, request: EngineCoreRequest) -> None: """Add request to the stage engine core.""" - logger.info(f"[StageEngineCoreClient] Stage-{self.stage_id} adding request: {request.request_id}") + logger.info( + "[%s] Stage-%s adding request: %s", + self.__class__.__name__, + self.stage_id, + request.request_id, + ) await super().add_request_async(request) # ==================== Stage Methods ==================== @@ -287,9 +337,9 @@ class StageEngineCoreClient(AsyncMPClient): ) -> Any: """Forward control RPCs to the underlying AsyncMPClient stage engine. - Each ``StageEngineCoreClient`` already represents one logical stage, so - stage-scoped control operations should be executed here and then fanned - in-core across the workers managed by this EngineCore client. + Each stage client already represents one logical stage, so stage-scoped + control operations should be executed here and then fanned in-core + across the workers managed by this EngineCore client. """ return await super().collective_rpc_async( method=method, @@ -299,10 +349,19 @@ class StageEngineCoreClient(AsyncMPClient): ) def shutdown(self) -> None: - """Shutdown ZMQ connections and the subprocess.""" + """Shutdown managed resources and any externally spawned subprocess.""" super().shutdown() if self._proc is not None and self._proc.is_alive(): self._proc.terminate() self._proc.join(timeout=5) if self._proc.is_alive(): self._proc.kill() + self._proc = None + + +class StageEngineCoreClient(StageEngineCoreClientBase, AsyncMPClient): + """Stage async client backed by vLLM's ``AsyncMPClient``.""" + + +class DPLBStageEngineCoreClient(StageEngineCoreClientBase, DPLBAsyncMPClient): + """Stage async client backed by vLLM's ``DPLBAsyncMPClient``.""" diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py new file mode 100644 index 000000000..6af66c71f --- /dev/null +++ b/vllm_omni/engine/stage_engine_startup.py @@ -0,0 +1,599 @@ +"""Helpers for launching and handshaking omni engine cores.""" + +from __future__ import annotations + +import contextlib +import dataclasses +import threading +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Any + +import msgspec +import zmq +from omegaconf import OmegaConf +from vllm.config import CacheConfig, VllmConfig +from vllm.logger import init_logger +from vllm.utils.network_utils import get_open_port, 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, + wait_for_engine_startup, +) +from vllm.v1.executor import Executor + +logger = init_logger(__name__) + +# Poll period (ms) used by the registration/handshake loop. +_POLL_PERIOD_MS = 5_000 +# Default timeout (s) for a stage to send READY. +_DEFAULT_STARTUP_TIMEOUT_S = 300 + + +def _serialize_stage_config(stage_config: Any) -> Any: + """Convert a stage config to msgpack-friendly builtins.""" + if stage_config is None or isinstance(stage_config, (str, bytes, int, float, bool)): + return stage_config + + if OmegaConf.is_config(stage_config): + return _serialize_stage_config(OmegaConf.to_container(stage_config, resolve=True)) + + if dataclasses.is_dataclass(stage_config): + return _serialize_stage_config(dataclasses.asdict(stage_config)) + + if isinstance(stage_config, dict): + return {key: _serialize_stage_config(value) for key, value in stage_config.items() if not callable(value)} + + if isinstance(stage_config, (list, tuple, set)): + return [_serialize_stage_config(item) for item in stage_config if not callable(item)] + + if hasattr(stage_config, "items"): + return {key: _serialize_stage_config(value) for key, value in stage_config.items() if not callable(value)} + + if hasattr(stage_config, "__dict__"): + return { + key: _serialize_stage_config(value) + for key, value in vars(stage_config).items() + if not key.startswith("_") and not callable(value) + } + + return stage_config + + +# --------------------------------------------------------------------------- +# Per-stage address allocation +# --------------------------------------------------------------------------- + + +@dataclass +class StageAllocation: + """ZMQ addresses reserved for a single stage.""" + + # Per-stage handshake socket (OmniMasterServer binds, engine connects) + handshake_bind_address: str + handshake_connect_address: str + # Input channel: client binds ROUTER, engine connects DEALER + input_bind_address: str + input_connect_address: str + # Output channel: client binds PULL, engine connects PUSH + output_bind_address: str + output_connect_address: str + + +@dataclass(frozen=True) +class StageCoordinatorAddresses: + """Optional DP coordinator addresses registered for a stage.""" + + coordinator_input: str | None = None + coordinator_output: str | None = None + frontend_stats_publish_address: str | None = None + + +# --------------------------------------------------------------------------- +# OmniMasterServer +# --------------------------------------------------------------------------- + + +class OmniMasterServer: + """Registration server for single-stage engine startup.""" + + def __init__( + self, + master_address: str, + master_port: int, + stage_ids: list[int], + ) -> None: + self._address = master_address + self._port = master_port + self._allocations: dict[int, StageAllocation] = {} + self._stage_configs: dict[int, Any] = {} + self._stage_coordinator_addresses: dict[int, StageCoordinatorAddresses] = {} + self._stage_config_events: dict[int, threading.Event] = {} + self._thread: threading.Thread | None = None + self._stop_event = threading.Event() + + for sid in stage_ids: + self._stage_config_events[sid] = threading.Event() + self._stage_coordinator_addresses[sid] = StageCoordinatorAddresses() + hs_port = get_open_port() + inp_port = get_open_port() + out_port = get_open_port() + self._allocations[sid] = StageAllocation( + handshake_bind_address=f"tcp://{master_address}:{hs_port}", + handshake_connect_address=f"tcp://{master_address}:{hs_port}", + input_bind_address=f"tcp://{master_address}:{inp_port}", + input_connect_address=f"tcp://{master_address}:{inp_port}", + output_bind_address=f"tcp://{master_address}:{out_port}", + output_connect_address=f"tcp://{master_address}:{out_port}", + ) + + logger.info( + "[OmniMasterServer] Pre-allocated addresses for stages %s (master=%s:%d)", + list(stage_ids), + master_address, + master_port, + ) + + # ------------------------------------------------------------------ + # Public helpers + # ------------------------------------------------------------------ + @property + def address(self) -> str: + """Return the registration address exposed to stage launchers.""" + return self._address + + @property + def port(self) -> int: + """Return the registration port exposed to stage launchers.""" + return self._port + + def get_allocation(self, stage_id: int) -> StageAllocation: + """Return the full address allocation for *stage_id*.""" + return self._allocations[stage_id] + + def register_stage_config( + self, + stage_id: int, + stage_config: Any, + coordinator_addresses: StageCoordinatorAddresses | None = None, + ) -> None: + """Store the latest stage registration payload for *stage_id*.""" + if stage_id not in self._allocations: + raise KeyError(stage_id) + self._stage_configs[stage_id] = stage_config + if coordinator_addresses is not None: + self._stage_coordinator_addresses[stage_id] = coordinator_addresses + self._stage_config_events[stage_id].set() + + def get_stage_config(self, stage_id: int, timeout_s: float | None = None) -> Any: + """Return the stage config for *stage_id*, waiting if necessary.""" + if stage_id not in self._allocations: + raise KeyError(stage_id) + + if stage_id in self._stage_configs: + return self._stage_configs[stage_id] + + if not self._stage_config_events[stage_id].wait(timeout=timeout_s): + raise TimeoutError(f"Timed out waiting for stage config for stage {stage_id}.") + + return self._stage_configs[stage_id] + + def get_stage_coordinator_addresses( + self, + stage_id: int, + timeout_s: float | None = None, + ) -> StageCoordinatorAddresses: + """Return the registered coordinator addresses for *stage_id*.""" + if stage_id not in self._allocations: + raise KeyError(stage_id) + + if not self._stage_config_events[stage_id].is_set(): + if not self._stage_config_events[stage_id].wait(timeout=timeout_s): + raise TimeoutError(f"Timed out waiting for stage registration for stage {stage_id}.") + + return self._stage_coordinator_addresses[stage_id] + + def get_client_addresses(self, stage_id: int) -> dict[str, str]: + """Return the addresses the client-side sockets should *bind* to.""" + alloc = self._allocations[stage_id] + return { + "input_address": alloc.input_bind_address, + "output_address": alloc.output_bind_address, + } + + def get_zmq_addresses(self, stage_id: int) -> EngineZmqAddresses: + """Return EngineZmqAddresses using the *bind* (client) side addresses.""" + alloc = self._allocations[stage_id] + return EngineZmqAddresses( + inputs=[alloc.input_bind_address], + outputs=[alloc.output_bind_address], + ) + + def get_engine_zmq_addresses(self, stage_id: int) -> EngineZmqAddresses: + """Return EngineZmqAddresses using the *connect* (engine) addresses.""" + alloc = self._allocations[stage_id] + return EngineZmqAddresses( + inputs=[alloc.input_connect_address], + outputs=[alloc.output_connect_address], + ) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def start(self) -> None: + """Start the background server thread.""" + self._thread = threading.Thread( + target=self._run, + name="OmniMasterServer", + daemon=True, + ) + self._thread.start() + logger.info( + "[OmniMasterServer] Listening on tcp://%s:%d", + self.address, + self.port, + ) + + def stop(self) -> None: + """Signal stop and join the background thread.""" + self._stop_event.set() + if self._thread is not None: + self._thread.join(timeout=10) + + # ------------------------------------------------------------------ + # Internal server logic + # ------------------------------------------------------------------ + + def _run(self) -> None: + ctx = zmq.Context() + try: + self._serve(ctx) + except Exception: + logger.exception("[OmniMasterServer] Server thread crashed") + finally: + ctx.term() + + def _serve(self, ctx: zmq.Context) -> None: # type: ignore[type-arg] + # Registration socket for the initial stage registration. + # Per-stage handshake sockets are bound by the launch helpers. + reg_socket: zmq.Socket = ctx.socket(zmq.ROUTER) # type: ignore[attr-defined] + reg_socket.bind(f"tcp://{self.address}:{self.port}") + + poller = zmq.Poller() + poller.register(reg_socket, zmq.POLLIN) + + pending: set[int] = set(self._allocations.keys()) + + while pending and not self._stop_event.is_set(): + events: list[tuple[zmq.Socket, int]] = poller.poll(_POLL_PERIOD_MS) # type: ignore[assignment] + if not events: + logger.debug("[OmniMasterServer] Still waiting for registration from stages: %s", pending) + continue + + for sock, _ in events: + if sock is reg_socket: + sid = self._handle_registration(reg_socket) + if sid is not None: + pending.discard(sid) + + # Cleanup + reg_socket.close(linger=0) + logger.info("[OmniMasterServer] All stages registered; server thread exiting.") + + def _handle_registration(self, reg_socket: zmq.Socket) -> int | None: # type: ignore[type-arg] + """Receive a stage registration and reply with the handshake address. + + Returns the registered stage_id on success, or None on failure. + """ + frames = reg_socket.recv_multipart() + if len(frames) < 2: + logger.warning( + "[OmniMasterServer] Unexpected registration frame count: %d", + len(frames), + ) + return None + identity = frames[0] + msg_bytes = frames[-1] + try: + msg = msgspec.msgpack.decode(msg_bytes) + except Exception as exc: + logger.warning("[OmniMasterServer] Failed to decode registration message: %s", exc) + return None + + stage_id: int | None = msg.get("stage_id") + if stage_id not in self._allocations: + logger.warning( + "[OmniMasterServer] Received registration for unknown stage_id=%s", + stage_id, + ) + return None + + self.register_stage_config( + stage_id, + msg.get("stage_config"), + coordinator_addresses=StageCoordinatorAddresses( + coordinator_input=msg.get("coordinator_input"), + coordinator_output=msg.get("coordinator_output"), + frontend_stats_publish_address=msg.get("frontend_stats_publish_address"), + ), + ) + + alloc = self._allocations[stage_id] + response = msgspec.msgpack.encode( + { + "handshake_address": alloc.handshake_connect_address, + "input_address": alloc.input_bind_address, + "output_address": alloc.output_bind_address, + } + ) + # ROUTER-DEALER: reply is [identity, payload] (no empty delimiter). + reg_socket.send_multipart([identity, response]) + logger.info( + "[OmniMasterServer] Stage %d registered; assigned handshake=%s", + stage_id, + alloc.handshake_connect_address, + ) + return stage_id + + +def register_stage_with_omni_master( + *, + omni_master_address: str, + omni_master_port: int, + omni_stage_id: int, + omni_stage_config: Any = None, + coordinator: DPCoordinator | None = None, + return_addresses: bool = False, +) -> str | tuple[str, str, str]: + """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. + """ + + reg_ctx = zmq.Context() + try: + reg_sock: zmq.Socket = reg_ctx.socket(zmq.DEALER) # type: ignore[attr-defined] + try: + reg_sock.connect(f"tcp://{omni_master_address}:{omni_master_port}") + payload = { + "stage_id": omni_stage_id, + "stage_config": _serialize_stage_config(omni_stage_config), + } + if coordinator is not None: + coordinator_input, coordinator_output = coordinator.get_engine_socket_addresses() + payload["coordinator_input"] = coordinator_input + payload["coordinator_output"] = coordinator_output + payload["frontend_stats_publish_address"] = coordinator.get_stats_publish_address() + + reg_sock.send(msgspec.msgpack.encode(payload)) + timeout_ms = _DEFAULT_STARTUP_TIMEOUT_S * 1_000 + if not reg_sock.poll(timeout=timeout_ms): + raise RuntimeError( + f"Timed out waiting for registration " + f"response from OmniMasterServer " + f"({omni_master_address}:{omni_master_port}) " + f"for stage {omni_stage_id}." + ) + response_bytes = reg_sock.recv() + response = msgspec.msgpack.decode(response_bytes) + handshake_address: str = response["handshake_address"] + input_address: str = response["input_address"] + output_address: str = response["output_address"] + logger.info( + "Stage %d registered; handshake_address=%s", + omni_stage_id, + handshake_address, + ) + finally: + reg_sock.close(linger=0) + finally: + reg_ctx.term() + + 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: + num_gpu_blocks = (cache_config.num_gpu_blocks or 0) + msg["num_gpu_blocks"] + cache_config.num_gpu_blocks = num_gpu_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=%d)", eng_index, msg["num_gpu_blocks"]) + + else: + raise RuntimeError(f"[omni] Unexpected status '{status}' from engine {eng_index} in state {engine.state}.") + + +@contextlib.contextmanager +def connect_remote_engine_cores( + vllm_config: VllmConfig, + omni_master_server: OmniMasterServer, + stage_id: int, +) -> Iterator[tuple[None, DPCoordinator | None, EngineZmqAddresses]]: + """Wait for remote engine cores to connect through the omni handshake.""" + addresses = omni_master_server.get_zmq_addresses(stage_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) + ) + start_index = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 + coordinator = None + + registered_coordinator_addresses = omni_master_server.get_stage_coordinator_addresses(stage_id) + addresses.coordinator_input = registered_coordinator_addresses.coordinator_input + addresses.coordinator_output = registered_coordinator_addresses.coordinator_output + addresses.frontend_stats_publish_address = registered_coordinator_addresses.frontend_stats_publish_address + + engines_to_handshake = [CoreEngine(index=start_index + i, local=False) for i in range(remote_engine_count)] + + logger.info( + "Waiting for %d remote engine(s) for stage %d", + remote_engine_count, + stage_id, + ) + + handshake_bind_address = omni_master_server.get_allocation(stage_id).handshake_bind_address + + with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: + yield None, coordinator, addresses + + _wait_for_omni_engine_startup( + handshake_socket, + addresses, + engines_to_handshake, + vllm_config.cache_config, + ) + + +@contextlib.contextmanager +def launch_omni_core_engines( + vllm_config: VllmConfig, + executor_class: type[Executor], + log_stats: bool, + omni_master_server: OmniMasterServer, + stage_id: int, + stage_config: Any = None, +) -> Iterator[tuple[CoreEngineProcManager, DPCoordinator | None, EngineZmqAddresses]]: + """Launch local engine cores using the omni registration flow.""" + addresses = omni_master_server.get_zmq_addresses(stage_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) + ) + 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 + + # Run the DP Coordinator process with rank 0 when in online DP mode. + # The coordinator is needed for: + # 1. Internal/hybrid LB: collecting and publishing queue stats + # 2. MoE models: wave coordination in addition to stats + run_coordinator = vllm_config.needs_dp_coordinator and dp_rank == 0 + + if run_coordinator: + coordinator = DPCoordinator( + parallel_config, + enable_wave_coordination=vllm_config.model_config.is_moe, + ) + + addresses.coordinator_input, addresses.coordinator_output = coordinator.get_engine_socket_addresses() + addresses.frontend_stats_publish_address = coordinator.get_stats_publish_address() + + logger.info( + "[omni] Started DP Coordinator process for stage %d (PID: %d)", + stage_id, + coordinator.proc.pid, + ) + else: + coordinator = None + + logger.info( + "Starting %d local engine(s) for stage %d (dp_rank=%d)", + local_engine_count, + stage_id, + dp_rank, + ) + + # 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( + omni_master_address=omni_master_server.address, + omni_master_port=omni_master_server.port, + omni_stage_id=stage_id, + omni_stage_config=stage_config, + coordinator=coordinator, + ) + + # One CoreEngine entry per local engine so wait_for_engine_startup can + # track the HELLO/READY handshake for each of them. + engines_to_handshake = [CoreEngine(index=start_index + i, local=True) for i in range(local_engine_count)] + + # Bind the pre-allocated handshake socket for this stage. + handshake_bind_address = omni_master_server.get_allocation(stage_id).handshake_bind_address + + with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: + 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, + vllm_config.cache_config, + local_engine_manager, + coordinator.proc if coordinator else None, + ) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index e6f603d2a..09195faec 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -101,6 +101,34 @@ def resolve_worker_cls(engine_args: dict[str, Any]) -> None: raise ValueError(f"Unknown worker_type: {worker_type}") +def inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None: + """Inject stage metadata into omni_kv_config when present.""" + try: + engine_args = stage_cfg.engine_args + if hasattr(engine_args, "get"): + omni_kv = engine_args.get("omni_kv_config", None) + else: + omni_kv = getattr(engine_args, "omni_kv_config", None) + + if omni_kv is None: + return + + if hasattr(omni_kv, "setdefault"): + omni_kv.setdefault("stage_id", stage_id) + elif hasattr(omni_kv, "__setitem__"): + if "stage_id" not in omni_kv: + omni_kv["stage_id"] = stage_id + + engine_input_source = getattr(stage_cfg, "engine_input_source", None) + if engine_input_source is not None: + if hasattr(omni_kv, "setdefault"): + omni_kv.setdefault("engine_input_source", list(engine_input_source)) + elif hasattr(omni_kv, "__setitem__") and "engine_input_source" not in omni_kv: + omni_kv["engine_input_source"] = list(engine_input_source) + except Exception as e: + logger.debug("Failed to inject stage info into omni_kv_config: %s", e) + + @dataclass class StageMetadata: """Lightweight stage attributes extracted from stage_config.""" @@ -129,8 +157,10 @@ class StartedLlmStage: metadata: Any vllm_config: Any executor_class: type - proc: Any addresses: Any + proc: Any = None + engine_manager: Any = None + coordinator: Any = None def extract_stage_metadata(stage_config: Any) -> StageMetadata: @@ -263,6 +293,7 @@ def build_vllm_config( model: str, stage_connector_spec: dict[str, Any] | None = None, engine_args_dict: dict[str, Any] | None = None, + headless: bool = False, ) -> tuple[Any, type]: """Build engine args, then create VllmConfig and executor_class. @@ -278,7 +309,10 @@ def build_vllm_config( filtered_engine_args_dict = filter_dataclass_kwargs(OmniEngineArgs, engine_args_dict) omni_engine_args = OmniEngineArgs(**filtered_engine_args_dict) - vllm_config = omni_engine_args.create_engine_config(usage_context=UsageContext.LLM_CLASS) + vllm_config = omni_engine_args.create_engine_config( + usage_context=UsageContext.LLM_CLASS, + headless=headless, + ) executor_class = Executor.get_class(vllm_config) return vllm_config, executor_class @@ -445,6 +479,37 @@ def get_stage_connector_spec( return {} +def build_diffusion_config( + model: str, + stage_cfg: Any, + metadata: StageMetadata, +) -> Any: + """Build diffusion config for a stage.""" + from vllm_omni.diffusion.data import OmniDiffusionConfig + + engine_args_dict = build_engine_args_dict(stage_cfg, model) + od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict) + + num_devices_per_stage = od_config.parallel_config.world_size + device_control_env = current_omni_platform.device_control_env_var + visible_devices_str = os.environ.get(device_control_env) if device_control_env else None + if visible_devices_str: + physical_devices = [device.strip() for device in visible_devices_str.split(",") if device.strip()] + else: + physical_devices = list(range(current_omni_platform.get_device_count())) + + if len(physical_devices) < num_devices_per_stage: + raise ValueError( + f"Stage {metadata.stage_id} requires {num_devices_per_stage} device(s) based on parallel_config, " + f"but {len(physical_devices)} device(s) are available: {physical_devices}" + ) + + od_config.num_gpus = num_devices_per_stage + if metadata.cfg_kv_collect_func is not None: + od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func + return od_config + + def initialize_diffusion_stage( model: str, stage_cfg: Any, @@ -461,30 +526,9 @@ def initialize_diffusion_stage( diffusion engine. Passed through to ``StageDiffusionClient`` and ultimately to ``AsyncOmni``. """ - from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient - od_config = OmniDiffusionConfig.from_kwargs( - model=model, - **_to_dict(stage_cfg.engine_args), - ) - num_devices_per_stage = od_config.parallel_config.world_size - device_control_env = current_omni_platform.device_control_env_var - visible_devices_str = os.environ.get(device_control_env) - if visible_devices_str: - physical_devices = [device.strip() for device in visible_devices_str.split(",") if device.strip()] - else: - physical_devices = list(range(current_omni_platform.get_device_count())) - - if len(physical_devices) < num_devices_per_stage: - raise ValueError( - f"Stage {metadata.stage_id} requires {num_devices_per_stage} device(s) based on parallel_config, " - f"but {len(physical_devices)} device(s) are available: {physical_devices}" - ) - - od_config.num_gpus = num_devices_per_stage - if metadata.cfg_kv_collect_func is not None: - od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func + od_config = build_diffusion_config(model, stage_cfg, metadata) return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size) @@ -518,17 +562,18 @@ def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int def close_started_llm_stage(started: StartedLlmStage) -> None: - """Terminate the subprocess owned by a launched stage that never attached.""" - if started.proc is None: - return - try: - terminate_alive_proc(started.proc) - except Exception as cleanup_error: - logger.warning( - "[stage_init] Failed to terminate process for stage %s: %s", - started.stage_id, - cleanup_error, - ) + """Release resources owned by a launched stage that never attached.""" + if started.proc is not None: + try: + terminate_alive_proc(started.proc) + except Exception as cleanup_error: + logger.warning( + "[stage_init] Failed to terminate process for stage %s: %s", + started.stage_id, + cleanup_error, + ) + _shutdown_or_close_resource(started.engine_manager, "engine manager", started.stage_id) + _shutdown_or_close_resource(started.coordinator, "coordinator", started.stage_id) def finalize_initialized_stages( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index b72df41cd..6e9adc246 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -8,6 +8,8 @@ diffusion models (e.g., Qwen-Image) through the same CLI interface. import argparse import json import os +import signal +from types import FrameType from typing import Any import uvloop @@ -419,22 +421,228 @@ def _create_default_diffusion_stage_cfg(args: argparse.Namespace) -> list[dict[s def run_headless(args: argparse.Namespace) -> None: - """Run a single stage in headless mode. + """Run a single stage in headless mode.""" + from vllm.v1.engine.coordinator import DPCoordinator + from vllm.v1.engine.utils import CoreEngineProcManager + from vllm.v1.executor.multiproc_executor import MultiprocExecutor + from vllm.version import __version__ as VLLM_VERSION - .. deprecated:: 0.x.x - Headless mode is deprecated and will be removed in a future version. - It is only compatible with the old OmniStage-based runtime. - The current AsyncOmniEngine-based runtime does not support headless mode. - - Raises: - RuntimeError: Always raises an error indicating headless mode is deprecated. - """ - raise RuntimeError( - "Headless mode is deprecated and not supported in the current runtime. " - "Please use the standard orchestrator mode (without --headless flag). " - "If you need distributed deployment, consider using Ray backend or " - "other distributed serving solutions." + 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.stage_engine_startup import register_stage_with_omni_master + 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, + inject_kv_stage_info, + load_omni_transfer_config_for_model, + prepare_engine_environment, + terminate_alive_proc, + ) + from vllm_omni.entrypoints.utils import inject_omni_kv_config, load_and_resolve_stage_configs + + model = args.model + stage_id: int | None = args.stage_id + omni_master_address: str | None = args.omni_master_address + omni_master_port: int | None = args.omni_master_port + + if stage_id is None: + raise ValueError("--stage-id is required in headless mode") + if omni_master_address is None or omni_master_port is None: + raise ValueError("--omni-master-address and --omni-master-port are required in headless mode") + if getattr(args, "api_server_count", 0) and args.api_server_count > 1: + raise ValueError("api_server_count can't be set in headless mode") + if args.worker_backend != "multi_process": + raise ValueError("headless mode requires worker_backend=multi_process") + + args_dict = vars(args).copy() + config_path, stage_configs = load_and_resolve_stage_configs( + model, + args_dict.get("stage_configs_path"), + args_dict, + ) + + # Locate the stage config that matches stage_id. + stage_cfg = None + for cfg in stage_configs: + if getattr(cfg, "stage_id", None) == stage_id: + stage_cfg = cfg + break + if stage_cfg is None: + raise ValueError( + f"No stage config found for stage_id={stage_id}. " + f"Available stage ids: {[getattr(c, 'stage_id', None) for c in stage_configs]}" + ) + + 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) + + if getattr(stage_cfg, "stage_type", "llm") == "diffusion": + metadata = extract_stage_metadata(stage_cfg) + if omni_conn_cfg: + inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) + inject_kv_stage_info(stage_cfg, stage_id) + od_config = build_diffusion_config(model, stage_cfg, metadata) + + logger.info( + "[Headless] Launching diffusion stage %d via OmniMasterServer at %s:%d", + stage_id, + omni_master_address, + omni_master_port, + ) + + proc = None + try: + handshake_address, request_address, response_address = 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, + return_addresses=True, + ) + proc, _, _, _ = spawn_diffusion_proc( + model, + od_config, + handshake_address=handshake_address, + request_address=request_address, + response_address=response_address, + ) + complete_diffusion_handshake(proc, handshake_address) + proc.join() + if proc.exitcode not in (None, 0): + raise RuntimeError(f"Diffusion stage {stage_id} exited with code {proc.exitcode}") + return + finally: + logger.info("[Headless] Shutting down stage %d.", stage_id) + if proc is not None and proc.is_alive(): + terminate_alive_proc(proc) + + stage_connector_spec = get_stage_connector_spec( + omni_transfer_config=omni_transfer_config, + stage_id=stage_id, + async_chunk=False, + ) + + # Device assignment is managed externally (e.g. CUDA_VISIBLE_DEVICES); + # runtime_cfg is intentionally ignored in headless mode. + engine_args_dict = build_engine_args_dict( + stage_cfg, + model, + stage_connector_spec=stage_connector_spec, + ) + + # 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 + + vllm_config, executor_class = build_vllm_config( + stage_cfg, + model, + stage_connector_spec=stage_connector_spec, + engine_args_dict=engine_args_dict, + 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 + + def signal_handler(signum: int, frame: FrameType | None) -> None: + nonlocal shutdown_requested + logger.debug("Received %d signal.", signum) + if not shutdown_requested: + shutdown_requested = True + raise SystemExit + + signal.signal(signal.SIGTERM, signal_handler) + signal.signal(signal.SIGINT, signal_handler) + + if parallel_config.node_rank_within_dp > 0: + head_node_address = f"{parallel_config.master_addr}:{parallel_config.master_port}" + logger.info( + "Launching vLLM-Omni (v%s) headless multiproc executor, " + "with head node address %s for torch.distributed process group.", + VLLM_VERSION, + head_node_address, + ) + + executor = MultiprocExecutor(vllm_config, monitor_workers=False) + 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 engine core(s) for stage %d via OmniMasterServer at %s:%d", + local_engine_count, + stage_id, + omni_master_address, + omni_master_port, + ) + + # Headless mode launches all local engine cores for a single stage. + # The OmniMasterServer allocates one handshake endpoint per stage, so we + # register the stage once here and let every local engine core reuse the + # returned handshake address directly. + handshake_address = 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, + ) + + engine_manager = None + log_stats = bool(getattr(args, "log_stats", False)) + if getattr(args, "disable_log_stats", False): + log_stats = False + + try: + engine_manager = CoreEngineProcManager( + local_engine_count=local_engine_count, + start_index=dp_rank, + local_start_index=0, + vllm_config=vllm_config, + local_client=False, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + ) + engine_manager.join_first() + finally: + logger.info("[Headless] Shutting down stage %d.", stage_id) + if engine_manager is not None: + engine_manager.shutdown() + if coordinator is not None: + coordinator.shutdown() def cmd_init() -> list[CLISubcommand]: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 0706b9898..4519ae8c0 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -353,7 +353,8 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, try: await shutdown_task finally: - serving_speech = getattr(getattr(app, "state", None), "openai_serving_speech", None) + state = getattr(app, "state", None) + serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None if serving_speech is not None: serving_speech.shutdown() sock.close()