[feature] Enable multi-node inference through Ray backend (#7938)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Co-authored-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

docs/configuration/stage_configs.md

@@ -78,7 +78,44 @@ explicit CLI or stage override, or the downstream vLLM engine default.
 Note: for the diffusion path, an omitted `distributed_executor_backend` selects
 `uni` on a single GPU (in-process worker, no MessageQueue / `/dev/shm` output
 segments) and `mp` when `num_gpus > 1`. Set `mp` explicitly to keep a worker
-subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet.
+subprocess on one GPU. Set `ray` for Ray-managed single- or multi-node
+diffusion workers. `external_launcher` is not supported yet.
+
+The diffusion `ray` backend runs one worker actor per GPU and supports workers
+spanning Ray nodes. Enable it in the deploy YAML:
+
+```yaml
+distributed_executor_backend: ray
+```
+
+The executor reuses Ray when already initialized; otherwise it calls
+`ray.init()`, which handles cluster discovery or local startup and respects
+`RAY_ADDRESS`. There is no diffusion-specific address setting. Ray must be
+installed on every node.
+
+Ray diffusion workers receive the driver's `VLLM_*`, `OMNI_*`, `DIFFUSION_*`,
+`NCCL_*`, `TORCH_NCCL_*`, `UCX_*`, and Hugging Face environment settings,
+plus `PYTHONPATH`, `CUDA_LAUNCH_BLOCKING`, and `OMP_NUM_THREADS`.
+Explicit stage `runtime.env` entries are also forwarded and take precedence.
+Device visibility, rank/rendezvous settings, worker host addresses, and `RAY_*`
+variables are excluded so each worker retains its own assignment and identity.
+
+For multiple nodes, start the cluster before launching Omni. For example,
+with head IP `10.0.0.1`:
+
+```bash
+# On the head node:
+ray start --head --node-ip-address=10.0.0.1 --port=6379
+
+# On each additional worker node:
+ray start --address=10.0.0.1:6379
+
+# On the host launching Omni (MODEL and my_deploy.yaml are placeholders):
+RAY_ADDRESS=10.0.0.1:6379 vllm serve MODEL --omni --deploy-config my_deploy.yaml
+```
+
+Omni creates actors on the resources of that cluster; it does not start Ray
+on additional machines.
 
 ### Stage-level runner selection
 

docs/design/module/diffusion/diffusion_runtime.md

@@ -36,6 +36,7 @@ validation_paths:
   - tests/diffusion/test_diffusion_worker.py
   - tests/diffusion/test_multiproc_engine_concurrency.py
   - tests/diffusion/test_uniproc_executor.py
+  - tests/diffusion/test_ray_executor.py
   - tests/diffusion/test_result_pump.py
   - tests/diffusion/test_async_output_worker.py
   - tests/diffusion/test_diffusion_ipc.py
@@ -96,8 +97,10 @@ The executor then chooses how the worker runs:
   RPC.
 - **`mp`** (default when `num_gpus > 1`, or when set explicitly):
   `MultiprocDiffusionExecutor` starts one `WorkerProc` per device.
+- **`ray`** (explicit): `RayDiffusionExecutor` starts one GPU actor per worker
+  across one or more Ray nodes.
 
-Either way, each worker owns its device, distributed state,
+In all backends, each worker owns its device, distributed state,
 `DiffusionWorker`, `DiffusionModelRunner`, and pipeline instance.
 
 ```mermaid
@@ -108,7 +111,7 @@ flowchart TB
         stageProc["StageDiffusionProc<br/>stage transport and lifecycle"]
         engine["DiffusionEngine<br/>admission and control loop"]
         scheduler["RequestScheduler or StepScheduler<br/>request state and policy"]
-        executor["DiffusionExecutor<br/>uni or mp backend"]
+        executor["DiffusionExecutor<br/>uni, mp, or ray backend"]
         streams["Per-request output streams"]
 
         stageProc --> engine
@@ -132,6 +135,8 @@ flowchart TB
     client -.->|"inline: direct call"| engine
     executor -->|"mp: broadcast / result queues"| workerProc
     executor -.->|"uni: in-process call"| worker
+    executor -->|"ray: actor RPC"| rayActor["RayDiffusionWorkerWrapper"]
+    rayActor --> worker
     streams -->|"process-backed"| stageProc
     stageProc -->|"ZMQ results"| client
     streams -.->|"inline"| client
@@ -154,6 +159,7 @@ flowchart TB
 | `DiffusionExecutor` | Execution backend contract, worker RPC, health, shutdown | Admission and request-state transitions |
 | `UniProcDiffusionExecutor` | Single in-process worker for `num_gpus == 1`; no IPC or async output pump | Multi-GPU execution |
 | `MultiprocDiffusionExecutor` | Worker processes, shared-memory message queues, result dispatch, worker monitoring | In-process single-GPU path |
+| `RayDiffusionExecutor` | GPU actors, placement groups, remote RPC, actor monitoring | Local worker subprocesses and shared-memory IPC |
 | `WorkerProc` | One `mp` worker process's IPC loop and reply rules | Scheduling policy; unused by `uni` |
 | `DiffusionWorker` | Device/distributed setup, LoRA activation, profiling, sleep/wake, runner delegation | Request admission |
 | `DiffusionModelRunner` | Pipeline loading, request-local model state, cache/compile setup, request or step execution | Queueing and cross-stage routing |
@@ -219,12 +225,12 @@ uses fused `execute_model_batch()` and requires the pipeline to declare
 request-batch support.
 
 **Distributed layerwise offload with AllGather (DLO DP concurrency)** is a
-separate multiproc dispatch path. It activates when
+separate dispatch path supported by `mp` and `ray`. It activates when
 `data_parallel_size > 1`, `enable_distributed_layerwise_offload` is set, and
 `dlo_use_allgather` is true. The engine then sets `dp_concurrent = True` and
 raises `scheduler.max_num_running_reqs` to `dp_size`, overriding the
 `max_num_seqs` value from `initialize()`. The scheduler still emits a
-multi-request `DiffusionSchedulerOutput`; the multiproc executor routes that
+multi-request `DiffusionSchedulerOutput`; these executors route that
 wave through `execute_request()` instead of fused `execute_model_batch()`.
 
 Optional admission coalescing can fill a `dp_size` wave: when
@@ -233,7 +239,7 @@ the first schedule of a wave (under `dp_concurrent`, the stable window is
 `min(0.3s, wait/2)`). With the default `request_batch_max_wait_ms == 0`,
 there is no wait.
 
-Before dispatch, `MultiprocDiffusionExecutor.execute_request()` rejects waves
+Before dispatch, both backends reject waves
 whose requests differ in sampling-parameter compatibility or `extra_args`
 (AllGather requires every DP rank to follow the same forward schedule). Each
 worker rank then picks one envelope from the wave:
@@ -297,8 +303,9 @@ head-of-line blocking.
 | --- | --- | --- |
 | `uni` | Default for `num_gpus == 1`, or set explicitly | One in-process worker |
 | `mp` | Default for `num_gpus > 1`, or set explicitly | One worker process per device |
+| `ray` | Set explicitly | One GPU actor per worker across Ray nodes |
 
-Ray and external-launcher diffusion backends are not implemented. A custom
+The external-launcher diffusion backend is not implemented. A custom
 `DiffusionExecutor` subclass or import path is also accepted.
 
 ### Uniproc backend
@@ -337,11 +344,27 @@ result path and does not start those pumps.
 See [Async diffusion output](../../feature/async_diffusion_output.md) for the
 multiproc request-mode timeline.
 
+### Ray backend
+
+`RayDiffusionExecutor` reuses initialized Ray or calls `ray.init()` (respecting
+`RAY_ADDRESS`). It reuses the current placement group or reserves one GPU
+bundle per worker with `PACK` placement. Ray owns GPU placement; PyTorch
+owns distributed collectives through a shared TCP rendezvous. Global ranks
+are grouped by host, while each actor uses local device 0.
+
+Request batches, step execution, and DLO DP waves use actor RPCs with
+rank-aware replies. Returned tensors move to CPU before Ray serialization;
+this backend does not use the multiproc async output pumps. Driver inference
+settings and explicit stage `runtime.env` are forwarded, excluding worker
+identity and device assignment. See [Stage configuration](../../../configuration/stage_configs.md)
+for cluster setup and environment details.
+
 ## Worker and runner boundary
 
 On the `mp` path, `WorkerProc` receives messages and calls methods through
 `WorkerWrapperBase`. On the `uni` path, the executor owns that wrapper
-directly. In both cases, `DiffusionWorker` handles the parts tied to a device:
+directly. Ray actors also construct `WorkerWrapperBase`. In all cases,
+`DiffusionWorker` handles the parts tied to a device:
 distributed initialization, model-runner construction, LoRA activation,
 profiling, and memory sleep/wake. It delegates actual model work to
 `DiffusionModelRunner`.
@@ -399,6 +422,10 @@ streams with an error, closes scheduler state, and shuts down the executor.
   async-output futures as failed, then clears them. Completed async outputs
   cached for later `wait_output_ready()` calls are dropped with the executor
   object.
+- **`ray`**: the executor gives actors up to five seconds to shut down, then
+  kills them and removes placement groups it owns. Actor death is monitored
+  even while idle; actor errors and RPC timeouts fail the group and notify the
+  stage. The stage client allows subprocess cleanup before signal fallback.
 - **`uni`**: the executor shuts down the in-process worker, drops the final
   model reference, and empties the accelerator cache so a later engine can reuse
   the device.
@@ -407,7 +434,7 @@ streams with an error, closes scheduler state, and shuts down the executor.
 
     Do not continue using a worker group after a fatal collective timeout or
     unexpected worker exit. Distributed state may be incomplete. The
-    multiprocess executor deliberately fails closed so the stage can restart.
+    `mp` and `ray` executors deliberately fail closed so the stage can restart.
     The uniproc executor fails closed when the accelerator context itself is
     poisoned.
 
@@ -417,10 +444,11 @@ streams with an error, closes scheduler state, and shuts down the executor.
   `diffusion_model_runner_cls` configuration still wins.
 - Tests and custom engine integrations may inject a `BaseScheduler` subclass.
   `SchedulerInterface` remains only as a deprecated compatibility name.
-- `distributed_executor_backend` may be `"uni"`, `"mp"`, a custom `DiffusionExecutor`
+- `distributed_executor_backend` may be `"uni"`, `"mp"`, `"ray"`, a custom `DiffusionExecutor`
   subclass, or an import path. A backend must preserve scheduler output,
   request identity, health, and cleanup contracts. Use `"mp"` when you need
-  process isolation or multi-GPU; `"uni"` is the single-GPU default.
+  local process isolation or multi-GPU, or `"ray"` for workers across nodes;
+  `"uni"` is the single-GPU default.
 - Worker extensions go through `WorkerWrapperBase`. They add worker methods;
   they do not become a second scheduler or request lifecycle.
 
@@ -466,7 +494,7 @@ Test the smallest affected slice, then cover its neighboring boundary:
 | --- | --- |
 | Admission or state transitions | `test_diffusion_scheduler.py`, including duplicate IDs, compatibility, abort, and missing output |
 | Engine loop or output delivery | `test_diffusion_engine.py`, `test_diffusion_engine_cleanup.py`, `test_diffusion_engine_rpc_routing.py` |
-| Executor or IPC | `test_multiproc_engine_concurrency.py`, `test_uniproc_executor.py`, `test_result_pump.py`, `test_diffusion_ipc.py`, `test_async_output_worker.py` |
+| Executor or IPC | `test_multiproc_engine_concurrency.py`, `test_uniproc_executor.py`, `test_ray_executor.py`, `test_result_pump.py`, `test_diffusion_ipc.py`, `test_async_output_worker.py` |
 | Worker or runner | `test_diffusion_worker.py`, `test_diffusion_model_runner.py` |
 | Stage boundary | `test_stage_diffusion_proc.py`, `test_inline_stage_diffusion_client.py` |
 | Warmup or streaming | `test_diffusion_engine_dummy_run.py`, `test_diffusion_streaming_output.py` |

tests/config/test_omni_config.py

@@ -1948,7 +1948,9 @@ def test_diffusion_config_field_classification_covers_current_fields():
 
     assert classified_fields == {f.name for f in fields(omni_config_module._DiffusionConfigProjection)}
     assert {f.name for f in fields(OmniDiffusionConfig)} <= (
-        classified_fields | omni_config_module._DIFFUSION_MOVED_SHARED_FIELDS
+        classified_fields
+        | omni_config_module._DIFFUSION_MOVED_SHARED_FIELDS
+        | omni_config_module._DIFFUSION_INTERNAL_FIELDS
     )
     assert {
         "enable_prompt_embed_cache",

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -31,6 +31,7 @@ def _make_request(request_id: str) -> OmniDiffusionRequest:
 
 def _make_engine() -> DiffusionEngine:
     engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(distributed_executor_backend=None)
     engine.scheduler = RequestScheduler()
     engine.scheduler.initialize(SimpleNamespace())
     engine.executor = SimpleNamespace(shutdown=Mock())

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -1290,6 +1290,44 @@ class TestStageDiffusionClientErrorPropagation:
         client = self._make_client()
         client.check_health()
 
+    def test_local_shutdown_uses_existing_signal_path(self):
+        client = self._make_client()
+        client._zmq_ctx = MagicMock()
+        client._proc_manager.distributed_executor_backend = "mp"
+        client._proc_manager.wait_for_shutdown = MagicMock()
+        client._proc_manager.shutdown = MagicMock()
+        client.shutdown()
+        client._proc_manager.wait_for_shutdown.assert_not_called()
+        client._proc_manager.shutdown.assert_called_once_with(timeout=10)
+
+    def test_shutdown_waits_for_clean_subprocess_exit(self):
+        client = self._make_client()
+        client._zmq_ctx = MagicMock()
+        client._proc_manager.distributed_executor_backend = "ray"
+        client._proc_manager.wait_for_shutdown = MagicMock(return_value=True)
+        client._proc_manager.shutdown = MagicMock()
+
+        client.shutdown()
+
+        client._request_socket.send.assert_called_once()
+        client._proc_manager.wait_for_shutdown.assert_called_once_with(timeout=10)
+        client._proc_manager.shutdown.assert_not_called()
+        client._request_socket.close.assert_called_once_with(linger=0)
+        client._response_socket.close.assert_called_once_with(linger=0)
+        client._zmq_ctx.term.assert_called_once_with()
+
+    def test_shutdown_uses_signals_only_after_clean_exit_timeout(self):
+        client = self._make_client()
+        client._zmq_ctx = MagicMock()
+        client._proc_manager.distributed_executor_backend = "ray"
+        client._proc_manager.wait_for_shutdown = MagicMock(return_value=False)
+        client._proc_manager.shutdown = MagicMock()
+
+        client.shutdown()
+
+        client._proc_manager.wait_for_shutdown.assert_called_once_with(timeout=10)
+        client._proc_manager.shutdown.assert_called_once_with(timeout=10)
+
     def test_get_output_raises_engine_dead_when_dead(self):
         """When ``_engine_dead`` is True and the output queue is empty,
         ``get_diffusion_output_nowait`` must raise ``EngineDeadError``."""
@@ -1783,3 +1821,29 @@ class TestDrainResponsesDeathSentinel:
         assert output.request_id == "req-fail"
         assert output.error == "gpu fault"
         assert output.finished is True
+
+
+@pytest.mark.parametrize("backend", ["ray", "mp"])
+def test_engine_close_interrupts_only_ray_before_join(backend):
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(distributed_executor_backend=backend)
+    engine._cv = threading.Condition()
+    engine._closed = False
+    engine._shutdown_complete = False
+    engine.stop_event = threading.Event()
+    engine._out_streams = {}
+    interrupted = threading.Event()
+    engine.executor = MagicMock()
+    engine.executor.shutdown.side_effect = interrupted.set
+    engine.scheduler = MagicMock(_native_prefetch_enabled=False)
+    engine.worker_thread = MagicMock()
+    engine.worker_thread.is_alive.side_effect = [True, False]
+
+    def join(timeout):
+        assert engine._closed
+        assert interrupted.is_set() == (backend == "ray")
+
+    engine.worker_thread.join.side_effect = join
+    engine.close()
+    assert engine._shutdown_complete
+    engine.scheduler.close.assert_called_once()

tests/diffusion/test_ray_executor.py

@@ -0,0 +1,586 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU tests of Ray placement, rank assignment, RPC output, and cleanup."""
+
+import threading
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.executor import ray_executor as ray_module
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_cross_node_workers_get_global_rank_and_explicit_rendezvous(monkeypatch):
+    monkeypatch.setenv("NCCL_SOCKET_IFNAME", "eth-test")
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,4,5")
+    actors = []
+    for ip in ["10.0.0.2", "10.0.0.1", "10.0.0.2"]:
+        actor = Mock()
+        actor.get_node_ip.remote.return_value = ip
+        actor.get_open_port.remote.return_value = 23456
+        actors.append(actor)
+    factory = Mock()
+    factory.remote.side_effect = actors
+    fake_ray = Mock()
+    fake_ray.remote.return_value = lambda cls: factory
+    fake_ray.get.side_effect = lambda values, **kwargs: values
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    monkeypatch.setattr(ray_module, "get_ip", lambda: "10.0.0.1")
+    monkeypatch.setattr(ray_module, "PlacementGroupSchedulingStrategy", Mock())
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor.od_config = SimpleNamespace(num_gpus=3, ray_worker_env={"CUSTOM_PLUGIN_SETTING": "enabled"})
+    executor.workers = []
+    executor._init_workers(object())
+    assert fake_ray.remote.call_args.kwargs["max_concurrency"] == 1
+    assert fake_ray.remote.call_args.kwargs["concurrency_groups"] == {"health": 1, "monitor": 1}
+    worker_env = fake_ray.remote.call_args.kwargs["runtime_env"]["env_vars"]
+    assert worker_env["NCCL_SOCKET_IFNAME"] == "eth-test"
+    assert worker_env["CUSTOM_PLUGIN_SETTING"] == "enabled"
+    assert "CUDA_VISIBLE_DEVICES" not in worker_env
+    assert len(fake_ray.get.call_args_list) == 3
+    assert all(call.kwargs["timeout"] == ray_module._WORKER_INIT_TIMEOUT_S for call in fake_ray.get.call_args_list)
+    assert [worker.ip for worker in executor.workers] == ["10.0.0.1", "10.0.0.2", "10.0.0.2"]
+    for rank, metadata in enumerate(executor.workers):
+        metadata.worker.init_worker.remote.assert_called_once_with(executor.od_config, rank, "tcp://10.0.0.1:23456")
+
+
+def test_worker_env_forwards_stage_overrides_without_driver_identity(monkeypatch):
+    from unittest.mock import patch
+
+    driver_env = {
+        "HF_TOKEN": "test-token",
+        "VLLM_PLUGINS": "test-plugin",
+        "NCCL_DEBUG": "WARN",
+        "DIFFUSION_ATTENTION_BACKEND": "FLASH_ATTN",
+        "PYTHONPATH": "/test/modules",
+        "UNRELATED_SECRET": "do-not-copy",
+        "HOME": "/driver/home",
+    }
+    protected = {
+        "CUDA_VISIBLE_DEVICES": "7",
+        "VLLM_HOST_IP": "10.0.0.1",
+        "VLLM_HOST_PORT": "1234",
+        "VLLM_NIXL_SIDE_CHANNEL_HOST": "10.0.0.1",
+        "LOCAL_RANK": "7",
+        "RANK": "7",
+        "WORLD_SIZE": "8",
+        "MASTER_ADDR": "driver-host",
+        "MASTER_PORT": "1234",
+        "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
+    }
+    config = SimpleNamespace(ray_worker_env={**protected, "NCCL_DEBUG": "INFO", "CUSTOM_PLUGIN_SETTING": "stage"})
+    with patch.dict(ray_module.os.environ, {**driver_env, **protected}, clear=True):
+        env = ray_module._worker_env(config)
+        assert env == {key: value for key, value in driver_env.items() if key not in {"UNRELATED_SECRET", "HOME"}} | {
+            "NCCL_DEBUG": "INFO",
+            "CUSTOM_PLUGIN_SETTING": "stage",
+        }
+        assert ray_module.os.environ["NCCL_DEBUG"] == "WARN"
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_collective_executes_all_ranks_but_only_output_rank_returns_tensors(rank):
+    actor = ray_module.RayDiffusionWorkerWrapper(rank)
+    value = torch.ones(2, requires_grad=True)
+    actor.worker = Mock()
+    actor.worker.execute_method.return_value = DiffusionOutput(output=value)
+    result = actor.execute_rpc("execute_model", output_rank=0, exec_all_ranks=True)
+    actor.worker.execute_method.assert_called_once()
+    assert result["replied"] == (rank == 0)
+    if rank == 0:
+        assert not result["result"].output.requires_grad
+        torch.testing.assert_close(result["result"].output, value)
+    else:
+        assert result["result"] is None
+
+
+def test_finalizer_kills_workers_without_waiting_and_releases_owned_group(monkeypatch):
+    events = []
+    fake_ray = Mock()
+    fake_ray.get.side_effect = lambda futures, **kw: events.append("wait")
+    fake_ray.kill.side_effect = lambda worker, **kw: events.append("kill")
+    fake_ray.util.remove_placement_group.side_effect = lambda group: events.append("remove")
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actor = Mock()
+    actor.shutdown.remote.side_effect = lambda: events.append("shutdown")
+    resources = ray_module._RayExecutorResources([ray_module._RayWorkerMetadata(actor)], object(), True)
+    resources()
+    assert events == ["kill", "remove"]
+    fake_ray.kill.assert_called_once_with(actor, no_restart=True)
+    actor.shutdown.remote.assert_not_called()
+    fake_ray.get.assert_not_called()
+    resources()
+    assert events == ["kill", "remove"]
+    assert resources.workers == []
+
+
+def test_wrapper_shutdown_releases_worker_once():
+    actor = ray_module.RayDiffusionWorkerWrapper(0)
+    worker = Mock()
+    actor.worker = worker
+    actor.shutdown()
+    actor.shutdown()
+    worker.shutdown.assert_called_once_with()
+    assert not actor.check_alive()
+
+
+def _make_executor(fake_ray, actors, group=None, owns_group=True):
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor._closed = False
+    executor._is_failed = False
+    executor._state_lock = threading.Lock()
+    executor._failure_callbacks = []
+    executor.workers = [ray_module._RayWorkerMetadata(actor, rank=rank) for rank, actor in enumerate(actors)]
+    executor._resources = ray_module._RayExecutorResources(executor.workers, group, owns_group)
+    executor._finalizer = executor._resources
+    return executor
+
+
+@pytest.mark.parametrize("stuck_rank", [None, 1])
+def test_shutdown_drains_workers_before_killing_them(monkeypatch, stuck_rank):
+    events = []
+    actors = [Mock(), Mock()]
+    refs = []
+    for rank, actor in enumerate(actors):
+        ref = f"shutdown-{rank}"
+        refs.append(ref)
+        actor.shutdown.remote.side_effect = lambda ref=ref: events.append(ref) or ref
+    fake_ray = Mock()
+
+    def wait(pending, num_returns, timeout):
+        events.append("wait")
+        assert num_returns == len(pending)
+        assert timeout == ray_module._WORKER_SHUTDOWN_TIMEOUT_S
+        return [ref for ref in pending if ref != f"shutdown-{stuck_rank}"], [
+            ref for ref in pending if ref == f"shutdown-{stuck_rank}"
+        ]
+
+    fake_ray.wait.side_effect = wait
+    fake_ray.kill.side_effect = lambda worker, **kw: events.append("kill")
+    fake_ray.util.remove_placement_group.side_effect = lambda group: events.append("remove")
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    executor = _make_executor(fake_ray, actors, group=object())
+    callback = Mock()
+    executor._failure_callbacks = [callback]
+
+    executor.shutdown()
+
+    assert events == [*refs, "wait", "kill", "kill", "remove"]
+    assert [call.args[0] for call in fake_ray.get.call_args_list] == [
+        ref for ref in refs if ref != f"shutdown-{stuck_rank}"
+    ]
+    assert executor._closed
+    assert not executor._is_failed
+    callback.assert_not_called()
+    executor.shutdown()
+    assert fake_ray.kill.call_count == len(actors)
+
+
+def test_shutdown_kills_workers_when_drain_fails(monkeypatch):
+    fake_ray = Mock()
+    fake_ray.wait.side_effect = RuntimeError("Ray is shutting down")
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actors = [Mock(), Mock()]
+    executor = _make_executor(fake_ray, actors)
+    executor.shutdown()
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors
+
+
+def test_shutdown_preserves_borrowed_placement_group(monkeypatch):
+    fake_ray = Mock()
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    ray_module._RayExecutorResources([], object(), False)()
+    fake_ray.util.remove_placement_group.assert_not_called()
+
+
+@pytest.mark.parametrize("failure_kind", ["GetTimeoutError", "RayActorError", "RayTaskError"])
+@pytest.mark.parametrize("output_rank", [None, 0])
+def test_rpc_failure_kills_peers_and_rejects_subsequent_work(monkeypatch, failure_kind, output_rank):
+    fake_ray = Mock()
+    fake_ray.exceptions = SimpleNamespace(
+        **{name: type(name, (Exception,), {}) for name in ("GetTimeoutError", "RayActorError", "RayTaskError")}
+    )
+    failure = getattr(fake_ray.exceptions, failure_kind)("rank failed while peers are still running")
+    fake_ray.get.side_effect = failure
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actors = [Mock(), Mock()]
+    group = object()
+    executor = _make_executor(fake_ray, actors, group=group)
+    callback = Mock()
+    executor._failure_callbacks = [callback]
+
+    expected_error = TimeoutError if failure_kind == "GetTimeoutError" else ray_module.EngineDeadError
+    with pytest.raises(expected_error) as raised:
+        executor.collective_rpc("execute_model", unique_reply_rank=output_rank, exec_all_ranks=True)
+
+    assert raised.value.__cause__ is failure
+    assert executor.is_dead
+    assert executor._is_failed
+    assert executor._closed
+    callback.assert_called_once_with()
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+    # Peers may be stuck in a collective, so failure skips the graceful drain.
+    fake_ray.wait.assert_not_called()
+    for actor in actors:
+        actor.shutdown.remote.assert_not_called()
+
+    with pytest.raises(RuntimeError, match="closed"):
+        executor.collective_rpc("execute_model", unique_reply_rank=output_rank, exec_all_ranks=True)
+    for actor in actors:
+        actor.execute_rpc.remote.assert_called_once()
+    fake_ray.get.assert_called_once()
+    executor.shutdown()
+    assert fake_ray.kill.call_count == len(actors)
+    callback.assert_called_once_with()
+
+
+def test_actor_death_during_shutdown_does_not_report_failure():
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor._closed = True
+    executor._is_failed = False
+    executor._state_lock = threading.Lock()
+    executor._finalizer = Mock()
+    callback = Mock()
+    executor._failure_callbacks = [callback]
+    executor._fail(RuntimeError("actor terminated"))
+    assert not executor._is_failed
+    callback.assert_not_called()
+    executor._finalizer.assert_not_called()
+
+
+class _InlineThread:
+    """Run the monitor synchronously so tests observe its effects."""
+
+    def __init__(self, target, **kwargs):
+        self._target = target
+
+    def start(self):
+        self._target()
+
+
+@pytest.mark.parametrize("closed_first", [False, True])
+def test_worker_monitor_fails_executor_when_idle_actor_dies(monkeypatch, closed_first):
+    fake_ray = Mock()
+    death = RuntimeError("actor died")
+    actors = [Mock(), Mock()]
+    sentinels = [actor.wait_for_exit.options.return_value.remote.return_value for actor in actors]
+
+    def wait(refs, num_returns):
+        assert refs == sentinels
+        assert num_returns == 1
+        if closed_first:
+            executor._closed = True
+        return [sentinels[1]], [sentinels[0]]
+
+    fake_ray.wait.side_effect = wait
+    fake_ray.get.side_effect = death
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    monkeypatch.setattr(ray_module.threading, "Thread", _InlineThread)
+    group = object()
+    executor = _make_executor(fake_ray, actors, group=group)
+    callback = Mock()
+    executor.register_failure_callback(callback)
+
+    executor._start_worker_monitor()
+
+    for actor in actors:
+        actor.wait_for_exit.options.assert_called_once_with(concurrency_group="monitor")
+    if closed_first:
+        assert not executor._is_failed
+        callback.assert_not_called()
+        fake_ray.kill.assert_not_called()
+        return
+    fake_ray.get.assert_called_once_with(sentinels[1])
+    assert executor._is_failed
+    assert executor._closed
+    callback.assert_called_once_with()
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+    fake_ray.wait.assert_called_once()
+    with pytest.raises(RuntimeError, match="closed"):
+        executor.check_health()
+
+
+def test_worker_monitor_does_not_keep_executor_alive(monkeypatch):
+    fake_ray = Mock()
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    started = []
+    monkeypatch.setattr(
+        ray_module.threading,
+        "Thread",
+        lambda target, **kwargs: started.append(target) or SimpleNamespace(start=lambda: None),
+    )
+    executor = _make_executor(fake_ray, [Mock()])
+    executor._start_worker_monitor()
+    (monitor,) = started
+    assert all(executor is not cell.cell_contents for cell in monitor.__closure__ or ())
+
+
+@pytest.mark.parametrize(
+    "bundles, required, error",
+    [
+        ([{"GPU": 1}] * 4, 4, None),
+        ([{"GPU": 1}] * 8, 4, None),
+        ([{"CPU": 1}, {"GPU": 0.5}, {"GPU": 1}], 1, None),
+        ([{"GPU": 1}] * 2, 4, "requires 4 full GPUs.*only 2"),
+        ([{"GPU": 0.5}] * 4, 4, "requires 4 full GPUs.*only 0"),
+        ([{"CPU": 4}], 1, "requires 1 full GPUs.*only 0"),
+        ([{"GPU": 4}], 4, None),
+        ([{"GPU": 2}, {"GPU": 2}], 4, None),
+        ([{"GPU": 1.5}, {"GPU": 1.5}], 3, "requires 3 full GPUs.*only 2"),
+    ],
+)
+def test_reused_placement_group_requires_full_gpu_capacity(monkeypatch, bundles, required, error):
+    group = SimpleNamespace(bundle_specs=bundles)
+    fake_ray = Mock()
+    fake_ray.util.get_current_placement_group.return_value = group
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor.od_config = SimpleNamespace(num_gpus=required)
+
+    if error is not None:
+        with pytest.raises(ValueError, match=error):
+            executor._get_or_create_placement_group()
+    else:
+        actual_group, owned = executor._get_or_create_placement_group()
+        assert actual_group is group
+        assert not owned
+
+    fake_ray.util.placement_group.assert_not_called()
+    fake_ray.util.remove_placement_group.assert_not_called()
+
+
+@pytest.mark.parametrize("failure_phase", ["creation", "discovery", "port", "initialization"])
+def test_worker_startup_failure_cleans_up_registered_actors(monkeypatch, failure_phase):
+    actors = [Mock(), Mock()]
+    for actor in actors:
+        actor.get_node_ip.remote.return_value = "10.0.0.1"
+        actor.get_open_port.remote.return_value = 23456
+    factory = Mock()
+    error = RuntimeError("startup failed")
+    factory.remote.side_effect = [actors[0], error] if failure_phase == "creation" else actors
+    fake_ray = Mock()
+    fake_ray.is_initialized.return_value = True
+    fake_ray.remote.return_value = lambda cls: factory
+    call_count = 0
+
+    def get(values, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        assert kwargs["timeout"] == ray_module._WORKER_INIT_TIMEOUT_S
+        if call_count == {"discovery": 1, "port": 2, "initialization": 3}.get(failure_phase):
+            raise error
+        return values
+
+    fake_ray.get.side_effect = get
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    monkeypatch.setattr(ray_module, "get_ip", lambda: "10.0.0.1")
+    monkeypatch.setattr(ray_module, "PlacementGroupSchedulingStrategy", Mock())
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor.od_config = SimpleNamespace(num_gpus=2)
+    group = object()
+    executor._get_or_create_placement_group = Mock(return_value=(group, True))
+
+    with pytest.raises(RuntimeError, match="startup failed"):
+        executor._init_executor()
+
+    expected = actors[:1] if failure_phase == "creation" else actors
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == expected
+    for actor in expected:
+        actor.shutdown.remote.assert_not_called()
+        actor.wait_for_exit.options.assert_not_called()
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+    assert executor.workers == []
+    assert executor._closed
+
+
+@pytest.mark.parametrize("failure_type", [TimeoutError, RuntimeError, KeyboardInterrupt])
+@pytest.mark.parametrize("cleanup_fails", [False, True])
+def test_placement_group_wait_failure_cleans_up_and_preserves_error(monkeypatch, failure_type, cleanup_fails):
+    fake_ray = Mock()
+    fake_ray.exceptions.GetTimeoutError = TimeoutError
+    fake_ray.util.get_current_placement_group.return_value = None
+    group = fake_ray.util.placement_group.return_value
+    original_error = failure_type("readiness failed")
+    fake_ray.get.side_effect = original_error
+    if cleanup_fails:
+        fake_ray.util.remove_placement_group.side_effect = RuntimeError("cleanup failed")
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor.od_config = SimpleNamespace(num_gpus=2)
+
+    expected_type = ValueError if failure_type is TimeoutError else failure_type
+    with pytest.raises(expected_type) as raised:
+        executor._get_or_create_placement_group()
+
+    if failure_type is TimeoutError:
+        assert raised.value.__cause__ is original_error
+        assert "Cannot reserve 2 Ray GPU bundle(s)" in str(raised.value)
+    else:
+        assert raised.value is original_error
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+
+
+def test_health_probes_use_separate_concurrency_group(monkeypatch):
+    fake_ray = Mock()
+    fake_ray.get.return_value = [True, True]
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actors = [Mock(), Mock()]
+    executor = _make_executor(fake_ray, actors)
+
+    executor.check_health()
+
+    for actor in actors:
+        actor.check_alive.options.assert_called_once_with(concurrency_group="health")
+        actor.check_alive.options.return_value.remote.assert_called_once_with()
+        actor.check_alive.remote.assert_not_called()
+    fake_ray.get.assert_called_once_with(
+        [actor.check_alive.options.return_value.remote.return_value for actor in actors],
+        timeout=ray_module._HEALTH_CHECK_TIMEOUT_S,
+    )
+    assert not executor._is_failed
+
+
+@pytest.mark.parametrize("probe_result", ["error", "not_alive"])
+def test_health_failure_releases_actors_and_notifies(monkeypatch, probe_result):
+    fake_ray = Mock()
+    if probe_result == "error":
+        fake_ray.get.side_effect = RuntimeError("actor died")
+    else:
+        fake_ray.get.return_value = [True, False]
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actors = [Mock(), Mock()]
+    group = object()
+    executor = _make_executor(fake_ray, actors, group=group)
+    callback = Mock()
+    executor.register_failure_callback(callback)
+
+    with pytest.raises(ray_module.EngineDeadError):
+        executor.check_health()
+
+    assert executor._is_failed
+    assert executor._closed
+    callback.assert_called_once_with()
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+
+
+def _make_request_executor(monkeypatch, prompts, component="text_encoder", allgather=True):
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor._closed = False
+    executor._is_failed = False
+    executor._state_lock = threading.Lock()
+    executor._failure_callbacks = []
+    executor._finalizer = Mock()
+    executor.od_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(data_parallel_size=len(prompts)),
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": [component],
+            "layer_options": {component: {"weight_transfer": "allgather" if allgather else "rank-local"}},
+        },
+    )
+    monkeypatch.setattr(ray_module, "build_request_batch_sampling_params_key", lambda req: ())
+    requests = []
+    for index, prompt in enumerate(prompts):
+        request_id = str(index)
+        req = SimpleNamespace(request_id=request_id, prompt=prompt, sampling_params=SimpleNamespace(extra_args=None))
+        requests.append(SimpleNamespace(request_id=request_id, req=req, diffusion_kv_metadata=None))
+    return executor, SimpleNamespace(scheduled_new_reqs=requests, kv_prefetch_job=None)
+
+
+@pytest.mark.parametrize("field", ["prompt_embeds", "negative_prompt_embeds"])
+def test_text_encoder_allgather_rejects_mismatched_embedding_presence(monkeypatch, field):
+    executor, scheduled = _make_request_executor(monkeypatch, ["first", {"prompt": "second", field: object()}])
+    executor.collective_rpc = Mock()
+    with pytest.raises(ValueError, match="same positive/negative prompt embedding fields"):
+        executor.execute_request(scheduled)
+    executor.collective_rpc.assert_not_called()
+
+
+def test_text_encoder_allgather_accepts_matching_precomputed_embedding_paths(monkeypatch):
+    executor, scheduled = _make_request_executor(
+        monkeypatch, [{"prompt_embeds": object()}, {"prompt_embeds": object()}]
+    )
+    results = [DiffusionOutput(output="first"), DiffusionOutput(output="second")]
+    executor.collective_rpc = Mock(return_value=[{"dp_rank": rank, "output": out} for rank, out in enumerate(results)])
+    output = executor.execute_request(scheduled)
+    assert [item.result for item in output.runner_outputs] == results
+    executor.collective_rpc.assert_called_once()
+
+
+@pytest.mark.parametrize("allgather", [False, True])
+def test_single_request_allgather_timeout_option(monkeypatch, allgather):
+    executor, scheduled = _make_request_executor(monkeypatch, ["prompt"], allgather=allgather)
+    executor.collective_rpc = Mock(return_value=DiffusionOutput(output="done"))
+    executor.execute_request(scheduled)
+    kwargs = executor.collective_rpc.call_args.kwargs
+    if allgather:
+        assert kwargs["timeout"] == ray_module._DLO_DP_WAVE_TIMEOUT_S
+    else:
+        assert "timeout" not in kwargs
+
+
+def test_single_request_allgather_timeout_fails_executor_and_cleans_up(monkeypatch):
+    executor, scheduled = _make_request_executor(monkeypatch, ["prompt"])
+    callback = Mock()
+    executor.register_failure_callback(callback)
+    executor.collective_rpc = Mock(side_effect=TimeoutError("AllGather timed out"))
+    output = executor.execute_request(scheduled)
+    assert executor._is_failed
+    assert executor._closed
+    executor._finalizer.assert_called_once()
+    callback.assert_called_once()
+    assert "AllGather timed out" in output.runner_outputs[0].result.error
+
+
+@pytest.mark.parametrize("wrapper", ["media", "output", "batch"])
+def test_move_to_cpu_handles_typed_video_media(wrapper):
+    from vllm_omni.diffusion.media import (
+        DiffusionMediaOutput,
+        VideoMediaOutput,
+        VideoTensorEncoding,
+        VideoTensorLayout,
+        VideoTensorSpec,
+        VideoValueRange,
+    )
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    tensor = torch.ones((1, 3, 2, 4, 4), requires_grad=True)
+    media = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=tensor,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.ZERO_TO_ONE,
+            ),
+        ),
+        prepared_for_transport=True,
+    )
+    value = media
+    if wrapper != "media":
+        value = DiffusionOutput(media=media, to_cpu=False)
+    if wrapper == "batch":
+        value = BatchRunnerOutput.from_list([RunnerOutput(request_id="video", finished=True, result=value)])
+
+    moved = ray_module._move_to_cpu(value)
+    if wrapper == "batch":
+        moved = moved.runner_outputs[0].result
+    if wrapper != "media":
+        moved = moved.media
+
+    moved.validate()
+    assert moved.video.tensor.device.type == "cpu"
+    assert not moved.video.tensor.requires_grad
+    assert moved.video.spec == media.video.spec
+    assert moved.prepared_for_transport
+    torch.testing.assert_close(moved.video.tensor, tensor.detach())
+    assert media.video.tensor is tensor
+    assert tensor.requires_grad

tests/diffusion/test_stage_diffusion_proc.py

@@ -191,3 +191,62 @@ async def test_proc_process_request_with_batching_async_output():
         time_gap = elapsed_time - base_time
         assert time_gap > time_gap_std - eps and time_gap < time_gap_std + eps
         base_time = elapsed_time
+
+
+class _FailingExecutor:
+    def __init__(self, is_dead: bool = False) -> None:
+        self.is_dead = is_dead
+        self.callbacks = []
+
+    def register_failure_callback(self, callback) -> None:
+        self.callbacks.append(callback)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failed_before_loop", [False, True])
+async def test_run_loop_exits_when_executor_fails_between_requests(failed_before_loop):
+    """An idle executor failure must reach the client without waiting for a request."""
+    import threading
+
+    import zmq
+    import zmq.asyncio
+
+    ctx = zmq.asyncio.Context()
+    request_socket = ctx.socket(zmq.PUSH)
+    response_socket = ctx.socket(zmq.PULL)
+    request_port = request_socket.bind_to_random_port("tcp://127.0.0.1")
+    response_port = response_socket.bind_to_random_port("tcp://127.0.0.1")
+    executor = _FailingExecutor(is_dead=failed_before_loop)
+    stage_proc = StageDiffusionProc.__new__(StageDiffusionProc)
+    stage_proc._engine = type("_Engine", (), {"executor": executor})()
+    stage_proc._active_tasks = None
+    stage_proc._fatal_event = None
+    try:
+        loop_task = asyncio.create_task(
+            stage_proc.run_loop(f"tcp://127.0.0.1:{request_port}", f"tcp://127.0.0.1:{response_port}")
+        )
+        if not failed_before_loop:
+            for _ in range(500):
+                if executor.callbacks:
+                    break
+                await asyncio.sleep(0.01)
+            assert executor.callbacks, "run_loop did not register an executor failure callback"
+            # Executor monitors report failures from their own threads.
+            monitor = threading.Thread(target=executor.callbacks[0])
+            monitor.start()
+            monitor.join()
+
+        with pytest.raises(RuntimeError, match="permanent failure"):
+            await asyncio.wait_for(loop_task, timeout=5)
+        assert await asyncio.wait_for(response_socket.recv(), timeout=5) == StageDiffusionProc.DIFFUSION_PROC_DEAD
+        assert len(executor.callbacks) == 1
+
+        # A late failure after run_loop has exited must not raise.
+        closed_loop = asyncio.new_event_loop()
+        closed_loop.close()
+        stage_proc._watch_executor_failure(closed_loop)
+        executor.callbacks[-1]()
+    finally:
+        request_socket.close(linger=0)
+        response_socket.close(linger=0)
+        ctx.term()

tests/diffusion/test_worker_wrapper_base.py

@@ -770,3 +770,60 @@ class TestCustomPipelineWorkerExtension:
         # Verify load_model was called twice with different pipelines
         assert wrapper.worker.load_model.call_count == 2
         assert wrapper.worker.init_lora_manager.call_count == 2
+
+
+@pytest.mark.parametrize(
+    ("init_method", "rank", "local_rank", "expected_addr", "expected_port"),
+    [
+        (None, 1, 1, "localhost", "12345"),
+        ("tcp://10.0.0.2:23456", 7, 0, "unrelated-host", "9999"),
+    ],
+)
+def test_worker_rendezvous_preserves_local_device_rank(
+    monkeypatch, init_method, rank, local_rank, expected_addr, expected_port
+):
+    """TCP rendezvous must not confuse a global rank with the actor's GPU index."""
+    import os
+    from contextlib import nullcontext
+    from unittest.mock import Mock
+
+    from vllm_omni.diffusion.worker import diffusion_worker as worker_module
+
+    class ReachedRendezvousError(Exception):
+        pass
+
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = rank
+    worker.local_rank = local_rank
+    worker.distributed_init_method = init_method
+    worker.od_config = SimpleNamespace(num_gpus=8, master_port=12345, moe_backend="auto")
+    config = SimpleNamespace(kernel_config=SimpleNamespace(ir_op_priority=None))
+    platform = Mock()
+    platform.get_torch_device.return_value = "cuda:0"
+    rendezvous = Mock(side_effect=ReachedRendezvousError)
+    monkeypatch.setenv("MASTER_ADDR", "unrelated-host")
+    monkeypatch.setenv("MASTER_PORT", "9999")
+    # Restore these variables after init_device publishes the worker identity.
+    for key in ("LOCAL_RANK", "RANK", "WORLD_SIZE"):
+        monkeypatch.setenv(key, "stale")
+    monkeypatch.setattr(worker_module, "current_omni_platform", platform)
+    monkeypatch.setattr(worker_module, "create_diffusion_vllm_config", Mock(return_value=config))
+    monkeypatch.setattr(worker_module, "_resolve_ir_op_priority", Mock(return_value=[]))
+    monkeypatch.setattr(worker_module, "set_forward_context", lambda **kwargs: nullcontext())
+    monkeypatch.setattr(worker_module, "set_current_vllm_config", lambda config: nullcontext())
+    monkeypatch.setattr(worker_module, "init_distributed_environment", rendezvous)
+
+    with pytest.raises(ReachedRendezvousError):
+        worker.init_device()
+
+    rendezvous.assert_called_once_with(
+        world_size=8,
+        rank=rank,
+        distributed_init_method=init_method or "env://",
+        local_rank=local_rank,
+    )
+    platform.get_torch_device.assert_called_once_with(local_rank)
+    assert os.environ["MASTER_ADDR"] == expected_addr
+    assert os.environ["MASTER_PORT"] == expected_port
+    assert os.environ["LOCAL_RANK"] == str(local_rank)
+    assert os.environ["RANK"] == str(rank)

tests/engine/test_async_omni_engine_stage_init.py

@@ -71,6 +71,34 @@ def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch):
     assert env_key not in os.environ
 
 
+@pytest.mark.parametrize("typed", [False, True])
+def test_build_ray_diffusion_config_preserves_explicit_stage_env(monkeypatch, typed):
+    from vllm_omni.engine import stage_init_utils as init_mod
+
+    env = {"CUSTOM_PLUGIN_SETTING": "stage", "OMP_NUM_THREADS": 2}
+    runtime_cfg = OmniStageRuntimeConfig(env=env) if typed else {"env": env}
+    metadata = types.SimpleNamespace(
+        runtime_cfg=runtime_cfg, stage_id=1, cfg_kv_collect_func=None, default_sampling_params=None
+    )
+    config = types.SimpleNamespace(
+        distributed_executor_backend="ray", parallel_config=types.SimpleNamespace(world_size=2)
+    )
+    monkeypatch.setattr(init_mod, "build_engine_args_dict", lambda *args: {})
+    monkeypatch.setattr(init_mod, "OmniDiffusionConfig", lambda **kwargs: config)
+    monkeypatch.setattr(
+        init_mod,
+        "current_omni_platform",
+        types.SimpleNamespace(device_control_env_var=None, get_device_count=lambda: 0),
+    )
+
+    result = init_mod.build_diffusion_config("model", {}, metadata)
+
+    assert result is config
+    assert result.ray_worker_env == {"CUSTOM_PLUGIN_SETTING": "stage", "OMP_NUM_THREADS": "2"}
+    assert result.num_gpus == 2
+    assert env["OMP_NUM_THREADS"] == 2
+
+
 @pytest.mark.parametrize(
     "invalid_key,invalid_value",
     [("INVALID=ENV", "value"), ("INVALID\0ENV", "value"), ("INVALID_ENV", "value\0")],
@@ -702,7 +730,11 @@ def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkey
         assert kwargs["visible_devices"] == "1"
         return [42]
 
-    od_config = types.SimpleNamespace(max_num_seqs=4, parallel_config=types.SimpleNamespace(world_size=1))
+    od_config = types.SimpleNamespace(
+        max_num_seqs=4,
+        distributed_executor_backend="mp",
+        parallel_config=types.SimpleNamespace(world_size=1),
+    )
     monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
     monkeypatch.setattr(startup_mod, "acquire_device_locks", acquire)
     monkeypatch.setattr(
@@ -796,6 +828,7 @@ def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(mo
     od_config = types.SimpleNamespace(
         max_num_seqs=8,
         step_execution=True,
+        distributed_executor_backend="mp",
         parallel_config=types.SimpleNamespace(world_size=1),
     )
     monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)

tests/engine/test_single_stage_mode.py

@@ -1219,7 +1219,11 @@ class TestSingleStageReplicaInitialization:
         runtime._init_visible_devices_baseline = "0"
 
         mocker.patch.object(runtime_mod, "inject_kv_stage_info")
-        od_config = SimpleNamespace(max_num_seqs=4, parallel_config=SimpleNamespace(world_size=1))
+        od_config = SimpleNamespace(
+            max_num_seqs=4,
+            parallel_config=SimpleNamespace(world_size=1),
+            distributed_executor_backend="mp",
+        )
         mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
         mock_register = mocker.patch(
             "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
@@ -1352,7 +1356,11 @@ class TestSingleStageReplicaInitialization:
         runtime._init_visible_devices_baseline = "0"
 
         mocker.patch.object(runtime_mod, "inject_kv_stage_info")
-        od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
+        od_config = SimpleNamespace(
+            max_num_seqs=None,
+            parallel_config=SimpleNamespace(world_size=1),
+            distributed_executor_backend="mp",
+        )
         mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
         mocker.patch(
             "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",

vllm_omni/config/omni_config.py

@@ -1145,6 +1145,9 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset(
         "disable_autocast",
     }
 )
+# Runtime-populated OmniDiffusionConfig state (init=False) that is never
+# user-configurable and therefore not part of the projection.
+_DIFFUSION_INTERNAL_FIELDS = frozenset({"ray_worker_env"})
 
 
 _STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS)

vllm_omni/diffusion/data.py

@@ -1005,6 +1005,9 @@ class OmniDiffusionConfig:
     # Worker extension class for custom functionality
     worker_extension_cls: str | None = None
 
+    # Internal transport of explicit stage runtime.env to remote Ray actors.
+    ray_worker_env: dict[str, str] = field(default_factory=dict, init=False, repr=False)
+
     # Custom pipeline arguments for custom pipelines
     custom_pipeline_args: dict[str, Any] | None = None
 

vllm_omni/diffusion/diffusion_engine.py

@@ -669,8 +669,10 @@ class DiffusionEngine:
                 self._poll_native_kv()
                 worker_execution_completed = True
             except Exception as exc:
+                # Shutdown can interrupt an in-flight executor call. Exit through
+                # the loop tail so pending RPCs are failed instead of left waiting.
                 if self._closed:
-                    return
+                    break
                 worker_execution_completed = False
                 logger.error(
                     "Execution failed for diffusion requests %s", sched_output.scheduled_request_ids, exc_info=True
@@ -1569,6 +1571,10 @@ class DiffusionEngine:
         for stream in pending_streams:
             self._put_queue_output(stream, closed_output)
 
+        # Interrupt remote inference before waiting for the thread blocked on Ray.
+        if self.od_config.distributed_executor_backend == "ray":
+            self.executor.shutdown()
+
         worker_thread = self.worker_thread
         if worker_thread is not None:
             if worker_thread.is_alive():

vllm_omni/diffusion/executor/abstract.py

@@ -49,7 +49,9 @@ class DiffusionExecutor(ABC):
                 )
             executor_class = distributed_executor_backend
         elif distributed_executor_backend == "ray":
-            raise NotImplementedError("ray backend is not yet supported.")
+            from vllm_omni.diffusion.executor.ray_executor import RayDiffusionExecutor
+
+            executor_class = RayDiffusionExecutor
         elif distributed_executor_backend == "mp":
             from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 

vllm_omni/diffusion/executor/ray_executor.py

@@ -0,0 +1,723 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Ray-based distributed executor for diffusion models.
+
+Ray owns worker placement and process lifecycle. PyTorch distributed still
+owns the model-parallel data plane: the executor gives every actor a global
+rank and a common rendezvous address before the actor constructs its worker.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import threading
+import weakref
+from collections import defaultdict
+from collections.abc import Callable
+from dataclasses import dataclass, fields
+from typing import TYPE_CHECKING, Any
+
+import torch
+from vllm.logger import init_logger
+from vllm.utils.import_utils import resolve_obj_by_qualname
+from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from vllm_omni.diffusion.data import (
+    DiffusionOutput,
+    OmniDiffusionConfig,
+)
+from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+from vllm_omni.diffusion.media import DiffusionMediaOutput
+from vllm_omni.diffusion.offloader.config import (
+    TEXT_ENCODER_COMPONENT,
+    any_selected_component_uses_allgather,
+    resolve_offload,
+)
+from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key
+
+if TYPE_CHECKING:
+    from ray.actor import ActorHandle
+    from ray.util.placement_group import PlacementGroup
+
+    from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
+    from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
+
+logger = init_logger(__name__)
+
+_PLACEMENT_GROUP_WAIT_TIMEOUT_S = 1800
+_WORKER_INIT_TIMEOUT_S = 7200
+_HEALTH_CHECK_TIMEOUT_S = 10
+# Must stay below StageDiffusionClient's wait for the Ray subprocess to exit.
+_WORKER_SHUTDOWN_TIMEOUT_S = 5.0
+_DLO_DP_WAVE_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", 600.0))
+
+try:
+    import ray
+    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+except ImportError:
+    ray = None  # type: ignore[assignment]
+    PlacementGroupSchedulingStrategy = None  # type: ignore[assignment,misc]
+
+
+def _worker_env(od_config: OmniDiffusionConfig) -> dict[str, str]:
+    """Copy inference settings, keeping worker identity local to each actor."""
+    prefixes = (
+        "VLLM_",
+        "OMNI_",
+        "DIFFUSION_",
+        "NCCL_",
+        "TORCH_NCCL_",
+        "UCX_",
+        "HF_",
+        "HUGGINGFACE_",
+        "HUGGING_FACE_HUB_",
+    )
+    names = {"PYTHONPATH", "CUDA_LAUNCH_BLOCKING", "OMP_NUM_THREADS"}
+    env = {key: value for key, value in os.environ.items() if key.startswith(prefixes) or key in names}
+    # Explicit stage settings may include plugin-specific variables outside
+    # the standard prefixes, and take precedence over the driver's defaults.
+    env.update(getattr(od_config, "ray_worker_env", {}))
+    worker_specific = {
+        "CUDA_VISIBLE_DEVICES",
+        "HIP_VISIBLE_DEVICES",
+        "ROCR_VISIBLE_DEVICES",
+        "ASCEND_RT_VISIBLE_DEVICES",
+        "MUSA_VISIBLE_DEVICES",
+        "ONEAPI_DEVICE_SELECTOR",
+        "VLLM_HOST_IP",
+        "VLLM_HOST_PORT",
+        "VLLM_NIXL_SIDE_CHANNEL_HOST",
+        "VLLM_LOCAL_RANK",
+        "MASTER_ADDR",
+        "MASTER_PORT",
+        "LOCAL_RANK",
+        "LOCAL_WORLD_SIZE",
+        "RANK",
+        "WORLD_SIZE",
+    }
+    return {key: value for key, value in env.items() if key not in worker_specific and not key.startswith("RAY_")}
+
+
+def _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool:
+    parallel_config = getattr(od_config, "parallel_config", None)
+    return bool(
+        (getattr(parallel_config, "data_parallel_size", 1) or 1) > 1
+        and any_selected_component_uses_allgather(od_config)
+    )
+
+
+def _is_empty_dp_prompt(prompt: object) -> bool:
+    """Return whether a DP request has no usable text prompt."""
+    if prompt is None:
+        return True
+    if isinstance(prompt, (str, list, tuple)):
+        return not prompt
+    if isinstance(prompt, dict):
+        return (
+            not prompt.get("prompt")
+            and not prompt.get("prompt_token_ids")
+            and not prompt.get("prompt_ids")
+            and prompt.get("prompt_embeds") is None
+        )
+    return False
+
+
+def _text_encoder_input_signature(prompt: object) -> tuple[bool, bool]:
+    """Describe precomputed embeddings that change encoder forward counts."""
+    if not isinstance(prompt, dict):
+        return False, False
+    return prompt.get("prompt_embeds") is not None, prompt.get("negative_prompt_embeds") is not None
+
+
+def _uses_text_encoder_allgather(config: object) -> bool:
+    resolved = resolve_offload(config)
+    return resolved.offloads(TEXT_ENCODER_COMPONENT) and resolved.uses_allgather(TEXT_ENCODER_COMPONENT)
+
+
+def _move_to_cpu(value: Any) -> Any:
+    """Recursively detach tensors before Ray serializes a worker response."""
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    if isinstance(value, DiffusionMediaOutput):
+        return value.to_cpu()
+    if isinstance(value, DiffusionOutput):
+        for output_field in fields(value):
+            setattr(value, output_field.name, _move_to_cpu(getattr(value, output_field.name)))
+        return value
+    if isinstance(value, dict):
+        return {key: _move_to_cpu(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_move_to_cpu(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_move_to_cpu(item) for item in value)
+
+    # RunnerOutput and BatchRunnerOutput are kept out of the module imports to
+    # avoid pulling the complete worker stack into the driver at import time.
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    if isinstance(value, RunnerOutput):
+        value.result = _move_to_cpu(value.result)
+        return value
+    if isinstance(value, BatchRunnerOutput):
+        value.runner_outputs = [_move_to_cpu(output) for output in value.runner_outputs]
+        return value
+    return value
+
+
+@dataclass
+class _RayWorkerMetadata:
+    worker: ActorHandle
+    rank: int = -1
+    ip: str = ""
+
+
+@dataclass
+class _RayExecutorResources:
+    """Resources owned by the executor's finalizer."""
+
+    workers: list[_RayWorkerMetadata]
+    placement_group: PlacementGroup | None = None
+    owns_placement_group: bool = False
+
+    def drain(self, timeout_s: float) -> None:
+        """Run worker shutdown on every actor, waiting at most ``timeout_s``.
+
+        Shutdown RPCs queue behind in-flight inference, so actors that do not
+        finish in time are left for ``__call__`` to kill.
+        """
+        if ray is None or not self.workers:
+            return
+
+        pending: dict[Any, int] = {}
+        for metadata in self.workers:
+            try:
+                pending[metadata.worker.shutdown.remote()] = metadata.rank
+            except Exception as exc:
+                logger.warning("Failed to request shutdown of Ray diffusion worker rank %d: %s", metadata.rank, exc)
+        if not pending:
+            return
+
+        done, not_done = ray.wait(list(pending), num_returns=len(pending), timeout=timeout_s)
+        for ref in done:
+            try:
+                ray.get(ref)
+            except Exception as exc:
+                logger.warning("Ray diffusion worker rank %d did not shut down cleanly: %s", pending[ref], exc)
+        if not_done:
+            logger.warning(
+                "Ray diffusion worker rank(s) %s did not shut down within %.1fs; killing them",
+                sorted(pending[ref] for ref in not_done),
+                timeout_s,
+            )
+
+    def __call__(self) -> None:
+        if ray is None:
+            return
+
+        # Actors that were drained have already released their resources. The
+        # rest may be blocked in inference or a collective, so kill them all.
+        for metadata in self.workers:
+            try:
+                ray.kill(metadata.worker, no_restart=True)
+            except Exception as exc:
+                logger.warning("Failed to kill Ray diffusion worker rank %d: %s", metadata.rank, exc)
+        self.workers.clear()
+
+        if self.owns_placement_group and self.placement_group is not None:
+            try:
+                ray.util.remove_placement_group(self.placement_group)
+            except Exception as exc:
+                logger.warning("Failed to remove Ray diffusion placement group: %s", exc)
+            self.placement_group = None
+
+
+class RayDiffusionWorkerWrapper:
+    """Ray actor that constructs one rank of the existing DiffusionWorker."""
+
+    def __init__(self, provisional_rank: int) -> None:
+        self.rank = provisional_rank
+        self.worker = None
+
+    def get_node_ip(self) -> str:
+        return get_ip()
+
+    def get_open_port(self) -> int:
+        return get_open_port()
+
+    def init_worker(self, od_config: OmniDiffusionConfig, rank: int, distributed_init_method: str) -> None:
+        from vllm_omni.platforms import current_omni_platform
+        from vllm_omni.plugins import load_omni_general_plugins
+
+        load_omni_general_plugins()
+
+        from vllm_omni.diffusion.worker.diffusion_worker import WorkerWrapperBase
+
+        worker_cls_path = current_omni_platform.get_diffusion_worker_cls()
+        self.rank = rank
+        self.worker = WorkerWrapperBase(
+            gpu_id=0,
+            rank=self.rank,
+            distributed_init_method=distributed_init_method,
+            od_config=od_config,
+            worker_extension_cls=od_config.worker_extension_cls,
+            custom_pipeline_args=getattr(od_config, "custom_pipeline_args", None),
+            base_worker_class=resolve_obj_by_qualname(worker_cls_path),
+        )
+
+    def _is_primary_in_replica(self) -> bool:
+        from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
+
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_classifier_free_guidance_rank,
+            get_pipeline_parallel_rank,
+            get_sequence_parallel_rank,
+        )
+
+        return (
+            get_sequence_parallel_rank() == 0
+            and get_classifier_free_guidance_rank() == 0
+            and get_tensor_model_parallel_rank() == 0
+            and get_pipeline_parallel_rank() == 0
+        )
+
+    def execute_rpc(
+        self,
+        method: str,
+        args: tuple = (),
+        kwargs: dict[str, Any] | None = None,
+        output_rank: int | None = None,
+        exec_all_ranks: bool = False,
+        primary_replies_only: bool = False,
+    ) -> dict[str, Any]:
+        if self.worker is None:
+            raise RuntimeError("Ray diffusion worker is not initialized")
+
+        should_execute = exec_all_ranks or output_rank is None or output_rank == self.rank
+        should_reply = output_rank is None or output_rank == self.rank
+        if primary_replies_only:
+            should_reply = self._is_primary_in_replica()
+
+        result = None
+        if should_execute:
+            result = self.worker.execute_method(method, *(args or ()), **(kwargs or {}))
+        if should_reply:
+            result = _move_to_cpu(result)
+        else:
+            result = None
+        return {"rank": self.rank, "replied": should_reply, "result": result}
+
+    def check_alive(self) -> bool:
+        return self.worker is not None
+
+    def wait_for_exit(self) -> None:
+        """Block until the actor dies; the executor monitor waits on this call."""
+        threading.Event().wait()
+
+    def shutdown(self) -> None:
+        worker, self.worker = self.worker, None
+        if worker is not None:
+            worker.shutdown()
+
+
+class RayDiffusionExecutor(DiffusionExecutor):
+    """Execute one diffusion worker per Ray GPU actor."""
+
+    def _init_executor(self) -> None:
+        if ray is None:
+            raise ImportError(
+                "Ray is required for distributed_executor_backend='ray'. "
+                'Install the project dependencies or `pip install "ray[default]"`.'
+            )
+
+        self._closed = False
+        self._is_failed = False
+        self._state_lock = threading.Lock()
+        self._failure_callbacks: list[Callable[[], None]] = []
+        self.workers: list[_RayWorkerMetadata] = []
+
+        if not ray.is_initialized():
+            logger.info("Initializing Ray for the diffusion executor")
+            ray.init()
+
+        placement_group, owns_placement_group = self._get_or_create_placement_group()
+        self._resources = _RayExecutorResources(
+            workers=self.workers,
+            placement_group=placement_group,
+            owns_placement_group=owns_placement_group,
+        )
+        self._finalizer = weakref.finalize(self, self._resources)
+
+        try:
+            self._init_workers(placement_group)
+        except Exception:
+            self._shutdown(graceful=False)
+            raise
+        self._start_worker_monitor()
+
+    @property
+    def is_dead(self) -> bool:
+        return self._closed or self._is_failed
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("RayDiffusionExecutor is closed")
+        if self._is_failed:
+            raise EngineDeadError()
+
+    def _get_or_create_placement_group(self) -> tuple[PlacementGroup, bool]:
+        num_gpus = int(self.od_config.num_gpus or 1)
+        current_group = ray.util.get_current_placement_group()
+        if current_group is not None:
+            # A worker needs one full GPU within a single bundle.
+            gpu_capacity = sum(int(bundle.get("GPU", 0)) for bundle in current_group.bundle_specs)
+            if gpu_capacity < num_gpus:
+                raise ValueError(
+                    f"Ray diffusion requires {num_gpus} full GPUs, "
+                    f"but the placement group can accommodate only {gpu_capacity} one-GPU workers."
+                )
+            logger.info("Using the current Ray placement group for diffusion workers")
+            return current_group, False
+
+        placement_group = ray.util.placement_group([{"GPU": 1} for _ in range(num_gpus)], strategy="PACK")
+        logger.info("Waiting for a Ray placement group with %d GPU bundle(s)", num_gpus)
+        try:
+            ray.get(placement_group.ready(), timeout=_PLACEMENT_GROUP_WAIT_TIMEOUT_S)
+        except BaseException as exc:
+            # The executor finalizer is not registered until this method returns.
+            try:
+                ray.util.remove_placement_group(placement_group)
+            except Exception:
+                logger.warning("Failed to remove placement group after startup failure", exc_info=True)
+            if isinstance(exc, ray.exceptions.GetTimeoutError):
+                raise ValueError(
+                    f"Cannot reserve {num_gpus} Ray GPU bundle(s) within "
+                    f"{_PLACEMENT_GROUP_WAIT_TIMEOUT_S}s; check `ray status`."
+                ) from exc
+            raise
+        return placement_group, True
+
+    def _init_workers(self, placement_group: PlacementGroup) -> None:
+        num_gpus = int(self.od_config.num_gpus or 1)
+        driver_ip = get_ip()
+        actor_cls = ray.remote(
+            num_cpus=0,
+            num_gpus=1,
+            runtime_env={"env_vars": _worker_env(self.od_config)},
+            max_concurrency=1,
+            concurrency_groups={"health": 1, "monitor": 1},
+            scheduling_strategy=PlacementGroupSchedulingStrategy(
+                placement_group=placement_group,
+                placement_group_capture_child_tasks=True,
+            ),
+        )(RayDiffusionWorkerWrapper)
+
+        # Register actors immediately so partial startup failures can clean them up.
+        for rank in range(num_gpus):
+            actor = actor_cls.remote(provisional_rank=rank)
+            self.workers.append(_RayWorkerMetadata(worker=actor, rank=rank))
+
+        worker_ips = ray.get(
+            [item.worker.get_node_ip.remote() for item in self.workers],
+            timeout=_WORKER_INIT_TIMEOUT_S,
+        )
+        for item, worker_ip in zip(self.workers, worker_ips):
+            item.ip = worker_ip
+
+        ip_counts: dict[str, int] = defaultdict(int)
+        for item in self.workers:
+            ip_counts[item.ip] += 1
+
+        # Keep ranks from the same host contiguous. Prefer the driver host for
+        # rank 0 when it owns a GPU; otherwise choose deterministically by IP.
+        self.workers.sort(key=lambda item: (0 if item.ip == driver_ip else 1, ip_counts[item.ip], item.ip))
+        for rank, item in enumerate(self.workers):
+            item.rank = rank
+
+        unique_ips = {item.ip for item in self.workers}
+        master_addr = "127.0.0.1" if len(unique_ips) == 1 else self.workers[0].ip
+        master_port = ray.get(self.workers[0].worker.get_open_port.remote(), timeout=_WORKER_INIT_TIMEOUT_S)
+        distributed_init_method = get_distributed_init_method(master_addr, master_port)
+
+        logger.info(
+            "Initializing %d Ray diffusion worker(s) across %d node(s)",
+            num_gpus,
+            len(unique_ips),
+        )
+        futures = [
+            item.worker.init_worker.remote(self.od_config, item.rank, distributed_init_method) for item in self.workers
+        ]
+        ray.get(futures, timeout=_WORKER_INIT_TIMEOUT_S)
+        logger.info("All %d Ray diffusion workers initialized", num_gpus)
+
+    def _start_worker_monitor(self) -> None:
+        """Fail the executor as soon as any actor dies, even while idle."""
+        sentinels = {
+            item.worker.wait_for_exit.options(concurrency_group="monitor").remote(): item.rank for item in self.workers
+        }
+        if not sentinels:
+            return
+        # Keep the thread from holding the executor alive; its finalizer kills
+        # the actors, which releases the wait below.
+        self_ref = weakref.ref(self)
+
+        def _monitor() -> None:
+            try:
+                done, _ = ray.wait(list(sentinels), num_returns=1)
+            except Exception:
+                return
+            executor = self_ref()
+            if executor is None or executor._closed:
+                return
+            ref = done[0]
+            try:
+                ray.get(ref)
+                exc: BaseException = RuntimeError(f"Ray diffusion worker rank {sentinels[ref]} exited unexpectedly")
+            except Exception as actor_exc:
+                exc = actor_exc
+            executor._fail(exc)
+
+        threading.Thread(target=_monitor, daemon=True, name="ray-diffusion-worker-monitor").start()
+
+    def _fail(self, exc: BaseException) -> None:
+        """Mark the executor dead, release every actor, then notify listeners."""
+        with self._state_lock:
+            if self._closed or self._is_failed:
+                return
+            self._is_failed = True
+        logger.error("Ray diffusion executor failed: %s", exc)
+        self._shutdown(graceful=False)
+        for callback in self._failure_callbacks:
+            try:
+                callback()
+            except Exception:
+                logger.exception("Diffusion executor failure callback raised")
+
+    def register_failure_callback(self, callback: Callable[[], None]) -> None:
+        self._failure_callbacks.append(callback)
+
+    def collective_rpc(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple = (),
+        kwargs: dict | None = None,
+        unique_reply_rank: int | None = None,
+        exec_all_ranks: bool = False,
+    ) -> Any:
+        self._ensure_open()
+
+        # ``None + exec_all_ranks`` is the multiprocessing executor's
+        # rank-local-DP response mode: all ranks compute, but only the primary
+        # rank of each DP replica returns an output.
+        primary_replies_only = unique_reply_rank is None and exec_all_ranks
+        execute_all_ranks = unique_reply_rank is None or exec_all_ranks
+        try:
+            futures = [
+                item.worker.execute_rpc.remote(
+                    method=method,
+                    args=args,
+                    kwargs=kwargs or {},
+                    output_rank=unique_reply_rank,
+                    exec_all_ranks=execute_all_ranks,
+                    primary_replies_only=primary_replies_only,
+                )
+                for item in self.workers
+            ]
+            envelopes = ray.get(futures, timeout=timeout)
+        except ray.exceptions.GetTimeoutError as exc:
+            self._fail(exc)
+            raise TimeoutError(f"RPC call to {method} timed out") from exc
+        except ray.exceptions.RayActorError as exc:
+            self._fail(exc)
+            raise EngineDeadError() from exc
+        except ray.exceptions.RayTaskError as exc:
+            # ray.get raises on the first failed rank without waiting for its
+            # peers. They may still be blocked in a collective, so no worker
+            # can safely accept another RPC from this executor.
+            self._fail(exc)
+            raise EngineDeadError() from exc
+
+        replies = sorted(
+            (envelope for envelope in envelopes if envelope["replied"]),
+            key=lambda envelope: envelope["rank"],
+        )
+        if unique_reply_rank is not None:
+            for envelope in replies:
+                if envelope["rank"] == unique_reply_rank:
+                    return envelope["result"]
+            raise RuntimeError(f"Ray diffusion rank {unique_reply_rank} did not return a response for {method}")
+        return [envelope["result"] for envelope in replies]
+
+    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
+        from vllm_omni.diffusion.sched.interface import validate_new_request_data_identity
+        from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+        self._ensure_open()
+        new_reqs = scheduler_output.scheduled_new_reqs
+        for new_req in new_reqs:
+            validate_new_request_data_identity(new_req)
+
+        if len(new_reqs) > 1 and _uses_dlo_dp_concurrency(self.od_config):
+            compatibility_keys = [build_request_batch_sampling_params_key(item.req) for item in new_reqs]
+            if any(key != compatibility_keys[0] for key in compatibility_keys[1:]):
+                raise ValueError(
+                    "Rank-local DP concurrency requires compatible shape, CFG, denoise schedule, "
+                    "output count, and LoRA settings for every request in a collective wave."
+                )
+            extra_args_signatures = {
+                json.dumps(getattr(item.req.sampling_params, "extra_args", None), sort_keys=True, default=repr)
+                for item in new_reqs
+            }
+            if len(extra_args_signatures) > 1:
+                raise ValueError("DP multi-concurrency requires identical extra_args for every request")
+            if _uses_text_encoder_allgather(self.od_config):
+                encoder_signatures = {_text_encoder_input_signature(item.req.prompt) for item in new_reqs}
+                if len(encoder_signatures) > 1:
+                    raise ValueError(
+                        "DLO text_encoder AllGather requires every concurrent request "
+                        "to provide the same positive/negative prompt embedding fields."
+                    )
+            empty_prompt_ids = [item.request_id for item in new_reqs if _is_empty_dp_prompt(item.req.prompt)]
+            if empty_prompt_ids:
+                raise ValueError(
+                    "DP multi-concurrency requires a non-empty prompt for every request; "
+                    f"empty prompt request IDs: {empty_prompt_ids}."
+                )
+
+            try:
+                results = self.collective_rpc(
+                    "execute_model",
+                    timeout=_DLO_DP_WAVE_TIMEOUT_S,
+                    args=(new_reqs, self.od_config, scheduler_output.kv_prefetch_job),
+                    unique_reply_rank=None,
+                    exec_all_ranks=True,
+                )
+                tagged_results = sorted(results, key=lambda item: item["dp_rank"])
+                outputs = []
+                for index, new_req in enumerate(new_reqs):
+                    result = tagged_results[index]["output"]
+                    if not isinstance(result, DiffusionOutput):
+                        raise RuntimeError(f"Unexpected DP response type [{index}]: {type(result)!r}")
+                    outputs.append(
+                        RunnerOutput(
+                            request_id=new_req.request_id,
+                            step_index=None,
+                            finished=True,
+                            result=result,
+                        )
+                    )
+                return BatchRunnerOutput.from_list(outputs)
+            except Exception as exc:
+                if isinstance(exc, TimeoutError):
+                    self._fail(exc)
+                return BatchRunnerOutput.from_list(
+                    [
+                        RunnerOutput(
+                            request_id=item.request_id,
+                            step_index=None,
+                            finished=True,
+                            result=DiffusionOutput.from_exception(exc),
+                        )
+                        for item in new_reqs
+                    ]
+                )
+
+        outputs = []
+        for new_req in new_reqs:
+            args: tuple = (new_req.req, self.od_config, scheduler_output.kv_prefetch_job)
+            if new_req.diffusion_kv_metadata is not None:
+                args += (new_req.diffusion_kv_metadata,)
+            try:
+                timeout_options: dict[str, Any] = {}
+                if any_selected_component_uses_allgather(self.od_config):
+                    timeout_options["timeout"] = _DLO_DP_WAVE_TIMEOUT_S
+                result = self.collective_rpc(
+                    "execute_model",
+                    args=args,
+                    unique_reply_rank=0,
+                    exec_all_ranks=True,
+                    **timeout_options,
+                )
+                if not isinstance(result, DiffusionOutput):
+                    raise RuntimeError(f"Unexpected response type: {type(result)!r}")
+            except Exception as exc:
+                if isinstance(exc, TimeoutError):
+                    self._fail(exc)
+                result = DiffusionOutput.from_exception(exc)
+            outputs.append(
+                RunnerOutput(
+                    request_id=new_req.request_id,
+                    step_index=None,
+                    finished=True,
+                    result=result,
+                )
+            )
+        return BatchRunnerOutput.from_list(outputs)
+
+    def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
+        from vllm_omni.diffusion.worker.utils import BatchRunnerOutput
+
+        self._ensure_open()
+        if len(scheduler_output.scheduled_new_reqs) <= 1:
+            return self.execute_request(scheduler_output)
+        if _uses_dlo_dp_concurrency(self.od_config):
+            return self.execute_request(scheduler_output)
+
+        result = self.collective_rpc(
+            "execute_model_batch",
+            args=(scheduler_output, self.od_config),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if not isinstance(result, BatchRunnerOutput):
+            raise RuntimeError(f"Unexpected response type for execute_batch: {type(result)!r}")
+        return result
+
+    def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
+        from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
+
+        self._ensure_open()
+        result = self.collective_rpc(
+            "execute_stepwise",
+            args=(scheduler_output,),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if not isinstance(result, BaseRunnerOutput):
+            raise RuntimeError(f"Unexpected response type for execute_step: {type(result)!r}")
+        return result
+
+    def check_health(self) -> None:
+        self._ensure_open()
+        try:
+            alive = ray.get(
+                [item.worker.check_alive.options(concurrency_group="health").remote() for item in self.workers],
+                timeout=_HEALTH_CHECK_TIMEOUT_S,
+            )
+        except Exception as exc:
+            self._fail(exc)
+            raise EngineDeadError() from exc
+        dead_ranks = [item.rank for item, is_alive in zip(self.workers, alive) if not is_alive]
+        if dead_ranks:
+            exc = RuntimeError(f"Ray diffusion worker ranks are not healthy: {dead_ranks}")
+            self._fail(exc)
+            raise EngineDeadError() from exc
+
+    def shutdown(self) -> None:
+        self._shutdown(graceful=True)
+
+    def _shutdown(self, graceful: bool) -> None:
+        with self._state_lock:
+            if self._closed:
+                return
+            self._closed = True
+        try:
+            if graceful:
+                self._resources.drain(_WORKER_SHUTDOWN_TIMEOUT_S)
+        except Exception as exc:
+            logger.warning("Failed to drain Ray diffusion workers: %s", exc)
+        finally:
+            self._finalizer()

vllm_omni/diffusion/stage_diffusion_client.py

@@ -524,14 +524,28 @@ class StageDiffusionClient(StageClientBase):
 
     def shutdown(self) -> None:
         self._shutting_down = True
+        shutdown_requested = False
         try:
             # A blocking send never returns once the subprocess is gone.
             self._request_socket.send(self._encoder.encode({"type": "shutdown"}), flags=zmq.NOBLOCK)
+            shutdown_requested = True
         except Exception:
             pass
 
         if self._proc_manager is not None and self._proc_manager.proc.is_alive():
-            self._proc_manager.shutdown(timeout=10)
+            if self._proc_manager.distributed_executor_backend == "ray":
+                # Let the subprocess kill its remote actors before terminating it.
+                stopped = shutdown_requested and self._proc_manager.wait_for_shutdown(timeout=10.0)
+                if not stopped:
+                    logger.warning(
+                        "Stage-%d Ray diffusion subprocess did not stop within %.1fs; "
+                        "falling back to signal-based shutdown",
+                        self.stage_id,
+                        10.0,
+                    )
+                    self._proc_manager.shutdown(timeout=10.0)
+            else:
+                self._proc_manager.shutdown(timeout=10.0)
 
         self._request_socket.close(linger=0)
         self._response_socket.close(linger=0)

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -130,6 +130,25 @@ class StageDiffusionProc:
         )
         self._fatal_event.set()
 
+    def _watch_executor_failure(self, loop: asyncio.AbstractEventLoop) -> None:
+        """Tear down when the executor fails between requests.
+
+        Executor monitors fire failure callbacks from their own threads, so
+        the signal is marshalled onto ``run_loop``'s event loop.
+        """
+        executor: DiffusionExecutor | None = getattr(self._engine, "executor", None)
+        if executor is None:
+            return
+
+        def _on_executor_failure() -> None:
+            # The loop is closed once run_loop has exited; nothing to signal.
+            with contextlib.suppress(RuntimeError):
+                loop.call_soon_threadsafe(self._signal_fatal_engine_failure, "diffusion executor failed")
+
+        executor.register_failure_callback(_on_executor_failure)
+        if executor.is_dead:
+            self._signal_fatal_engine_failure("diffusion executor failed before run_loop started")
+
     # ------------------------------------------------------------------
     # Initialization
     # ------------------------------------------------------------------
@@ -357,6 +376,7 @@ class StageDiffusionProc:
         # "DiffusionExecutor is closed" on every subsequent request.
         fatal_event = asyncio.Event()
         self._fatal_event = fatal_event
+        self._watch_executor_failure(asyncio.get_running_loop())
 
         async def _dispatch_request(
             request_id: str,
@@ -760,6 +780,7 @@ class StageDiffusionProcManager:
         )
         proc.start()
         self.proc = proc
+        self.distributed_executor_backend = od_config.distributed_executor_backend
         self.addresses = addresses
         self.manager_stopped = False
         self.failed_proc_name: str | None = None
@@ -797,6 +818,7 @@ class StageDiffusionProcManager:
         )
         proc.start()
         self.proc = proc
+        self.distributed_executor_backend = od_config.distributed_executor_backend
         self.addresses = addresses
         self.manager_stopped = False
         self.failed_proc_name = None
@@ -835,6 +857,12 @@ class StageDiffusionProcManager:
         self.manager_stopped = True
         shutdown([self.proc], timeout=timeout)
 
+    def wait_for_shutdown(self, timeout: float | None = None) -> bool:
+        """Wait for subprocess cleanup without sending a termination signal."""
+        self.manager_stopped = True
+        self.proc.join(timeout)
+        return not self.proc.is_alive()
+
     def sentinels(self) -> list[int]:
         return [self.proc.sentinel]
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -251,7 +251,9 @@ class DiffusionWorker:
         rank: int,
         od_config: OmniDiffusionConfig,
         skip_load_model: bool = False,
+        distributed_init_method: str | None = None,
     ):
+        self.distributed_init_method = distributed_init_method
         self.local_rank = local_rank
         self.rank = rank
         self.od_config = od_config
@@ -320,15 +322,16 @@ class DiffusionWorker:
         world_size = self.od_config.num_gpus
         rank = self.rank
 
-        # Set environment variables for distributed initialization
-        os.environ["MASTER_ADDR"] = "localhost"
-        os.environ["MASTER_PORT"] = str(self.od_config.master_port)
+        # Set environment variables for local distributed initialization.
+        if self.distributed_init_method is None:
+            os.environ["MASTER_ADDR"] = "localhost"
+            os.environ["MASTER_PORT"] = str(self.od_config.master_port)
         os.environ["LOCAL_RANK"] = str(self.local_rank)
         os.environ["RANK"] = str(rank)
         os.environ["WORLD_SIZE"] = str(world_size)
 
         # Setup device
-        self.device = current_omni_platform.get_torch_device(rank)
+        self.device = current_omni_platform.get_torch_device(self.local_rank)
         current_omni_platform.set_device(self.device)
 
         # Create vllm_config for parallel configuration. Pass explicit device_config
@@ -354,7 +357,12 @@ class DiffusionWorker:
             set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config),
             set_current_vllm_config(self.vllm_config),
         ):
-            init_distributed_environment(world_size=world_size, rank=rank)
+            init_distributed_environment(
+                world_size=world_size,
+                rank=rank,
+                distributed_init_method=self.distributed_init_method or "env://",
+                local_rank=self.local_rank,
+            )
             logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
 
             parallel_config = self.od_config.parallel_config
@@ -1682,16 +1690,22 @@ class WorkerWrapperBase:
         wake_event: mp.Event = None,
         worker_extension_cls: str | None = None,
         custom_pipeline_args: dict[str, Any] | None = None,
+        rank: int | None = None,
+        distributed_init_method: str | None = None,
     ):
         """
         Initialize WorkerWrapperBase with support for worker extensions.
 
         Args:
-            gpu_id: GPU device ID
+            gpu_id: Local GPU device ID
             od_config: OmniDiffusionConfig configuration
             worker_extension_cls: Optional qualified name of worker extension class
             custom_pipeline_args: Optional arguments passed to native pipelines.
                 A ``pipeline_class`` entry triggers custom pipeline initialization.
+            rank: Global distributed rank. Defaults to ``gpu_id`` for local
+                multiprocessing.
+            distributed_init_method: Explicit rendezvous URL, or None for the
+                local launcher's environment-based rendezvous.
         """
         self.gpu_id = gpu_id
         self.od_config = od_config
@@ -1706,12 +1720,15 @@ class WorkerWrapperBase:
         # Create the actual worker instance. Only dynamic custom pipelines skip
         # initial loading; native pipelines may also use custom_pipeline_args
         # for model-specific component paths.
-        worker = worker_class(
-            local_rank=gpu_id,
-            rank=gpu_id,
-            od_config=od_config,
-            skip_load_model=self.uses_custom_pipeline,
-        )
+        worker_init_kwargs: dict[str, Any] = {
+            "local_rank": gpu_id,
+            "rank": gpu_id if rank is None else rank,
+            "od_config": od_config,
+            "skip_load_model": self.uses_custom_pipeline,
+        }
+        if distributed_init_method is not None:
+            worker_init_kwargs["distributed_init_method"] = distributed_init_method
+        worker = worker_class(**worker_init_kwargs)
         try:
             # Re-initialize pipeline with custom pipeline if provided.
             if self.uses_custom_pipeline:

vllm_omni/engine/stage_engine_startup.py

@@ -1620,11 +1620,17 @@ def launch_diffusion_stage_replica(
         world_size = max(1, int(world_size))
     except (TypeError, ValueError):
         world_size = 1
-    lock_fds = acquire_device_locks(
-        metadata.stage_id,
-        {"tensor_parallel_size": world_size},
-        stage_init_timeout,
-        visible_devices=stage_visible_devices,
+    # These locks coordinate local stage initialization; the orchestrator cannot
+    # acquire them on other hosts where Ray workers may run.
+    lock_fds = (
+        []
+        if od_config.distributed_executor_backend == "ray"
+        else acquire_device_locks(
+            metadata.stage_id,
+            {"tensor_parallel_size": world_size},
+            stage_init_timeout,
+            visible_devices=stage_visible_devices,
+        )
     )
     proc_manager = None
     try:

vllm_omni/engine/stage_init_utils.py

@@ -1994,6 +1994,10 @@ def build_diffusion_config(
         if isinstance(value, int) and value > 0:
             od_config.additional_config.setdefault(f"diffusion_kv_profile_{dimension}", value)
 
+    if od_config.distributed_executor_backend == "ray":
+        runtime_env = _to_dict(_get_attr_or_item(metadata.runtime_cfg, "env", {}) or {})
+        od_config.ray_worker_env = {str(key): str(value) for key, value in runtime_env.items()}
+
     num_devices_per_stage = od_config.parallel_config.world_size
     device_control_env = current_omni_platform.device_control_env_var
     visible_devices_str = os.environ.get(device_control_env) if device_control_env else None
@@ -2003,7 +2007,10 @@ def build_diffusion_config(
     else:
         physical_devices = list(range(current_omni_platform.get_device_count()))
 
-    if len(physical_devices) < num_devices_per_stage:
+    # Ray validates cluster-wide GPU availability through its placement
+    # group. The stage driver only sees the GPUs on its own node, so a local
+    # device-count check would reject every valid multi-node configuration.
+    if od_config.distributed_executor_backend != "ray" and len(physical_devices) < num_devices_per_stage:
         raise ValueError(
             f"Stage {metadata.stage_id} requires {num_devices_per_stage} device(s) based on parallel_config, "
             f"but {len(physical_devices)} device(s) are available: {physical_devices}"