diff --git a/tests/conftest.py b/tests/conftest.py index 0a3c350d7..e65df3d49 100644 --- a/tests/conftest.py +++ b/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}") diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 474834d16..d2501f0a5 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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: diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 488b986d8..c91f4a7e1 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 1f14d0ab4..e556782ca 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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.