# SPDX-License-Identifier: Apache-2.0 # Adapt from https://github.com/vllm-project/vllm/blob/releases/v0.11.0/vllm/executor/ray_distributed_executor.py import asyncio from collections import defaultdict from queue import Queue import os import cloudpickle import fastvideo.envs as envs from dataclasses import dataclass from typing import Any, TYPE_CHECKING from collections.abc import Callable from fastvideo.utils import (DEPRECATED_HF_TOKEN_ENV_VARS, HF_TOKEN_ENV_VARS, get_ip, get_distributed_init_method, get_open_port, get_loopback_ip) from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.worker.executor import Executor from fastvideo.worker.ray_utils import ( initialize_ray_cluster, RayWorkerWrapper, ray, ) from fastvideo.worker.ray_env import get_env_vars_to_copy from fastvideo.logger import init_logger if ray is not None: from ray.actor import ActorHandle from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy else: ActorHandle = None if TYPE_CHECKING: from ray.util.placement_group import PlacementGroup logger = init_logger(__name__) def should_use_gloo_loopback(worker_ips: list[str]) -> bool: """Loopback is only safe when every worker shares one host IP. Single-node Ray (one or many GPUs on the same box) can dial the Gloo store on loopback. Two Sparks already have distinct worker IPs, so this returns False and Gloo stays on the fabric address. Per-node NIC names are a separate issue: do not copy ``NCCL_SOCKET_IFNAME`` / ``GLOO_SOCKET_IFNAME`` from the driver onto those workers. """ return len(set(worker_ips)) <= 1 # Ray narrows each GPU actor's CUDA_VISIBLE_DEVICES to the actor's own GPU unless this is set. RAY_NOSET_CUDA_VISIBLE_DEVICES = "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES" def keep_raylet_cuda_devices(ray_remote_kwargs: dict[str, Any]) -> dict[str, Any]: """Return actor options under which Ray leaves CUDA_VISIBLE_DEVICES as the raylet's. Ray merges an actor's ``env_vars`` into the job's runtime environment, so only this one variable is added. """ runtime_env = dict(ray_remote_kwargs.get("runtime_env") or {}) runtime_env["env_vars"] = {**(runtime_env.get("env_vars") or {}), RAY_NOSET_CUDA_VISIBLE_DEVICES: "1"} return {**ray_remote_kwargs, "runtime_env": runtime_env} def ray_worker_device_ordinal(node_gpu_ids: list[int], node_local_index: int, visible_devices: str | None) -> int: """Return the CUDA ordinal of the ``node_local_index``-th worker on a node. The worker uses ``node_gpu_ids[node_local_index]`` (the sorted Ray GPU IDs of the node's workers), addressed in the device list it inherited from its raylet. Ray reports GPU IDs as entries of the raylet's CUDA_VISIBLE_DEVICES, or as ordinals when that variable is unset. """ gpu_id = node_gpu_ids[node_local_index] if not visible_devices: return gpu_id visible = [entry.strip() for entry in visible_devices.split(",")] if str(gpu_id) not in visible: raise RuntimeError(f"Ray assigned GPU {gpu_id}, which is not in the worker's " f"CUDA_VISIBLE_DEVICES={visible_devices!r}.") return visible.index(str(gpu_id)) @dataclass class RayWorkerMetaData: """ Metadata for a Ray worker. The order of ray worker creation can be random, and we need to reset the rank after creating all workers. """ worker: ActorHandle created_rank: int adjusted_rank: int = -1 ip: str = "" class RayDistributedExecutor(Executor): """Ray-based distributed executor""" # These env vars are worker-specific, therefore are NOT copied # from the driver to the workers WORKER_SPECIFIC_ENV_VARS = { "FASTVIDEO_HOST_IP", "LOCAL_RANK", "CUDA_VISIBLE_DEVICES", } # Per-node fabric names. spark_pair_env.sh / ibdev2netdev can differ across # boxes; pushing the driver's value overwrites the export set before ray start. WORKER_LOCAL_NIC_ENV_VARS = { "NCCL_SOCKET_IFNAME", "NCCL_IB_HCA", "GLOO_SOCKET_IFNAME", } # These non-vLLM env vars are copied from the driver to workers. # NCCL_* knobs present on the driver are added dynamically in # ``_env_vars_to_copy_from_driver``, except the per-node NIC trio above. ADDITIONAL_ENV_VARS = { *HF_TOKEN_ENV_VARS, *DEPRECATED_HF_TOKEN_ENV_VARS, "NCCL_IB_DISABLE", "NCCL_P2P_DISABLE", "NCCL_CUMEM_ENABLE", "NCCL_NVLS_ENABLE", "NCCL_DEBUG", "NCCL_DEBUG_SUBSYS", "LD_LIBRARY_PATH", # MiniMax-H3 encoder split: the driver resolves these into # fastvideo_args, but keep the environment aligned on the workers so # code paths that consult envs directly agree with the driver. "FASTVIDEO_H3_ENCODER_SPLIT", "FASTVIDEO_H3_ENCODER_NODES", } def _init_executor(self) -> None: initialize_ray_cluster(self.fastvideo_args) placement_group = self.fastvideo_args.ray_placement_group # Disable Ray usage stats collection. ray_usage = os.environ.get("RAY_USAGE_STATS_ENABLED", "0") if ray_usage != "1": envs.set_external("RAY_USAGE_STATS_ENABLED", "0") self._init_workers_ray(placement_group) # child class could overwrite this to return actual env vars. def _get_env_vars_to_be_updated(self) -> list[dict[str, str]]: return self._env_vars_for_all_workers def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwargs): from fastvideo.platforms import current_platform num_gpus = envs.FASTVIDEO_RAY_PER_WORKER_GPUS.get() # On NVIDIA GPUs every actor keeps its raylet's CUDA_VISIBLE_DEVICES. # Otherwise Ray narrows it to the actor's own GPU before the actor # imports fastvideo, and that import initializes CUDA (diffusers and # Triton query the device at import time). CUDA then keeps the one-GPU # view, a later CUDA_VISIBLE_DEVICES update has no effect, and # cuda: is an invalid ordinal for every worker but the first # on a multi-GPU node. Each worker addresses its GPU by position in the # inherited list instead (ray_worker_device_ordinal). keep_raylet_devices = current_platform.is_cuda() if keep_raylet_devices: ray_remote_kwargs = keep_raylet_cuda_devices(ray_remote_kwargs) # The remaining workers are the actual ray actors. self.workers: list[RayWorkerWrapper] = [] # Create the workers. # use the first N bundles that have GPU resources. bundle_indices: list[int] = [] for bundle_id, bundle in enumerate(placement_group.bundle_specs): if bundle.get(current_platform.ray_device_key, 0): bundle_indices.append(bundle_id) bundle_indices = bundle_indices[:self.fastvideo_args.num_gpus] worker_metadata: list[RayWorkerMetaData] = [] driver_ip = get_ip() for rank, bundle_id in enumerate(bundle_indices): scheduling_strategy = PlacementGroupSchedulingStrategy( placement_group=placement_group, placement_group_capture_child_tasks=True, placement_group_bundle_index=bundle_id, ) if current_platform.ray_device_key == "GPU": # NV+AMD GPUs, and Intel XPUs worker = ray.remote( num_cpus=0, num_gpus=num_gpus, scheduling_strategy=scheduling_strategy, **ray_remote_kwargs, )(RayWorkerWrapper).remote(fastvideo_args=self.fastvideo_args, rpc_rank=rank) else: worker = ray.remote( num_cpus=0, num_gpus=0, resources={current_platform.ray_device_key: num_gpus}, scheduling_strategy=scheduling_strategy, **ray_remote_kwargs, )(RayWorkerWrapper).remote(fastvideo_args=self.fastvideo_args, rpc_rank=rank) worker_metadata.append(RayWorkerMetaData(worker=worker, created_rank=rank)) worker_ips = ray.get([ each.worker.get_node_ip.remote() # type: ignore[attr-defined] for each in worker_metadata ]) for each, ip in zip(worker_metadata, worker_ips, strict=False): each.ip = ip logger.debug("workers: %s", worker_metadata) ip_counts: dict[str, int] = {} for ip in worker_ips: ip_counts[ip] = ip_counts.get(ip, 0) + 1 def sort_by_driver_then_worker_ip(item: RayWorkerMetaData): """ Sort the workers based on 3 properties: 1. If the worker is on the same node as the driver (vllm engine), it should be placed first. 2. Then, if the worker is on a node with fewer workers, it should be placed first. 3. Finally, if the work is on a node with smaller IP address, it should be placed first. """ ip = item.ip return (0 if ip == driver_ip else 1, ip_counts[ip], ip) # After sorting, the workers on the same node will be # close to each other, and the workers on the driver # node will be placed first. sorted_worker_metadata = sorted(worker_metadata, key=sort_by_driver_then_worker_ip) start_rank = 0 for i, item in enumerate(sorted_worker_metadata): item.adjusted_rank = i + start_rank self.workers = [item.worker for item in sorted_worker_metadata] rerank_mapping = {item.created_rank: item.adjusted_rank for item in sorted_worker_metadata} self._run_ray_workers("adjust_rank", rerank_mapping) if getattr(self.fastvideo_args, "h3_encoder_split", False): # The sort keeps every node's workers contiguous, so the first # ``h3_encoder_nodes`` distinct IPs are exactly the encoder group. # Stamped before init_worker pickles the args to every worker. self._setup_h3_encoder_split([item.ip for item in sorted_worker_metadata]) # Get the set of GPU IDs used on each node. worker_node_and_gpu_ids = self._run_ray_workers("get_node_and_gpu_ids") node_workers = defaultdict(list) # node id -> list of worker ranks node_gpus = defaultdict(list) # node id -> list of gpu ids for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids): node_workers[node_id].append(i) # `gpu_ids` can be a list of strings or integers. # convert them to integers for consistency. # NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs), # string sorting is not sufficient. # see https://github.com/vllm-project/vllm/issues/5590 gpu_ids = [int(x) for x in gpu_ids] node_gpus[node_id].extend(gpu_ids) for node_id, gpu_ids in node_gpus.items(): node_gpus[node_id] = sorted(gpu_ids) all_ips = set(worker_ips + [driver_ip]) n_ips = len(all_ips) n_nodes = len(node_workers) if n_nodes != n_ips: raise RuntimeError(f"Every node should have a unique IP address. Got {n_nodes}" f" nodes with node ids {list(node_workers.keys())} and " f"{n_ips} unique IP addresses {all_ips}. Please check your" " network configuration. If you set `FASTVIDEO_HOST_IP`" " environment variable, make sure it is unique for" " each node.") # Set environment variables for the driver and workers. all_args_to_update_environment_variables: list[dict[str, str]] if keep_raylet_devices: # The inherited device list stays; each worker gets its ordinal in it. worker_visible_devices = self._run_ray_workers("get_cuda_visible_devices") all_args_to_update_environment_variables = [{} for _ in worker_node_and_gpu_ids] else: all_args_to_update_environment_variables = [{ current_platform.device_control_env_var: ",".join(map(str, node_gpus[node_id])), } for (node_id, _) in worker_node_and_gpu_ids] # Environment variables to copy from driver to workers extra_nccl = {k for k in os.environ if k.startswith("NCCL_") and k not in self.WORKER_LOCAL_NIC_ENV_VARS} env_vars_to_copy = get_env_vars_to_copy( exclude_vars=self.WORKER_SPECIFIC_ENV_VARS | self.WORKER_LOCAL_NIC_ENV_VARS, additional_vars=set(current_platform.additional_env_vars).union(self.ADDITIONAL_ENV_VARS).union(extra_nccl), destination="workers", ) # Copy existing env vars to each worker's args for args in all_args_to_update_environment_variables: # TODO: refactor platform-specific env vars for name in env_vars_to_copy: if name in os.environ: args[name] = os.environ[name] self._env_vars_for_all_workers: list[dict[str, str]] = (all_args_to_update_environment_variables) self._run_ray_workers("update_environment_variables", self._get_env_vars_to_be_updated()) _lb = should_use_gloo_loopback(worker_ips) if _lb: driver_ip = get_loopback_ip() distributed_init_method = get_distributed_init_method(driver_ip, get_open_port()) # Initialize the actual workers inside worker wrapper. all_kwargs = [] for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids): local_rank = node_workers[node_id].index(rank) if keep_raylet_devices: local_rank = ray_worker_device_ordinal(node_gpus[node_id], local_rank, worker_visible_devices[rank]) kwargs = dict( fastvideo_args=self.fastvideo_args, local_rank=local_rank, rank=rank, distributed_init_method=distributed_init_method, ) all_kwargs.append(kwargs) self._run_ray_workers("init_worker", all_kwargs) self._run_ray_workers("init_device") # This is the list of workers that are rank 0 of each TP group EXCEPT # global rank 0. These are the workers that will broadcast to the # rest of the workers. self.tp_driver_workers: list[RayWorkerWrapper] = [] # This is the list of workers that are not drivers and not the first # worker in a TP group. These are the workers that will be # broadcasted to. self.non_driver_workers: list[RayWorkerWrapper] = [] # Enforce rank order for correct rank to return final output. for index, worker in enumerate(self.workers): # The driver worker is rank 0 and not in self.workers. rank = index + 1 if rank % self.fastvideo_args.tp_size == 0: self.tp_driver_workers.append(worker) else: self.non_driver_workers.append(worker) def _setup_h3_encoder_split(self, sorted_worker_ips: list[str]) -> None: """Reserve the leading encoder nodes' workers and stamp the count onto the args. ``sorted_worker_ips`` follows the driver-then-node sort, so workers of one node are contiguous; the first ``h3_encoder_nodes`` distinct IPs form the encoder group (ranks ``0..n-1``). Workers must be one-per-node style placements; multi-GPU nodes simply contribute all of their workers. """ args = self.fastvideo_args # ``_fold_h3_encoder_split_env`` already normalized 0/unset to 1. nodes_needed = int(args.h3_encoder_nodes) ordered_ips: list[str] = [] for ip in sorted_worker_ips: if ip not in ordered_ips: ordered_ips.append(ip) if len(ordered_ips) < nodes_needed + 1: raise RuntimeError(f"MiniMax-H3 encoder split needs at least {nodes_needed + 1} distinct nodes " f"({nodes_needed} encoder + 1 denoise), got {len(ordered_ips)}.") encoder_ips = set(ordered_ips[:nodes_needed]) encoder_workers = sum(1 for ip in sorted_worker_ips if ip in encoder_ips) if not 1 <= encoder_workers < len(sorted_worker_ips): raise RuntimeError(f"MiniMax-H3 encoder split resolved {encoder_workers} encoder workers out of " f"{len(sorted_worker_ips)}; need at least one denoise worker.") from fastvideo.fastvideo_args import h3_split_sp_error, probe_h3_attention_heads denoise_workers = len(sorted_worker_ips) - encoder_workers heads = probe_h3_attention_heads(args.model_path) if heads and heads % denoise_workers: raise RuntimeError(h3_split_sp_error(heads, denoise_workers, len(sorted_worker_ips), unit="workers")) args.h3_encoder_nodes = nodes_needed args.h3_encoder_workers = encoder_workers logger.info( "MiniMax-H3 encoder split: ranks 0..%d on nodes %s run the Qwen3-VL conditioner only; ranks %d..%d " "run DiT/VAE with sp_size=%d", encoder_workers - 1, sorted(encoder_ips), encoder_workers, len(sorted_worker_ips) - 1, len(sorted_worker_ips) - encoder_workers) def execute_streaming_reset(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> dict[str, Any]: responses: list[dict[str, Any]] = self.collective_rpc( "execute_streaming_reset", kwargs={ "forward_batch": forward_batch, "fastvideo_args": fastvideo_args, }, ) return responses[0] def execute_streaming_step(self, keyboard_action=None, mouse_action=None) -> ForwardBatch: responses: list[ForwardBatch] = self.collective_rpc( "execute_streaming_step", kwargs={ "keyboard_action": keyboard_action, "mouse_action": mouse_action, }, ) return responses[0] async def execute_streaming_step_async(self, keyboard_action=None, mouse_action=None) -> ForwardBatch: kwargs = { "keyboard_action": keyboard_action, "mouse_action": mouse_action, } futures = [w.execute_method.remote("execute_streaming_step", **kwargs) for w in self.workers] responses = await asyncio.gather(*futures) return responses[0] def execute_streaming_clear(self) -> None: self.collective_rpc("execute_streaming_clear") def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: responses: list[ForwardBatch] = self.collective_rpc( "execute_forward", kwargs={ "forward_batch": forward_batch, "fastvideo_args": fastvideo_args, }, ) # Under the MiniMax-H3 encoder split the encoder ranks answer with a # stub batch; the decoded output lives on the first denoise rank. src = 0 # The executor's own args carry h3_encoder_workers (stamped when the # workers were placed); fall back to the per-call args when absent. split_args = getattr(self, "fastvideo_args", fastvideo_args) if getattr(split_args, "h3_encoder_split", False): src = int(split_args.h3_encoder_workers) output = responses[src].output.cpu() logging_info = None if envs.FASTVIDEO_STAGE_LOGGING.get(): logging_info = responses[src].logging_info result_batch = ForwardBatch( data_type=forward_batch.data_type, output=output, logging_info=logging_info, extra=responses[src].extra, trajectory_latents=responses[src].trajectory_latents, trajectory_timesteps=responses[src].trajectory_timesteps, ) return result_batch def set_lora_adapter(self, lora_nickname: str, lora_path: str | None = None, strength: float = 1.0, accumulate: bool = False) -> None: responses = self.collective_rpc("set_lora_adapter", kwargs={ "lora_nickname": lora_nickname, "lora_path": lora_path, "strength": strength, "accumulate": accumulate }) for i, response in enumerate(responses): if response["status"] != "lora_adapter_set": raise RuntimeError(f"Worker {i} failed to set LoRA adapter to {lora_path}") def unmerge_lora_weights(self) -> None: responses = self.collective_rpc("unmerge_lora_weights", kwargs={}) for i, response in enumerate(responses): if response["status"] != "lora_adapter_unmerged": raise RuntimeError(f"Worker {i} failed to unmerge LoRA weights") def merge_lora_weights(self) -> None: responses = self.collective_rpc("merge_lora_weights", kwargs={}) for i, response in enumerate(responses): if response["status"] != "lora_adapter_merged": raise RuntimeError(f"Worker {i} failed to merge LoRA weights") def set_log_queue(self, log_queue: Queue | None) -> None: """Keep the driver-side queue locally. ``multiprocessing.Queue`` is not picklable across Ray nodes, so worker logs stay in the Ray session log dir instead of being forwarded. """ self._log_queue = log_queue def clear_log_queue(self) -> None: self._log_queue = None def collective_rpc(self, method: str | Callable, timeout: float | None = None, args: tuple = (), kwargs: dict | None = None) -> list[Any]: if timeout is not None: raise NotImplementedError("RayDistributedExecutor.collective_rpc does not support per-call timeouts; " "omit the timeout and rely on Ray's own task failure handling instead.") return self._run_ray_workers(method, *args, **(kwargs or {})) def _run_ray_workers( self, method: str | Callable, *args, **kwargs, ) -> Any: sent_method = method if isinstance(method, str) else cloudpickle.dumps(method) del method # Start the ray workers first. ray_workers = self.workers ray_worker_outputs = [worker.execute_method.remote(sent_method, *args, **kwargs) for worker in ray_workers] # Get the results of the ray workers. ray_worker_outputs = ray.get(ray_worker_outputs) return ray_worker_outputs def shutdown(self) -> None: logger.info("Shutting down Ray distributed executor. If you see error log " "from logging.cc regarding SIGTERM received, please ignore because " "this is the expected termination process in Ray.") import ray for worker in self.workers: ray.kill(worker) self.workers = [] def __del__(self): self.shutdown()