Refactor shutdown and cleanup logic in OmniRunner and OmniStage to improve process termination and resource management. The changes ensure that stage workers are stopped before closing queues, preventing potential orphaned subprocesses. Additionally, enhanced error handling during process termination provides clearer logging for failed shutdown attempts. This update aims to increase the reliability and stability of the shutdown process across the application.

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

tests/conftest.py

@@ -1911,6 +1911,7 @@ class OmniRunner:
     def _cleanup_process(self):
         try:
             keywords = ["enginecore"]
+            matched = []
 
             for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
                 try:
@@ -1923,16 +1924,32 @@ class OmniRunner:
 
                     if is_process:
                         print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
-
-                        try:
-                            proc.terminate()
-                            time.sleep(2)
-                        except Exception:
-                            proc.kill()
-
+                        matched.append(proc)
                 except (psutil.NoSuchProcess, psutil.AccessDenied):
                     pass
 
+            for proc in matched:
+                try:
+                    proc.terminate()
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+
+            _, still_alive = psutil.wait_procs(matched, timeout=5)
+            for proc in still_alive:
+                try:
+                    proc.kill()
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+
+            if still_alive:
+                _, stubborn = psutil.wait_procs(still_alive, timeout=3)
+                if stubborn:
+                    print(f"Warning: failed to kill residual vllm pids: {[p.pid for p in stubborn]}")
+                else:
+                    print(f"Force-killed residual vllm pids: {[p.pid for p in still_alive]}")
+            elif matched:
+                print(f"Terminated vllm pids: {[p.pid for p in matched]}")
+
         except Exception as e:
             print(f"Error in psutil vllm cleanup: {e}")
 

vllm_omni/entrypoints/async_omni.py

@@ -24,7 +24,7 @@ from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni import OmniBase
 from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
+from vllm_omni.entrypoints.stage_utils import OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
@@ -51,11 +51,14 @@ def _weak_close_cleanup_async(
         except Exception as e:
             logger.warning("Failed to close inline diffusion engine: %s", e)
     if stage_list:
-        for q in stage_in_queues:
+        # Stop workers first, then close queues. Closing queues early can
+        # prevent workers from receiving shutdown tasks.
+        for stage in stage_list:
             try:
-                q.put_nowait(SHUTDOWN_TASK)
+                stage.stop_stage_worker()
             except Exception as e:
-                logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
+                logger.warning(f"Failed to stop stage worker: {e}")
+        for q in stage_in_queues:
             close_fn = getattr(q, "close", None)
             if callable(close_fn):
                 close_fn()
@@ -63,11 +66,6 @@ def _weak_close_cleanup_async(
             close_fn = getattr(q, "close", None)
             if callable(close_fn):
                 close_fn()
-        for stage in stage_list:
-            try:
-                stage.stop_stage_worker()
-            except Exception as e:
-                logger.warning(f"Failed to stop stage worker: {e}")
     try_close_ray(ray_pg)
     # Cancel output handler
     if output_handler is not None:

vllm_omni/entrypoints/omni.py

@@ -38,7 +38,7 @@ from vllm_omni.distributed.ray_utils.utils import (
 )
 from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker
 from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
+from vllm_omni.entrypoints.stage_utils import OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     filter_dataclass_kwargs,
@@ -72,11 +72,14 @@ def _weak_close_cleanup(
 ):
     """Weak reference cleanup function for OmniBase instances."""
     if stage_list:
-        for q in stage_in_queues:
+        # Stop workers first, then close queues. Closing queues too early can
+        # drop shutdown signals and leave engine subprocesses orphaned.
+        for stage in stage_list:
             try:
-                q.put_nowait(SHUTDOWN_TASK)
+                stage.stop_stage_worker()
             except Exception as e:
-                logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
+                logger.warning(f"Failed to stop stage worker: {e}")
+        for q in stage_in_queues:
             close_fn = getattr(q, "close", None)
             if callable(close_fn):
                 close_fn()
@@ -84,11 +87,6 @@ def _weak_close_cleanup(
             close_fn = getattr(q, "close", None)
             if callable(close_fn):
                 close_fn()
-        for stage in stage_list:
-            try:
-                stage.stop_stage_worker()
-            except Exception as e:
-                logger.warning(f"Failed to stop stage worker: {e}")
     try_close_ray(ray_pg)
 
     # Gracefully shutdown handshake server thread

vllm_omni/entrypoints/omni_stage.py

@@ -616,28 +616,17 @@ class OmniStage:
         If graceful shutdown fails, forcefully terminates the process.
         Handles both multiprocessing Process and Ray Actor.
         """
-        if self._in_q is not None:
+        in_q = self._in_q
+        out_q = self._out_q
+
+        if in_q is not None:
             try:
-                self._in_q.put_nowait(SHUTDOWN_TASK)
+                # Prefer a blocking put to reduce the chance that shutdown
+                # is dropped when queue/socket buffers are momentarily full.
+                in_q.put(SHUTDOWN_TASK)
             except Exception as e:
-                # Queue may already be closed by the weak-ref cleanup
-                # (e.g. ZmqQueue socket closed) — this is expected.
+                # Queue may already be closed by cleanup paths.
                 logger.debug("Failed to send shutdown to in_q: %s", e)
-            close_fn = getattr(self._in_q, "close", None)
-            if callable(close_fn):
-                try:
-                    close_fn()
-                except Exception:
-                    pass
-            self._in_q = None
-        if self._out_q is not None:
-            close_fn = getattr(self._out_q, "close", None)
-            if callable(close_fn):
-                try:
-                    close_fn()
-                except Exception:
-                    pass
-            self._out_q = None
 
         if self._ray_actor is not None:
             kill_ray_actor(self._ray_actor)
@@ -645,7 +634,7 @@ class OmniStage:
             self._ray_task_ref = None
         elif self._proc is not None:
             try:
-                self._proc.join(timeout=5)
+                self._proc.join(timeout=10)
             except Exception as e:
                 logger.debug("join() failed: %s", e)
             if self._proc.is_alive():
@@ -653,6 +642,36 @@ class OmniStage:
                     self._proc.terminate()
                 except Exception as e:
                     logger.warning("terminate() failed: %s", e)
+                try:
+                    self._proc.join(timeout=5)
+                except Exception as e:
+                    logger.debug("join() after terminate failed: %s", e)
+            if self._proc.is_alive():
+                try:
+                    self._proc.kill()
+                except Exception as e:
+                    logger.warning("kill() failed: %s", e)
+                try:
+                    self._proc.join(timeout=3)
+                except Exception as e:
+                    logger.debug("join() after kill failed: %s", e)
+
+        if in_q is not None:
+            close_fn = getattr(in_q, "close", None)
+            if callable(close_fn):
+                try:
+                    close_fn()
+                except Exception:
+                    pass
+        if out_q is not None:
+            close_fn = getattr(out_q, "close", None)
+            if callable(close_fn):
+                try:
+                    close_fn()
+                except Exception:
+                    pass
+        self._in_q = None
+        self._out_q = None
 
     def submit(self, payload: dict[str, Any]) -> None:
         """Submit a task to the stage worker.