[Refactor] Omni Stage Runtime and Distributed Replica Control Plane (#3855)

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Chenguang ZHENG <645327136@qq.com>

tests/engine/test_orchestrator.py

@@ -11,11 +11,9 @@ from types import SimpleNamespace
 from typing import Any
 
 import janus
-import psutil
 import pytest
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
-from vllm.v1.engine.core_client import AsyncMPClient
 
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
@@ -27,7 +25,6 @@ from vllm_omni.engine.messages import (
     StageSubmissionMessage,
 )
 from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
-from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
 from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -180,42 +177,6 @@ class FakeOutputProcessor:
         return None
 
 
-class _FakeProc:
-    pid = 1234
-
-    def __init__(self):
-        self.terminated = False
-        self.killed = False
-        self.join_calls = []
-
-    def is_alive(self):
-        return not self.terminated and not self.killed
-
-    def terminate(self):
-        self.terminated = True
-
-    def kill(self):
-        self.killed = True
-
-    def join(self, timeout=None):
-        self.join_calls.append(timeout)
-
-
-class _FakeChildProc:
-    def __init__(self):
-        self.terminated = False
-        self.killed = False
-
-    def is_running(self):
-        return not self.terminated and not self.killed
-
-    def terminate(self):
-        self.terminated = True
-
-    def kill(self):
-        self.killed = True
-
-
 def _sampling_params(max_tokens: int = 4) -> SamplingParams:
     return SamplingParams(max_tokens=max_tokens)
 
@@ -437,61 +398,6 @@ async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, requ
     orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=request_ids))
 
 
-def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fails(monkeypatch):
-    fake_proc = _FakeProc()
-    fake_child = _FakeChildProc()
-
-    class FakePsutilProcess:
-        def __init__(self, pid):
-            assert pid == fake_proc.pid
-
-        def children(self, recursive=True):
-            assert recursive
-            return [fake_child]
-
-    def fail_base_shutdown(self, **kwargs):
-        raise RuntimeError("base shutdown failed")
-
-    monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
-    monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: (list(procs), []))
-    monkeypatch.setattr(AsyncMPClient, "shutdown", fail_base_shutdown)
-
-    client = object.__new__(StageEngineCoreClient)
-    client._proc = fake_proc
-
-    with pytest.raises(RuntimeError, match="base shutdown failed"):
-        client.shutdown()
-
-    assert fake_proc.terminated
-    assert fake_proc.join_calls == [5]
-    assert fake_child.terminated
-
-
-def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch):
-    fake_proc = _FakeProc()
-    fake_child = _FakeChildProc()
-
-    class FakePsutilProcess:
-        def __init__(self, pid):
-            assert pid == fake_proc.pid
-
-        def children(self, recursive=True):
-            assert recursive
-            return [fake_child]
-
-    monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
-    monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs)))
-    monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self, **kwargs: None)
-
-    client = object.__new__(StageEngineCoreClient)
-    client._proc = fake_proc
-
-    client.shutdown()
-
-    assert fake_child.terminated
-    assert fake_child.killed
-
-
 @pytest.fixture
 def orchestrator_factory():
     fixtures: list[OrchestratorFixture] = []