[Bugfix] Restrict OmniRunner cleanup to owned EngineCore workers

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

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

@@ -216,40 +216,44 @@ def test_video_chat_stream_invalid_requests(
 # WS /v1/realtime
 # ─────────────────────────────────────────────────────────────────────────────
 
+# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake).
+_REALTIME_WS_MODEL_MISMATCH = object()
+_REALTIME_WS_INVALID_AUDIO_APPEND = object()
+
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize(
     "send_frames_spec, ws_error_code, err_message",
     [
-        pytest.param("{not-json", "invalid_event", "Invalid or unrecognized client event", id="invalid_json"),
+        pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"),
         pytest.param(
             json.dumps({"type": "session.update"}),
             "invalid_event",
-            "Invalid or unrecognized client event",
-            id="session_update_missing_session",
+            ("model", "Missing required field"),
+            id="session_update_missing_model",
         ),
         pytest.param(
-            json.dumps({"type": "session.update", "session": []}),
-            "invalid_event",
-            "Invalid or unrecognized client event",
-            id="session_update_invalid_session",
+            _REALTIME_WS_MODEL_MISMATCH,
+            "model_not_found",
+            ("model", "does not exist"),
+            id="session_update_model_mismatch",
         ),
         pytest.param(
-            json.dumps({"type": "input_audio_buffer.commit"}),
-            "invalid_request_error",
-            "Input audio buffer is empty",
-            id="commit_without_audio",
+            json.dumps({"type": "input_audio_buffer.commit", "final": False}),
+            "model_not_validated",
+            ("session.update", "validate the model"),
+            id="commit_without_model_validation",
         ),
         pytest.param(
             json.dumps({"type": "unknown.event"}),
-            "invalid_event",
-            "Invalid or unrecognized client event",
+            "unknown_event",
+            "Unknown event type",
             id="unknown_event_type",
         ),
         pytest.param(
-            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
-            "invalid_request_error",
-            "Invalid base64 audio data",
+            _REALTIME_WS_INVALID_AUDIO_APPEND,
+            "invalid_audio",
+            "Invalid audio data",
             id="invalid_audio_append",
         ),
     ],
@@ -262,12 +266,27 @@ def test_realtime_invalid_requests(
     ws_error_code: str,
     err_message: str | tuple[str, ...],
 ) -> None:
+    send_frames: Any
+    if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
+        send_frames = json.dumps(
+            {
+                "type": "session.update",
+                "model": "this-model-is-not-served-on-this-instance",
+            }
+        )
+    elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND:
+        send_frames = [
+            json.dumps({"type": "session.update", "model": omni_server.model}),
+            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
+        ]
+    else:
+        send_frames = send_frames_spec
     online_client.send_realtime_ws_request(
         {
-            "send_frames": send_frames_spec,
+            "send_frames": send_frames,
             "timeout": 120,
             "ws_max_size": None,
-            "ws_skip_types": ["session.created", "conversation.created"],
+            "ws_skip_types": ["session.created"],
             "ws_json_type": "error",
             "ws_error_code": ws_error_code,
             "err_message": err_message,

tests/entrypoints/test_realtime_connection_helpers.py

@@ -4,10 +4,7 @@
 
 from __future__ import annotations
 
-import asyncio
 import base64
-import json
-from types import SimpleNamespace
 
 import numpy as np
 import pytest
@@ -87,56 +84,3 @@ class TestAsyncOmniStreamingParamsValidation:
         p = SamplingParams(n=1, stop=["\n"], output_kind=RequestOutputKind.DELTA)
         with pytest.raises(ValueError, match="Input streaming"):
             AsyncOmni._validate_streaming_input_sampling_params(p)
-
-
-@pytest.mark.parametrize(
-    "frame,code,message",
-    [
-        ("{not-json", "invalid_event", "Invalid or unrecognized client event"),
-        (json.dumps({"type": "session.update"}), "invalid_event", "Invalid or unrecognized client event"),
-        (
-            json.dumps({"type": "session.update", "session": []}),
-            "invalid_event",
-            "Invalid or unrecognized client event",
-        ),
-        (json.dumps({"type": "input_audio_buffer.commit"}), "invalid_request_error", "Input audio buffer is empty"),
-        (json.dumps({"type": "unknown.event"}), "invalid_event", "Invalid or unrecognized client event"),
-        (
-            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
-            "invalid_request_error",
-            "Invalid base64 audio data",
-        ),
-    ],
-)
-def test_full_duplex_invalid_request_contract(frame, code, message):
-    from starlette.websockets import WebSocketDisconnect
-
-    from vllm_omni.entrypoints.openai.realtime.connection import OpenAIFullDuplexConnection
-
-    class Socket:
-        def __init__(self):
-            self.sent = []
-            self.frames = iter([frame])
-
-        async def accept(self):
-            pass
-
-        async def receive_text(self):
-            value = next(self.frames, None)
-            if value is None:
-                raise WebSocketDisconnect()
-            return value
-
-        async def send_text(self, payload):
-            self.sent.append(json.loads(payload))
-
-    async def run():
-        socket = Socket()
-        connection = OpenAIFullDuplexConnection(socket, SimpleNamespace(), "test-model", None)
-        await connection.handle_connection()
-        assert [event["type"] for event in socket.sent] == ["session.created", "conversation.created", "error"]
-        error = socket.sent[-1]["error"]
-        assert error["code"] == code
-        assert error["message"] == message
-
-    asyncio.run(run())

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