[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/diffusion/test_inline_stage_diffusion_client.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -94,3 +95,24 @@ def test_inline_shutdown(client, mock_engine):
 
     assert client._shutting_down
     mock_engine.close.assert_called_once()
+
+
+def test_inline_client_requires_replica_id(mock_engine):
+    metadata = SimpleNamespace(
+        stage_id=0,
+        final_output=True,
+        final_output_type="image",
+        default_sampling_params={},
+        requires_multimodal_data=False,
+        custom_process_input_func=None,
+        engine_input_source=[],
+    )
+    with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+        od_config = MagicMock(spec=OmniDiffusionConfig)
+        with pytest.raises(AttributeError, match="replica_id"):
+            InlineStageDiffusionClient(
+                model="test_model",
+                od_config=od_config,
+                metadata=metadata,
+                batch_size=1,
+            )

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -621,6 +621,29 @@ class TestStageDiffusionClientErrorPropagation:
         assert client._shutting_down is True
         assert client._engine_dead is True
 
+    def test_initialize_client_requires_replica_id(self):
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        metadata = SimpleNamespace(
+            stage_id=0,
+            final_output=True,
+            final_output_type="image",
+            default_sampling_params=None,
+            requires_multimodal_data=False,
+            custom_process_input_func=None,
+            engine_input_source=[],
+        )
+
+        with pytest.raises(AttributeError, match="replica_id"):
+            client._initialize_client(
+                metadata,
+                "tcp://req",
+                "tcp://resp",
+                proc=None,
+                batch_size=1,
+            )
+
 
 # ───────── monitor thread & death sentinel integration tests ─────────