diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 5fedd17dd..2475431dd 100644 --- a/tests/helpers/runtime.py +++ b/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() # --------------------------------------------------------------------------- diff --git a/tests/helpers/tests/test_runtime_omni_runner.py b/tests/helpers/tests/test_runtime_omni_runner.py index aaada20f2..bcd29493d 100644 --- a/tests/helpers/tests/test_runtime_omni_runner.py +++ b/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