[Bugfix] Restrict OmniRunner cleanup to owned EngineCore workers

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/diffusion/distributed/test_comm.py

@@ -14,7 +14,6 @@ from typing import Literal
 import pytest
 import torch
 
-from tests.helpers.distributed import configure_cpu_collective_worker, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -57,8 +56,6 @@ def _init_worker(
             "MASTER_PORT": str(master_port),
         }
     )
-    if device_kind == "cpu":
-        configure_cpu_collective_worker()
     backend = "gloo" if device_kind == "cpu" else None
     init_distributed_environment(backend=backend)
 
@@ -230,10 +227,6 @@ def _run_ring_p2p(
             atol=atol,
             msg=f"[Rank {local_rank}] Ring P2P data mismatch",
         )
-        if device_kind == "cpu":
-            # P2P completion only synchronizes neighbors. Keep fast CPU ranks
-            # alive until every rank has finished constructing and using groups.
-            torch.distributed.barrier()
     finally:
         destroy_distributed_env()
 
@@ -250,7 +243,7 @@ def _spawn_4d_identity(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    start_collective_workers(
+    torch.multiprocessing.spawn(
         _run_4d_identity,
         args=(
             world_size,
@@ -264,7 +257,6 @@ def _spawn_4d_identity(
             master_port,
         ),
         nprocs=world_size,
-        device_kind=device_kind,
     )
 
 
@@ -280,7 +272,7 @@ def _spawn_5d_identity(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    start_collective_workers(
+    torch.multiprocessing.spawn(
         _run_5d_identity,
         args=(
             world_size,
@@ -294,7 +286,6 @@ def _spawn_5d_identity(
             master_port,
         ),
         nprocs=world_size,
-        device_kind=device_kind,
     )
 
 
@@ -308,11 +299,10 @@ def _spawn_ring_p2p(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    start_collective_workers(
+    torch.multiprocessing.spawn(
         _run_ring_p2p,
         args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
         nprocs=world_size,
-        device_kind=device_kind,
     )
 
 
@@ -509,9 +499,8 @@ def _run_fused_qkv_matches_three_4d(
 @pytest.mark.cpu
 @pytest.mark.parametrize("world_size", [2, 4])
 def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
-    start_collective_workers(
+    torch.multiprocessing.spawn(
         _run_fused_qkv_matches_three_4d,
         args=(world_size, "cpu", 29613),
         nprocs=world_size,
-        device_kind="cpu",
     )

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -19,7 +19,6 @@ from vllm.utils.network_utils import get_file_store_init_method
 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
 
 import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
-from tests.helpers.distributed import collective_context, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -40,25 +39,6 @@ DeviceKind = Literal["cpu", "cuda"]
 _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
-@pytest.fixture(autouse=True)
-def _restore_collective_platforms():
-    """CPU baselines run in the parent process; retain its original platforms."""
-    import vllm.platforms
-
-    from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
-
-    original_platform = vllm.platforms.current_platform
-    original_omni_platforms = {
-        module: module.current_omni_platform for module in (comm, group_coordinator, parallel_state)
-    }
-    try:
-        yield
-    finally:
-        vllm.platforms.current_platform = original_platform
-        for module, platform in original_omni_platforms.items():
-            module.current_omni_platform = platform
-
-
 def _cleanup_distributed() -> None:
     """Destroy omni distributed state left dangling after a test.
 
@@ -429,10 +409,6 @@ def init_dist(
     device = _worker_device(local_rank, device_kind)
     if device_kind == "cuda":
         current_omni_platform.set_device(device)
-    else:
-        from tests.helpers.distributed import configure_cpu_collective_worker
-
-        configure_cpu_collective_worker()
     backend = "gloo" if device_kind == "cpu" else None
     init_distributed_environment(
         world_size=world_size,
@@ -524,14 +500,13 @@ def isend_irecv_worker(
 
 
 def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None:
-    mp_context = collective_context(device_kind)
+    mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     q = manager.Queue()
-    start_collective_workers(
+    torch.multiprocessing.spawn(
         isend_irecv_worker,
         args=(pp_size, init_method, device_kind, q),
         nprocs=pp_size,
-        device_kind=device_kind,
     )
     results = {label: tensor for label, tensor in [q.get(), q.get()]}
     torch.testing.assert_close(
@@ -681,16 +656,15 @@ def _run_predict_noise(
 
     baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
 
-    mp_context = collective_context(device_kind)
+    mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     pp_q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    start_collective_workers(
+    torch.multiprocessing.spawn(
         predict_noise_worker,
         args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
         nprocs=world_size,
-        device_kind=device_kind,
     )
 
     pp_out = pp_q.get()
@@ -926,16 +900,15 @@ def _run_scheduler_step(
 
     baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
 
-    mp_context = collective_context(device_kind)
+    mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    start_collective_workers(
+    torch.multiprocessing.spawn(
         scheduler_step_worker,
         args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
         nprocs=world_size,
-        device_kind=device_kind,
     )
 
     result = q.get()

tests/helpers/cpu_collective_server.py

@@ -1,19 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Import preload for the CPU-only collective forkserver.
-
-This module is imported in the forkserver process, before it starts workers.
-Its CUDA visibility does not alter the pytest process or CUDA parity workers.
-"""
-
-import os
-
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-
-import torch  # noqa: E402
-
-# Retain the expensive worker-module imports across CPU process groups.
-from tests.diffusion.distributed import test_comm, test_pipeline_parallel  # noqa: E402, F401
-
-if torch.cuda.is_initialized():
-    raise RuntimeError("CPU collective forkserver must not initialize CUDA")

tests/helpers/distributed.py

@@ -1,62 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Process and platform setup for isolated collective test workers."""
-
-from collections.abc import Callable
-from multiprocessing.context import BaseContext
-
-import torch
-
-
-def collective_context(device_kind: str) -> BaseContext:
-    """Reuse imports in a clean CPU forkserver; CUDA workers use spawn."""
-    if device_kind == "cpu" and "forkserver" in torch.multiprocessing.get_all_start_methods():
-        torch.multiprocessing.set_forkserver_preload(["tests.helpers.cpu_collective_server"])
-        return torch.multiprocessing.get_context("forkserver")
-    return torch.multiprocessing.get_context("spawn")
-
-
-def start_collective_workers(
-    worker: Callable[..., None], *, args: tuple[object, ...], nprocs: int, device_kind: str
-) -> None:
-    context = collective_context(device_kind)
-    if context.get_start_method() == "spawn":
-        torch.multiprocessing.spawn(worker, args=args, nprocs=nprocs)
-    else:
-        torch.multiprocessing.start_processes(worker, args=args, nprocs=nprocs, start_method=context.get_start_method())
-
-
-def configure_cpu_collective_worker() -> None:
-    """Select CPU platforms inside a spawned worker that only uses Gloo.
-
-    Upstream coordinators resolve local ranks through the selected
-    platform's visible devices. CPU ranks can outnumber visible GPUs, so they
-    must use the CPU platform even when the test host has CUDA installed.
-    The worker still creates real process groups and runs real collectives.
-    """
-    import vllm.platforms
-    from vllm.platforms.cpu import CpuPlatform
-
-    from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
-    from vllm_omni.platforms.interface import UnspecifiedOmniPlatform
-
-    class CpuCollectivePlatform(UnspecifiedOmniPlatform):
-        dist_backend = "gloo"
-
-        @classmethod
-        def get_device_count(cls) -> int:
-            # CPU tensors share one logical device regardless of process rank.
-            return 1
-
-        @classmethod
-        def set_device(cls, device: torch.device) -> None:
-            assert device.type == "cpu"
-
-        @classmethod
-        def synchronize(cls) -> None:
-            pass
-
-    vllm.platforms.current_platform = CpuPlatform()
-    platform = CpuCollectivePlatform()
-    for module in (comm, group_coordinator, parallel_state):
-        module.current_omni_platform = platform

tests/helpers/runtime.py

@@ -659,6 +659,14 @@ class OmniRunner:
         self.seed = seed
         self._prompt_len_estimate_cache: dict[str, Any] = {}
         self.omni: Any = None
+        self._runner_process = psutil.Process()
+        self._initial_child_ids: set[tuple[int, float]] = set()
+        self._owned_engine_processes: dict[int, float] = {}
+        for child in self._runner_process.children(recursive=True):
+            try:
+                self._initial_child_ids.add((child.pid, child.create_time()))
+            except psutil.Error:
+                pass
         try:
             from vllm_omni.entrypoints.omni import Omni
 
@@ -670,12 +678,14 @@ class OmniRunner:
                 deploy_config=deploy_config,
                 **kwargs,
             )
+            self._remember_engine_processes(startup=True)
             startup_s = time.perf_counter() - startup_t0
             if log_stats:
                 print(f"OmniRunner startup took {startup_s:.3f}s (model={model_name})", flush=True)
         except BaseException:
             # ``with OmniRunner(...)`` never reaches ``__enter__``/``__exit__``
             # when construction fails after worker processes have started.
+            self._remember_engine_processes(startup=True)
             self.__exit__(None, None, None)
             raise
 
@@ -899,18 +909,51 @@ class OmniRunner:
     def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
         return self.omni.stop_profile(stages=stages)
 
+    def _remember_engine_processes(self, *, startup: bool = False) -> None:
+        """Track startup workers, then only descendants of those workers.
+
+        Existing children may belong to another runner or a shared transcriber.
+        Record identities before close() can reparent surviving workers.
+        """
+        try:
+            if startup:
+                candidates = self._runner_process.children(recursive=True)
+            else:
+                candidates = []
+                for pid, created in list(self._owned_engine_processes.items()):
+                    try:
+                        proc = psutil.Process(pid)
+                        if proc.create_time() == created:
+                            candidates.extend(proc.children(recursive=True))
+                    except psutil.Error:
+                        pass
+            for proc in candidates:
+                try:
+                    identity = (proc.pid, proc.create_time())
+                    if identity in self._initial_child_ids:
+                        continue
+                    if startup and any(
+                        (parent.pid, parent.create_time()) in self._initial_child_ids for parent in proc.parents()
+                    ):
+                        continue
+                    cmdline = " ".join(proc.cmdline()).lower()
+                    if "enginecore" in cmdline or "enginecore" in proc.name().lower():
+                        self._owned_engine_processes[identity[0]] = identity[1]
+                except psutil.Error:
+                    pass
+        except psutil.Error as error:
+            print(f"Warning: could not inspect runner workers: {error}")
+
     def _cleanup_process(self):
         try:
-            keywords = ["enginecore"]
             matched = []
-            for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
+            for pid, created in self._owned_engine_processes.items():
                 try:
-                    cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else ""
-                    name = proc.name().lower()
-                    if any(k in cmdline for k in keywords) or any(k in name for k in keywords):
-                        print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
+                    proc = psutil.Process(pid)
+                    if proc.create_time() == created:
+                        print(f"Found owned vllm process: PID={proc.pid}")
                         matched.append(proc)
-                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                except psutil.Error:
                     pass
             for proc in matched:
                 try:
@@ -938,11 +981,14 @@ class OmniRunner:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        omni = getattr(self, "omni", None)
-        if omni is not None and hasattr(omni, "close"):
-            omni.close()
-        self._cleanup_process()
-        cleanup_test_environment()
+        self._remember_engine_processes()
+        try:
+            omni = getattr(self, "omni", None)
+            if omni is not None and hasattr(omni, "close"):
+                omni.close()
+        finally:
+            self._cleanup_process()
+            cleanup_test_environment()
 
 
 # ---------------------------------------------------------------------------

tests/helpers/tests/test_runtime_omni_runner.py

@@ -1,10 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""Tests for OmniRunner startup rollback."""
+"""Tests for OmniRunner startup rollback and worker ownership."""
 
 from __future__ import annotations
 
+import subprocess
+import sys
+from collections.abc import Generator
+
+import psutil
 import pytest
 
 from tests.helpers.runtime import OmniRunner
@@ -38,3 +43,114 @@ def test_runner_init_rolls_back_on_omni_startup_failure(monkeypatch: pytest.Monk
     assert cleaned.count("process") == 1
     # Once at the start of ``__init__``, once from the constructor rollback.
     assert cleaned.count("env") == 2
+
+
+@pytest.fixture
+def live_workers() -> Generator[list[subprocess.Popen[bytes]], None, None]:
+    workers: list[subprocess.Popen[bytes]] = []
+    try:
+        yield workers
+    finally:
+        for worker in workers:
+            if worker.poll() is None:
+                worker.kill()
+            worker.wait(timeout=10)
+            if worker.stdout is not None:
+                worker.stdout.close()
+
+
+def _start_engine_worker(
+    workers: list[subprocess.Popen[bytes]], *, ignore_terminate: bool = False
+) -> subprocess.Popen[bytes]:
+    script = "import signal, time; "
+    if ignore_terminate:
+        script += "signal.signal(signal.SIGTERM, signal.SIG_IGN); "
+    script += "print('ready', flush=True); time.sleep(90)"
+    worker = subprocess.Popen([sys.executable, "-c", script, "enginecore"], stdout=subprocess.PIPE)
+    workers.append(worker)
+    assert worker.stdout is not None
+    assert worker.stdout.readline() == b"ready\n"
+    return worker
+
+
+@pytest.mark.parametrize("close_raises", [False, True])
+def test_runner_cleanup_preserves_unrelated_engine_workers(
+    monkeypatch: pytest.MonkeyPatch,
+    live_workers: list[subprocess.Popen[bytes]],
+    close_raises: bool,
+) -> None:
+    existing = _start_engine_worker(live_workers)
+    owned: list[subprocess.Popen[bytes]] = []
+    late_unrelated: list[subprocess.Popen[bytes]] = []
+    monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
+
+    class _OwnedOmni:
+        def __init__(self, *args: object, **kwargs: object) -> None:
+            owned.append(_start_engine_worker(live_workers))
+
+        def close(self) -> None:
+            late_unrelated.append(_start_engine_worker(live_workers))
+            if close_raises:
+                raise RuntimeError("close failed")
+
+    monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _OwnedOmni)
+    if close_raises:
+        with pytest.raises(RuntimeError, match="close failed"), OmniRunner("fake-model"):
+            pass
+    else:
+        with OmniRunner("fake-model"):
+            pass
+
+    assert owned[0].poll() is not None
+    assert existing.poll() is None
+    assert late_unrelated[0].poll() is None
+
+
+def test_runner_startup_failure_cleans_only_new_workers(
+    monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]]
+) -> None:
+    unrelated = _start_engine_worker(live_workers)
+    owned: list[subprocess.Popen[bytes]] = []
+    monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
+
+    class _FailingOmni:
+        def __init__(self, *args: object, **kwargs: object) -> None:
+            owned.append(_start_engine_worker(live_workers))
+            raise RuntimeError("startup failed")
+
+    monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _FailingOmni)
+    with pytest.raises(RuntimeError, match="startup failed"):
+        OmniRunner("fake-model")
+
+    assert owned[0].poll() is not None
+    assert unrelated.poll() is None
+
+
+def test_runner_cleanup_rejects_stale_process_identity(live_workers: list[subprocess.Popen[bytes]]) -> None:
+    unrelated = _start_engine_worker(live_workers)
+    runner = object.__new__(OmniRunner)
+    runner._owned_engine_processes = {unrelated.pid: psutil.Process(unrelated.pid).create_time() + 1}
+
+    runner._cleanup_process()
+
+    assert unrelated.poll() is None
+
+
+def test_runner_cleanup_kills_owned_worker_ignoring_terminate(
+    monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]]
+) -> None:
+    owned: list[subprocess.Popen[bytes]] = []
+    monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
+
+    class _StubbornOmni:
+        def __init__(self, *args: object, **kwargs: object) -> None:
+            owned.append(_start_engine_worker(live_workers, ignore_terminate=True))
+
+        def close(self) -> None:
+            pass
+
+    monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _StubbornOmni)
+    with OmniRunner("fake-model"):
+        assert owned[0].poll() is None
+
+    assert owned[0].poll() is not None