# 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()