[FEAT] support multi-stage deployment (#2396)

Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: Zheng Wengang <zwg0606@gmail.com>
Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Signed-off-by: yinpe <11810305@mail.sustech.edu.cn>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: yinpe <11810305@mail.sustech.edu.cn>
Co-authored-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Chenguang Zheng <645327136@qq.com>

tests/entrypoints/test_serve.py

@@ -38,6 +38,7 @@ def _make_headless_args() -> argparse.Namespace:
     return argparse.Namespace(
         model="fake-model",
         stage_id=3,
+        replica_id=0,
         omni_master_address="127.0.0.1",
         omni_master_port=26000,
         api_server_count=0,
@@ -45,6 +46,7 @@ def _make_headless_args() -> argparse.Namespace:
         stage_configs_path=None,
         log_stats=False,
         disable_log_stats=False,
+        stage_init_timeout=600,
     )
 
 
@@ -59,6 +61,7 @@ def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker
         node_rank_within_dp=0,
     )
     vllm_config = mocker.Mock(parallel_config=parallel_config)
+    vllm_config.needs_dp_coordinator = False
     executor_class = mocker.Mock()
     engine_manager = mocker.Mock()
 
@@ -99,6 +102,7 @@ def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker
         omni_stage_id=3,
         omni_stage_config=stage_cfg,
         coordinator=None,
+        replica_id=0,
     )
     mock_manager_cls.assert_called_once()
     manager_kwargs = mock_manager_cls.call_args.kwargs
@@ -155,13 +159,56 @@ def test_run_headless_honors_explicit_log_stats_flag(mocker: MockerFixture) -> N
     assert manager_kwargs["log_stats"] is True
 
 
+def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None:
+    args = _make_headless_args()
+    args.replica_id = 2
+    stage_cfg = mocker.Mock(stage_id=3)
+    parallel_config = mocker.Mock(
+        data_parallel_size_local=1,
+        data_parallel_rank=0,
+        data_parallel_rank_local=0,
+        node_rank_within_dp=0,
+    )
+    vllm_config = mocker.Mock(parallel_config=parallel_config)
+    vllm_config.needs_dp_coordinator = False
+    engine_manager = mocker.Mock()
+
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        return_value=("/fake/stages.yaml", [stage_cfg]),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
+    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
+    mocker.patch(
+        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+        return_value=(None, None, None),
+    )
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.build_vllm_config",
+        return_value=(vllm_config, mocker.Mock()),
+    )
+    mock_register = mocker.patch(
+        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+        return_value="tcp://127.0.0.1:26001",
+    )
+    mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
+    mocker.patch("signal.signal")
+
+    run_headless(args)
+
+    assert mock_register.call_args.kwargs["replica_id"] == 2
+
+
 def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFixture) -> None:
     args = _make_headless_args()
+    args.replica_id = 1
     stage_cfg = mocker.Mock(stage_id=3, stage_type="diffusion")
     stage_cfg.engine_args = mocker.Mock()
     stage_cfg.engine_input_source = []
     stage_cfgs = [stage_cfg]
-    metadata = mocker.Mock(stage_id=3)
+    metadata = mocker.Mock(stage_id=3, stage_type="diffusion")
     od_config = mocker.Mock()
     proc = mocker.Mock()
     proc.exitcode = 0
@@ -199,6 +246,7 @@ def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFix
         omni_stage_id=3,
         omni_stage_config=stage_cfg,
         return_addresses=True,
+        replica_id=1,
     )
     mock_spawn.assert_called_once_with(
         "fake-model",
@@ -207,5 +255,5 @@ def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFix
         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")
+    mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001", 600)
     proc.join.assert_called_once_with()

tests/entrypoints/test_utils.py

@@ -360,6 +360,47 @@ class TestLoadAndResolveStageConfigs:
         assert len(stage_configs) == 1
         assert "dtype" in stage_configs[0]["engine_args"]
 
+    def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas(
+        self, tmp_path, mocker: MockerFixture
+    ):
+        deploy_path = tmp_path / "qwen3_multi.yaml"
+        deploy_path.write_text(
+            'stages:\n  - stage_id: 0\n    devices: "0"\n  - stage_id: 1\n    devices: "1,2,3"\n    num_replicas: 3\n',
+            encoding="utf-8",
+        )
+
+        returned_stage_configs = [
+            create_config({"stage_id": 0, "runtime": {"devices": "0"}, "engine_args": {"model": "dummy"}}),
+            create_config(
+                {
+                    "stage_id": 1,
+                    "runtime": {"devices": "1,2,3", "num_replicas": 3},
+                    "engine_args": {"model": "dummy"},
+                }
+            ),
+        ]
+        load_stage_configs = mocker.patch(
+            "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+            return_value=returned_stage_configs,
+        )
+
+        config_path, stage_configs = load_and_resolve_stage_configs(
+            model="dummy-model",
+            stage_configs_path=str(deploy_path),
+            kwargs={},
+        )
+
+        load_stage_configs.assert_called_once_with(
+            "dummy-model",
+            base_engine_args={},
+            deploy_config_path=str(deploy_path),
+            stage_overrides=None,
+        )
+        assert config_path == str(deploy_path)
+        assert len(stage_configs) == 2
+        assert stage_configs[1].runtime.num_replicas == 3
+        assert stage_configs[1].runtime.devices == "1,2,3"
+
 
 class TestLoadStageConfigsFromYaml:
     """Regression tests for stage-config loading and merging."""