[Bugfix] Ensure stage and diffusion subprocesses exit when parent dies unexpectedly (#3751)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>

tests/core/sched/test_generation_scheduler_restore.py

@@ -28,7 +28,7 @@ class FakeAdapter:
             req = running.pop()
             self.waiting_for_chunk_running_requests.append(req)
 
-    def restore_queues(self, waiting, running):
+    def restore_queues(self, waiting, running, scheduler_requests=None):
         """Put requests back."""
         self.restore_called = True
         running.extend(self.waiting_for_chunk_running_requests)

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -533,6 +533,52 @@ def test_finish_requests_restores_status(build_adapter):
     assert req_id not in adapter.requests_origin_status
 
 
+def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
+    other = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
+    adapter.waiting_for_chunk_waiting_requests = deque([zombie, other])
+    adapter.waiting_for_chunk_running_requests = deque([other, zombie])
+    adapter.requests_with_ready_chunks.add("req-zombie")
+    adapter.finished_requests.add("req-zombie")
+    requests_map = {
+        "req-zombie": zombie,
+        "req-live": other,
+    }
+
+    adapter.finish_requests(
+        ["req-zombie"],
+        RequestStatus.FINISHED_ABORTED,
+        requests_map,
+    )
+
+    assert [req.request_id for req in adapter.waiting_for_chunk_waiting_requests] == ["req-live"]
+    assert [req.request_id for req in adapter.waiting_for_chunk_running_requests] == ["req-live"]
+    assert "req-zombie" not in adapter.requests_with_ready_chunks
+    assert "req-zombie" not in adapter.finished_requests
+
+
+def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
+    live = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
+    waiting_queue = DummyWaitingQueue()
+    running_queue = []
+    adapter.waiting_for_chunk_waiting_requests = deque([zombie, live])
+    adapter.waiting_for_chunk_running_requests = deque([zombie, live])
+
+    adapter.restore_queues(
+        waiting_queue,
+        running_queue,
+        scheduler_requests={"req-live": live},
+    )
+
+    assert [req.request_id for req in waiting_queue] == ["req-live"]
+    assert [req.request_id for req in running_queue] == ["req-live"]
+    assert not adapter.waiting_for_chunk_waiting_requests
+    assert not adapter.waiting_for_chunk_running_requests
+
+
 # ---------------------------------------------------------------
 # Scheduler trigger tests
 # ---------------------------------------------------------------

tests/engine/test_async_omni_engine_stage_init.py

@@ -1,3 +1,4 @@
+import concurrent.futures
 import importlib
 import os
 import threading
@@ -379,6 +380,44 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
     assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas)
 
 
+def test_initialize_stage_replicas_reuses_long_lived_executor(monkeypatch):
+    engine = object.__new__(AsyncOmniEngine)
+    engine._stage_init_executor = None
+
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg, num_replicas=1)]
+
+    created_executors = []
+
+    class FakeExecutor:
+        def __init__(self, *args, **kwargs):
+            self.shutdown_calls = []
+            created_executors.append(self)
+
+        def submit(self, fn, *args, **kwargs):
+            future = concurrent.futures.Future()
+            future.set_result(fn(*args, **kwargs))
+            return future
+
+        def shutdown(self, wait=False, cancel_futures=False):
+            self.shutdown_calls.append((wait, cancel_futures))
+
+    monkeypatch.setattr(concurrent.futures, "ThreadPoolExecutor", FakeExecutor)
+    monkeypatch.setattr(
+        engine,
+        "_initialize_replica",
+        lambda plan, _stage_init_timeout, _stage_launch_lock: f"client-{plan.replica_id}",
+    )
+
+    first = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
+    second = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
+
+    assert first == {0: ["client-0"]}
+    assert second == {0: ["client-0"]}
+    assert len(created_executors) == 1
+    assert engine._stage_init_executor is created_executors[0]
+
+
 def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch):
     engine = object.__new__(AsyncOmniEngine)
 
@@ -410,6 +449,42 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
     }
 
 
+def test_shutdown_closes_long_lived_stage_init_executor():
+    engine = object.__new__(AsyncOmniEngine)
+    engine._shutdown_called = False
+    engine._weak_finalizer = None
+    engine._omni_master_server = None
+    engine._coordinator_runtime = None
+    engine.orchestrator_thread = types.SimpleNamespace(join=lambda timeout: None, is_alive=lambda: False)
+
+    class FakeQueue:
+        def __init__(self):
+            self.closed = False
+            self.sync_q = types.SimpleNamespace(put_nowait=lambda _msg: None)
+
+        def close(self):
+            self.closed = True
+
+    class FakeExecutor:
+        def __init__(self):
+            self.shutdown_calls = []
+
+        def shutdown(self, wait=False, cancel_futures=False):
+            self.shutdown_calls.append((wait, cancel_futures))
+
+    engine.request_queue = FakeQueue()
+    engine.output_queue = FakeQueue()
+    engine.rpc_output_queue = FakeQueue()
+    engine._stage_init_executor = FakeExecutor()
+
+    engine.shutdown()
+
+    assert engine._stage_init_executor is None
+    assert engine.request_queue.closed is True
+    assert engine.output_queue.closed is True
+    assert engine.rpc_output_queue.closed is True
+
+
 def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
 

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -232,7 +232,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         finally:
             if self.chunk_transfer_adapter:
                 # Add request waiting for chunk to the waiting and running queue
-                self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+                self.chunk_transfer_adapter.restore_queues(
+                    self.waiting,
+                    self.running,
+                    scheduler_requests=self.requests,
+                )
             if self.input_coordinator:
                 self.input_coordinator.restore_queues(self.waiting, self.running)
         try:

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -213,7 +213,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             if self.chunk_transfer_adapter:
                 # Don't fall back: base scheduler doesn't handle async_chunk
                 # requests with empty prompt_token_ids.
-                self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+                self.chunk_transfer_adapter.restore_queues(
+                    self.waiting,
+                    self.running,
+                    scheduler_requests=self.requests,
+                )
             else:
                 res = super().schedule()
                 if self.input_coordinator:
@@ -345,7 +349,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # otherwise they are permanently orphaned in the adapter's
             # internal deques and never scheduled again.
             if self.chunk_transfer_adapter:
-                self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+                self.chunk_transfer_adapter.restore_queues(
+                    self.waiting,
+                    self.running,
+                    scheduler_requests=self.requests,
+                )
             if self.input_coordinator:
                 self.input_coordinator.restore_queues(self.waiting, self.running)
 

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -32,6 +32,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
     OmniMsgpackEncoder,
 )
 from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
+from vllm_omni.engine.stage_init_utils import set_death_signal
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -41,6 +42,14 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+_SIGNAL_EXIT_BASE = 128
+
+
+def _signal_exit_code(signum: int) -> int:
+    """Return the conventional process exit code for signal-driven exits."""
+    return _SIGNAL_EXIT_BASE + signum
+
+
 class StageDiffusionProc:
     """Subprocess entry point for diffusion inference.
 
@@ -624,11 +633,13 @@ class StageDiffusionProc:
         """
         shutdown_requested = False
 
+        set_death_signal(signal.SIGTERM)
+
         def signal_handler(signum: int, frame: Any) -> None:
             nonlocal shutdown_requested
             if not shutdown_requested:
                 shutdown_requested = True
-                raise SystemExit(128 + signum)
+                raise SystemExit(_signal_exit_code(signum))
 
         signal.signal(signal.SIGTERM, signal_handler)
         signal.signal(signal.SIGINT, signal_handler)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -11,6 +11,7 @@ to DiffusionModelRunner.
 import gc
 import multiprocessing as mp
 import os
+import signal
 from collections.abc import Iterable, Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
@@ -46,6 +47,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
+from vllm_omni.engine.stage_init_utils import set_death_signal
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.profiler import OmniTorchProfilerWrapper, create_omni_profiler
@@ -831,6 +833,8 @@ class WorkerProc:
         """Worker initialization and execution loops."""
         from vllm_omni.plugins import load_omni_general_plugins
 
+        set_death_signal(signal.SIGTERM)
+
         # Set process title for visibility in nvidia-smi / htop (optional, non-fatal)
         try:
             import setproctitle

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -389,17 +389,28 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             request.status = RequestStatus.PREEMPTED
             waiting_queue.prepend_requests([request])
 
-    def restore_queues(self, waiting_queue: Any, running_queue: list[Request]) -> None:
+    def restore_queues(
+        self,
+        waiting_queue: Any,
+        running_queue: list[Request],
+        scheduler_requests: dict[str, Request] | None = None,
+    ) -> None:
         """
         Restore requests waiting for chunk to the waiting and running queues.
         """
         # Add request waiting for chunk to the waiting and running queue
         for request in self.waiting_for_chunk_waiting_requests:
-            waiting_queue.add_request(request)
+            if scheduler_requests is None or request.request_id in scheduler_requests:
+                waiting_queue.add_request(request)
         self.waiting_for_chunk_waiting_requests = deque()
 
         if self.waiting_for_chunk_running_requests:
-            running_queue.extend(self.waiting_for_chunk_running_requests)
+            live_running_requests = [
+                request
+                for request in self.waiting_for_chunk_running_requests
+                if scheduler_requests is None or request.request_id in scheduler_requests
+            ]
+            running_queue.extend(live_running_requests)
         self.waiting_for_chunk_running_requests = deque()
 
     def postprocess_scheduler_output(
@@ -493,3 +504,20 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 continue
             if req_id in self.requests_origin_status:
                 request.status = self.requests_origin_status.pop(req_id)
+
+        request_ids = set(request_ids)
+
+        self.waiting_for_chunk_waiting_requests = deque(
+            request for request in self.waiting_for_chunk_waiting_requests if request.request_id not in request_ids
+        )
+        self.waiting_for_chunk_running_requests = deque(
+            request for request in self.waiting_for_chunk_running_requests if request.request_id not in request_ids
+        )
+
+        for req_id in request_ids:
+            self.requests_with_ready_chunks.discard(req_id)
+            self.finished_requests.discard(req_id)
+            self._finished_load_reqs.discard(req_id)
+            self._cancelled_load_reqs.add(req_id)
+
+        return []

vllm_omni/engine/async_omni_engine.py

@@ -354,6 +354,7 @@ class AsyncOmniEngine:
         self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue()
         self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue()
         self._shutdown_called = False
+        self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None
         self._weak_finalizer: weakref.finalize | None = None
         self._rpc_lock = threading.Lock()
 
@@ -1167,35 +1168,37 @@ class AsyncOmniEngine:
                 primary_exc = exc
                 break
 
-        # --- 2) LLM replicas: parallel init via a scoped ThreadPoolExecutor. ---
+        # --- 2) LLM replicas: parallel init via a long-lived ThreadPoolExecutor. ---
         if primary_exc is None and llm_replicas:
             future_to_replica: dict[concurrent.futures.Future[StagePoolClient], tuple[int, int]] = {}
-            with concurrent.futures.ThreadPoolExecutor(
-                max_workers=max(1, len(llm_replicas)),
-                thread_name_prefix="stage-init",
-            ) as init_executor:
-                for stage_idx, replica in llm_replicas:
-                    future = init_executor.submit(
-                        self._initialize_replica,
-                        replica,
-                        stage_init_timeout,
-                        stage_launch_lock,
-                    )
-                    future_to_replica[future] = (stage_idx, replica.replica_id)
+            if getattr(self, "_stage_init_executor", None) is None:
+                self._stage_init_executor = concurrent.futures.ThreadPoolExecutor(
+                    max_workers=max(1, len(llm_replicas)),
+                    thread_name_prefix="stage-init",
+                )
+            init_executor = self._stage_init_executor
+            for stage_idx, replica in llm_replicas:
+                future = init_executor.submit(
+                    self._initialize_replica,
+                    replica,
+                    stage_init_timeout,
+                    stage_launch_lock,
+                )
+                future_to_replica[future] = (stage_idx, replica.replica_id)
 
-                for future in concurrent.futures.as_completed(future_to_replica):
-                    stage_idx, replica_id = future_to_replica[future]
-                    try:
-                        initialized_clients_by_stage[stage_idx][replica_id] = future.result()
-                    except concurrent.futures.CancelledError:
-                        continue
-                    except Exception as exc:
-                        if primary_exc is None:
-                            primary_exc = exc
-                            for other_future in future_to_replica:
-                                if other_future is future:
-                                    continue
-                                other_future.cancel()
+            for future in concurrent.futures.as_completed(future_to_replica):
+                stage_idx, replica_id = future_to_replica[future]
+                try:
+                    initialized_clients_by_stage[stage_idx][replica_id] = future.result()
+                except concurrent.futures.CancelledError:
+                    continue
+                except Exception as exc:
+                    if primary_exc is None:
+                        primary_exc = exc
+                        for other_future in future_to_replica:
+                            if other_future is future:
+                                continue
+                            other_future.cancel()
 
         if primary_exc is not None:
             setattr(primary_exc, "_initialized_clients_by_stage", initialized_clients_by_stage)
@@ -2424,6 +2427,11 @@ class AsyncOmniEngine:
             if self.orchestrator_thread.is_alive():
                 logger.warning("[AsyncOmniEngine] Orchestrator thread did not exit in time")
 
+        stage_init_executor = getattr(self, "_stage_init_executor", None)
+        if stage_init_executor is not None:
+            stage_init_executor.shutdown(wait=False, cancel_futures=True)
+            self._stage_init_executor = None
+
         for q in (self.request_queue, self.output_queue, self.rpc_output_queue):
             try:
                 q.close()

vllm_omni/engine/stage_engine_core_client.py

@@ -6,11 +6,13 @@ Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture.
 
 from __future__ import annotations
 
+import contextlib
 import inspect
 import multiprocessing.connection
 import os
 import socket
 import threading
+import time
 import weakref
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
@@ -158,6 +160,7 @@ class StageEngineCoreClientBase(StageClientBase):
 
         self.engine_outputs: Any = None
         self._proc = proc
+        self._shutting_down = False
         self.client_addresses = dict(client_addresses or {})
         self._omni_kv_config = getattr(getattr(vllm_config, "model_config", None), "omni_kv_config", None)
         self._kv_sender_host = self._resolve_contact_host()
@@ -233,15 +236,37 @@ class StageEngineCoreClientBase(StageClientBase):
                 multiprocessing.connection.wait([proc.sentinel])
             except Exception:
                 return
+
+            # Give multiprocessing a brief chance to publish a stable exitcode
+            # before emitting the monitor log. Without this, the sentinel can
+            # fire first and the log may misleadingly print ``None``.
+            with contextlib.suppress(Exception):
+                proc.join(timeout=0)
+            exitcode = proc.exitcode
+            if exitcode is None:
+                deadline = time.monotonic() + 0.2
+                while exitcode is None and time.monotonic() < deadline:
+                    time.sleep(0.01)
+                    with contextlib.suppress(Exception):
+                        proc.join(timeout=0)
+                    exitcode = proc.exitcode
+
             resources = resources_ref()
-            if resources is None or resources.engine_dead:
+            if resources is None or resources.engine_dead or self._shutting_down:
+                return
+            if exitcode == 0:
+                logger.info(
+                    "[StageEngineCoreClient] stage-%s [rep-%s] subprocess exited cleanly.",
+                    stage_id,
+                    replica_id,
+                )
                 return
             resources.engine_dead = True
             logger.error(
                 "[StageEngineCoreClient] stage-%s [rep-%s] subprocess died unexpectedly (exit code %s).",
                 stage_id,
                 replica_id,
-                proc.exitcode,
+                exitcode,
             )
 
         t = threading.Thread(
@@ -462,6 +487,7 @@ class StageEngineCoreClientBase(StageClientBase):
 
     def shutdown(self, timeout: float | None = None) -> None:
         """Shutdown managed resources and any externally spawned subprocess."""
+        self._shutting_down = True
         child_procs: list[psutil.Process] = []
         if self._proc is not None and self._proc.pid is not None:
             try:

vllm_omni/engine/stage_engine_core_proc.py

@@ -36,6 +36,7 @@ from vllm.v1.engine.utils import (
 from vllm.v1.utils import shutdown
 
 from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
+from vllm_omni.engine.stage_init_utils import set_death_signal
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -44,6 +45,14 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+_SIGNAL_EXIT_BASE = 128
+
+
+def _signal_exit_code(signum: int) -> int:
+    """Return the conventional process exit code for signal-driven exits."""
+    return _SIGNAL_EXIT_BASE + signum
+
+
 class StageEngineCoreProc(EngineCoreProc):
     """Stage-specific engine core process for vLLM-Omni.
 
@@ -89,6 +98,7 @@ class StageEngineCoreProc(EngineCoreProc):
             # like upstream vLLM.
 
             stage_label = f"stage{omni_stage_id}" if omni_stage_id is not None else "noid"
+            set_death_signal(signal.SIGTERM)
             set_process_title(f"StageEngineCoreProc_{stage_label}_replica{omni_replica_id}_DP{dp_rank}")
             decorate_logs()
             os.environ["VLLM_OMNI_REPLICA_ID"] = str(max(int(omni_replica_id), 0))
@@ -142,6 +152,7 @@ class StageEngineCoreProc(EngineCoreProc):
             def signal_handler(signum: int, frame: Any) -> None:
                 engine_core.shutdown_state = EngineShutdownState.REQUESTED
                 signal_callback.trigger()
+                raise SystemExit(_signal_exit_code(signum))
 
             signal.signal(signal.SIGTERM, signal_handler)
             signal.signal(signal.SIGINT, signal_handler)

vllm_omni/engine/stage_init_utils.py

@@ -182,6 +182,19 @@ def terminate_alive_proc(proc, timeout=5):
             proc.kill()
 
 
+def set_death_signal(sig: int) -> None:
+    """Best-effort parent-death signal for Linux subprocesses."""
+    try:
+        import ctypes
+        import platform
+
+        if platform.system() != "Linux":
+            return
+        ctypes.CDLL("libc.so.6").prctl(1, sig)
+    except Exception:
+        pass
+
+
 def patch_generation_config_if_needed(model_config: Any) -> None:
     """Guard InputProcessor init for models whose config lacks model_type."""
     try: