[Rebase] Align vLLM-Omni with vLLM v0.31.0 (#8459)

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

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()
 
 
 # ---------------------------------------------------------------------------