diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md
index 214eb3187..57e4721b0 100644
--- a/docs/configuration/stage_configs.md
+++ b/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
diff --git a/docs/design/module/diffusion/diffusion_runtime.md b/docs/design/module/diffusion/diffusion_runtime.md
index 0d36d08c2..52fc97a34 100644
--- a/docs/design/module/diffusion/diffusion_runtime.md
+++ b/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
stage transport and lifecycle"]
engine["DiffusionEngine
admission and control loop"]
scheduler["RequestScheduler or StepScheduler
request state and policy"]
- executor["DiffusionExecutor
uni or mp backend"]
+ executor["DiffusionExecutor
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` |
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index e99da6ceb..7027b5596 100644
--- a/tests/config/test_omni_config.py
+++ b/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",
diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py
index 3f138e721..ce00dde42 100644
--- a/tests/diffusion/test_diffusion_engine_cleanup.py
+++ b/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())
diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py
index 1c13edccd..2c268f2e4 100644
--- a/tests/diffusion/test_multiproc_engine_concurrency.py
+++ b/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()
diff --git a/tests/diffusion/test_ray_executor.py b/tests/diffusion/test_ray_executor.py
new file mode 100644
index 000000000..9749cb755
--- /dev/null
+++ b/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
diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py
index c6b5761d9..3761384ac 100644
--- a/tests/diffusion/test_stage_diffusion_proc.py
+++ b/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()
diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py
index aeedd25b7..4c5405178 100644
--- a/tests/diffusion/test_worker_wrapper_base.py
+++ b/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)
diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py
index 46ac15f9b..a5314ea26 100644
--- a/tests/engine/test_async_omni_engine_stage_init.py
+++ b/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)
diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py
index 0d5e47960..95b7c1d6e 100644
--- a/tests/engine/test_single_stage_mode.py
+++ b/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",
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index 19d5ae71c..31f202aee 100644
--- a/vllm_omni/config/omni_config.py
+++ b/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)
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 0811d2634..be4b842b6 100644
--- a/vllm_omni/diffusion/data.py
+++ b/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
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 57dddc889..037ec9720 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/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():
diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py
index 4bbef2f27..de1e90db8 100644
--- a/vllm_omni/diffusion/executor/abstract.py
+++ b/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
diff --git a/vllm_omni/diffusion/executor/ray_executor.py b/vllm_omni/diffusion/executor/ray_executor.py
new file mode 100644
index 000000000..692628bca
--- /dev/null
+++ b/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()
diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py
index c8fcf683d..e8bffdf39 100644
--- a/vllm_omni/diffusion/stage_diffusion_client.py
+++ b/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)
diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py
index fc795b6b6..3f574be2c 100644
--- a/vllm_omni/diffusion/stage_diffusion_proc.py
+++ b/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]
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index fe344e4d0..76ca42a7b 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/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:
diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py
index 81af45d58..4823a192d 100644
--- a/vllm_omni/engine/stage_engine_startup.py
+++ b/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:
diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py
index 27d37bd59..da09b5235 100644
--- a/vllm_omni/engine/stage_init_utils.py
+++ b/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}"