[Frontend / Core] Support Running Frontend Only & Validate Stage IDs

docs/cli/serve.md

@@ -107,6 +107,9 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --omni-master-port 26000
 ```
 
+!!! note
+    you may pass `--stage-id none` to run only the API server and orchestrator. This is useful when you want to launch every stage, including stage 0, as a headless worker.
+
 When utilizing a custom deployment YAML, append `--deploy-config /path/to/override.yaml` to each command execution.
 
 In the standard execution paradigm, the `--stage-overrides` argument is utilized to apply stage-specific configurations from a single CLI command.

docs/configuration/stage_configs.md

@@ -336,6 +336,8 @@ The stage-based CLI paradigm facilitates the execution of discrete pipeline stag
 - **Worker Stages** operate without a distinct API server (i.e., using `--headless`), are assigned sequential `--stage-id` identifiers, and must reference the corresponding
   `--omni-master-address` and `--omni-master-port` parameters to successfully register with Stage 0.
 
+You may also explicitly pass `--stage-id none` to run only the API server and orchestrator if you want to run every stage as headless.
+
 For migrated architectures, the system automatically resolves and loads the bundled deployment YAML. Consequently, the primary execution path
 does **not** necessitate the explicit definition of `--deploy-config`:
 the example below uses `CUDA_VISIBLE_DEVICES=0` for Stage 0 and

docs/mkdocs/hooks/generate_argparse.py

@@ -127,6 +127,7 @@ def add_parser(self, name, **kwargs):
                         "_nonneg_finite_float": float,
                         "_json_object": json.loads,
                         "_parse_stage_overrides": __import__("json").loads,
+                        "_parse_stage_id": int,
                         "VLLM_SUBCMD_PARSER_EPILOG": "",
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,

tests/engine/test_single_stage_mode.py

@@ -23,6 +23,7 @@
     StageExecutionType,
     StagePipelineConfig,
 )
+from vllm_omni.engine.arg_utils import FRONTEND_ONLY_ID_FILTER
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
 from vllm_omni.engine.stage_engine_startup import (
@@ -459,14 +460,44 @@ def test_stage_id_kwarg_promotes_to_single_stage_mode(self, mocker: MockerFixtur
         )
         assert engine.single_stage_mode is True
 
-    def test_stage_id_kwarg_sets_filter(self, mocker: MockerFixture):
+    @pytest.mark.parametrize("stage_id", [1, FRONTEND_ONLY_ID_FILTER])
+    def test_stage_id_kwarg_sets_filter(self, mocker: MockerFixture, stage_id):
         engine = self._make_engine_no_thread(
             mocker,
-            stage_id=1,
-            omni_master_address="127.0.0.1",
-            omni_master_port=20002,
-        )
-        assert engine._single_stage_id_filter == 1
+            stage_cfgs=[_make_stage_cfg(0), _make_stage_cfg(1)],
+            stage_id=stage_id,
+        )
+        assert engine._single_stage_id_filter == stage_id
+
+    def test_stage_id_not_in_resolved_stages_raises(self, mocker: MockerFixture):
+        """Ensure a stage_id outside the resolved stages fails before launch."""
+        with pytest.raises(ValueError, match=r"Available stage ids: \[0, 1\]"):
+            self._make_engine_no_thread(
+                mocker,
+                stage_cfgs=[_make_stage_cfg(0), _make_stage_cfg(1)],
+                stage_id=999,
+            )
+
+    def test_frontend_only_head_plans_every_replica_remote(self, mocker: MockerFixture):
+        """Ensure that launching with the frontend placeholder makes every replica run as remote."""
+        stage_plans = []
+
+        def capture_plans_and_stop(runtime, plans):
+            stage_plans.extend(plans)
+            raise RuntimeError("stage plans captured")
+
+        mocker.patch.object(
+            DistStageRuntime, "_start_omni_master_server", autospec=True, side_effect=capture_plans_and_stop
+        )
+        with pytest.raises(RuntimeError, match="stage plans captured"):
+            AsyncOmniEngine(
+                model="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice",
+                stage_id=FRONTEND_ONLY_ID_FILTER,
+                omni_master_address="127.0.0.1",
+                omni_master_port=26000,
+            )
+
+        assert {replica.launch_mode for plan in stage_plans for replica in plan.replicas} == {"remote"}
 
     def test_no_stage_id_no_single_stage_mode(self, mocker: MockerFixture):
         engine = self._make_engine_no_thread(mocker)
@@ -654,6 +685,11 @@ def _build_runtime(self, stage_cfgs: list[Any], *, stage_id_filter: int | None)
             omni_master_port=26000,
         )
 
+    def test_frontend_only_filter_launches_every_stage_remotely(self):
+        """Ensure the frontend-only filter launches no stage locally."""
+        runtime = self._build_runtime([_make_stage_cfg(0), _make_stage_cfg(1)], stage_id_filter=FRONTEND_ONLY_ID_FILTER)
+        assert [runtime._get_launch_mode(stage_id) for stage_id in (0, 1)] == ["remote", "remote"]
+
     def test_build_logical_stage_init_plans_marks_non_matching_stage_remote(self, mocker: MockerFixture):
         import vllm_omni.engine.stage_runtime as runtime_mod
 

tests/entrypoints/test_serve.py

@@ -20,10 +20,12 @@
     normalize_and_validate_diffusion_engine_ingress_kwargs,
 )
 from vllm_omni.config.resolver import OmniConfigResolution
+from vllm_omni.engine.arg_utils import FRONTEND_ONLY_ID_FILTER
 from vllm_omni.engine.stage_engine_startup import StageReplicaResources
 from vllm_omni.engine.stage_runtime import StageEngineLaunch
 from vllm_omni.entrypoints.cli.serve import (
     OmniServeCommand,
+    _parse_stage_id,
     _parse_stage_overrides,
     run_headless,
 )
@@ -144,6 +146,46 @@ def test_serve_accepts_disabled_stage_metrics_when_logging_is_disabled_or_headle
     OmniServeCommand().validate(args)
 
 
+@pytest.mark.parametrize("value, expected", [("0", 0), ("none", FRONTEND_ONLY_ID_FILTER)])
+def test_parse_stage_id_valid_values(value: str, expected: int) -> None:
+    """Ensure --stage-id takes a non-negative int or an explicit 'none'."""
+    assert _parse_stage_id(value) == expected
+
+
+@pytest.mark.parametrize("value", ["-1", "abc"])
+def test_parse_stage_id_invalid_values(value: str) -> None:
+    """Ensure --stage-id rejects negatives and non-integers."""
+    with pytest.raises(argparse.ArgumentTypeError, match="non-negative integer or 'none'"):
+        _parse_stage_id(value)
+
+
+@pytest.mark.parametrize(
+    "extra_argv, match",
+    [
+        (["--headless"], "--stage-id none cannot be used with --headless"),
+        (["--omni-dp-size-local", "2"], "requires a numeric --stage-id"),
+    ],
+)
+def test_serve_validate_rejects_stage_id_none_combinations(extra_argv: list[str], match: str) -> None:
+    """Ensure --stage-id none is rejected where a local stage is required."""
+    args = _parse_serve_args(
+        [
+            "serve",
+            "fake-model",
+            "--omni",
+            "--stage-id",
+            "none",
+            "--omni-master-address",
+            "127.0.0.1",
+            "--omni-master-port",
+            "20000",
+            *extra_argv,
+        ]
+    )
+    with pytest.raises(ValueError, match=match):
+        OmniServeCommand().validate(args)
+
+
 def test_no_guardrails_is_only_forwarded_as_model_config(mocker: MockerFixture) -> None:
     """The CLI alias must not reach the strict diffusion config validator."""
     parser = TrackingArgumentParser()

vllm_omni/engine/arg_utils.py

@@ -20,6 +20,16 @@
 
 logger = init_logger(__name__)
 
+# Stage ids are non-negative, so this filter never matches a stage id; it must never be used as an index.
+# We use this as a sentinel value for --stage-id none to indicate all stages should be launched remotely.
+FRONTEND_ONLY_ID_FILTER = -1
+
+
+def is_frontend_only(single_stage_id_filter: int | None) -> bool:
+    """Return whether the stage filter asks the head to launch no local stage."""
+    return single_stage_id_filter == FRONTEND_ONLY_ID_FILTER
+
+
 # Maps model architecture names to their HuggingFace model_type values.
 # Used when auto-injecting hf_overrides for models with missing config.json.
 _ARCH_TO_MODEL_TYPE: dict[str, str] = {

vllm_omni/engine/omni_engine_base.py

@@ -13,7 +13,7 @@
 import time
 import uuid
 import weakref
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any, cast
@@ -26,6 +26,7 @@
 from vllm.v1.engine.input_processor import InputProcessor
 
 from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override
+from vllm_omni.config.omni_config import StageConfigType
 from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
 from vllm_omni.config.stage_config import (
     _DEPLOY_DIR,
@@ -38,6 +39,7 @@
     parse_attention_config,
 )
 from vllm_omni.diffusion.io_support import get_diffusion_output_type
+from vllm_omni.engine.arg_utils import is_frontend_only
 from vllm_omni.engine.async_engine_utils import (
     SHUTDOWN_ENQUEUE_TIMEOUT_S,
     SHUTDOWN_JOIN_TIMEOUT_S,
@@ -167,23 +169,20 @@ def __init__(
         self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False))
         self._client_config = client_config
 
-        logger.info(f"[OmniEngine] Initializing with model {model}")
+        self._single_stage_id_filter: int | None = kwargs.get("stage_id")
 
+        logger.info(f"[OmniEngine] Initializing with model {model}")
         # ------------------------------------------------------------------ #
         # 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:
+        if self._single_stage_id_filter is not None and not single_stage_mode:
             single_stage_mode = True
         if client_config is not None and int(client_config.get("client_count", 1)) > 1 and single_stage_mode:
             raise ValueError("Multiple API servers cannot be combined with single-stage distributed mode")
-
         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")
 
@@ -273,6 +272,10 @@ def __init__(
             kwargs,
             trust_remote_code=trust_remote_code,
         )
+
+        if self._single_stage_id_filter is not None:
+            self._validate_single_stage_choice(self._single_stage_id_filter, self.stage_configs)
+
         if self._config_resolution is None:
             # Same model, trust_remote_code and deploy path as the resolution
             # above, so reuse it rather than paying for the HF config again.
@@ -377,6 +380,17 @@ def get_diffusion_od_config(self) -> Any:
             )
         return self._diffusion_od_config_view
 
+    def _validate_single_stage_choice(self, stage_id: int, stage_configs: Sequence[StageConfigType]) -> None:
+        """Reject a `stage_id` that is not a stage of the resolved pipeline."""
+        if is_frontend_only(stage_id):
+            return
+        available = [stage_cfg.stage_id for stage_cfg in stage_configs]
+        if stage_id not in available:
+            raise ValueError(
+                f"No stage config found for stage_id={stage_id!r}. Available stage ids: {available}. "
+                "Use --stage-id none to launch no local stage."
+            )
+
     def _initialize_stages(self, stage_init_timeout: int) -> None:
         """Initialize stage clients/processors via StageRuntime and assign to self."""
         self._runtime = create_stage_runtime(

vllm_omni/engine/stage_runtime.py

@@ -30,6 +30,7 @@
     RandomBalancer,
     RoundRobinBalancer,
 )
+from vllm_omni.engine.arg_utils import is_frontend_only
 from vllm_omni.engine.cuda_mps import CudaMPSServer, physical_gpu_uuid
 from vllm_omni.engine.messages import (
     EngineQueueMessage,
@@ -1412,12 +1413,15 @@ def create_membership_controller(self) -> Any | None:
 
     def _prepare_stage_plans(self) -> list[LogicalStageInitPlan]:
         self._validate_single_stage_mode_replica_constraints()
-        return super()._prepare_stage_plans()
+        stage_plans = super()._prepare_stage_plans()
+        if all(replica.launch_mode == "remote" for plan in stage_plans for replica in plan.replicas):
+            logger.info("All stage replicas are remote; running only the frontend (API server + orchestrator)")
+        return stage_plans
 
     def _validate_single_stage_mode_replica_constraints(self) -> None:
         """Apply --omni-dp-size-local to the local stage's runtime.num_replicas."""
         target_stage_id = self._single_stage_id_filter
-        if target_stage_id is None:
+        if target_stage_id is None or is_frontend_only(target_stage_id):
             return
 
         for idx, stage_cfg in enumerate(self._stage_configs):
@@ -1467,7 +1471,8 @@ def _cleanup_distributed_infra(self) -> None:
     # ---- Distributed overrides ----
 
     def _get_launch_mode(self, stage_id: int) -> str:
-        if self._single_stage_id_filter is not None and stage_id != self._single_stage_id_filter:
+        stage_id_filter = self._single_stage_id_filter
+        if is_frontend_only(stage_id_filter) or (stage_id_filter is not None and stage_id != stage_id_filter):
             return "remote"
         return "local"
 

vllm_omni/entrypoints/cli/serve.py

@@ -26,6 +26,7 @@
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.registry import resolve_native_single_file
+from vllm_omni.engine.arg_utils import FRONTEND_ONLY_ID_FILTER, is_frontend_only
 from vllm_omni.entrypoints.api_server_shared_dir import (
     create_api_server_shared_dir,
     remove_api_server_shared_dir,
@@ -76,6 +77,20 @@ def _parse_stage_overrides(value: str) -> dict[str, dict[str, Any]]:
     return parsed
 
 
+def _parse_stage_id(value: str) -> int:
+    """Parse `--stage-id` as a non-negative int, or `none` (no local stage) as `FRONTEND_ONLY_ID_FILTER`."""
+    if value.lower() == "none":
+        return FRONTEND_ONLY_ID_FILTER
+    msg = f"Stage ID must be a non-negative integer or 'none', got {value!r}"
+    try:
+        stage_id = int(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError(msg) from exc
+    if stage_id < 0:
+        raise argparse.ArgumentTypeError(msg)
+    return stage_id
+
+
 def _nonneg_finite_float(value: str) -> float:
     """Argparse type for finite, non-negative floats (rejects nan/inf)."""
     try:
@@ -176,6 +191,8 @@ def validate(self, args: argparse.Namespace) -> None:
 
         if args.stage_id is not None and (args.omni_master_address is None or args.omni_master_port is None):
             raise ValueError("--stage-id requires both --omni-master-address and --omni-master-port to be set")
+        if args.headless and is_frontend_only(args.stage_id):
+            raise ValueError("--stage-id none cannot be used with --headless")
 
         # Require an explicit model under --omni. ``args.model`` always carries
         # vLLM's ModelConfig default (``Qwen/Qwen3-0.6B``), so an omit is silent:
@@ -266,8 +283,8 @@ def validate(self, args: argparse.Namespace) -> None:
         if omni_dp_size_local is not None:
             if omni_dp_size_local < 1:
                 raise ValueError(f"--omni-dp-size-local must be >= 1, got {omni_dp_size_local}")
-            if omni_dp_size_local != 1 and args.stage_id is None:
-                raise ValueError("--omni-dp-size-local != 1 requires --stage-id to be set")
+            if omni_dp_size_local != 1 and (args.stage_id is None or is_frontend_only(args.stage_id)):
+                raise ValueError("--omni-dp-size-local != 1 requires a numeric --stage-id")
 
         # vLLM CLI args that omni does not honor: parallelism comes from the
         # per-stage YAML (parallel_config:, enable_expert_parallel:) and the
@@ -430,9 +447,12 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu
         )
         omni_config_group.add_argument(
             "--stage-id",
-            type=int,
+            type=_parse_stage_id,
             default=None,
-            help="Select and launch a single stage by stage_id.",
+            help=(
+                "Select and launch a single stage by stage_id. Pass 'none' on the "
+                "head to launch no local stage (API server + orchestrator only)."
+            ),
         )
         omni_config_group.add_argument(
             "--replica-id",