diff --git a/tests/core/sched/test_generation_scheduler_restore.py b/tests/core/sched/test_generation_scheduler_restore.py index 5cc1cab70..096471e5b 100644 --- a/tests/core/sched/test_generation_scheduler_restore.py +++ b/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) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 998a8ebb9..24ba819b1 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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 # --------------------------------------------------------------- diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 71952e6cf..9d474e81c 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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 diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 61ee24547..09ee55ba9 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 1218d416b..2dfd85d3b 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/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) diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 871a29729..19e890abb 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/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) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index cd5fe246d..782e79855 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index bda2abeea..5686c7696 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/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 [] diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index fc682eb92..4986eae63 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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() diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 84bf72863..b5aecf124 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/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: diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 0b3dadcbd..0bb700369 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/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) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index bd727b930..e650373a9 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/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: