# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing as mp import os import time import torch import zmq from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.data import ( DiffusionOutput, OmniDiffusionConfig, set_current_omni_diffusion_config, ) from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, init_distributed_environment, initialize_model_parallel, ) from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.request import OmniDiffusionRequest logger = init_logger(__name__) class GPUWorker: """ A worker that executes the model on a single GPU. """ def __init__( self, local_rank: int, rank: int, od_config: OmniDiffusionConfig, ): self.local_rank = local_rank self.rank = rank self.od_config = od_config self.pipeline = None self.device = None self.init_device_and_model() def init_device_and_model(self) -> None: """Initialize the device and load the model.""" 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) os.environ["LOCAL_RANK"] = str(self.local_rank) os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) self.device = torch.device(f"cuda:{rank}") torch.cuda.set_device(self.device) # hack vllm_config = VllmConfig() vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size self.vllm_config = vllm_config with ( set_current_omni_diffusion_config(self.od_config), set_current_vllm_config(vllm_config), ): init_distributed_environment(world_size=world_size, rank=rank) logger.info(f"Worker {self.rank}: Initialized device and distributed environment.") parallel_config = self.od_config.parallel_config initialize_model_parallel( data_parallel_size=parallel_config.data_parallel_size, cfg_parallel_size=parallel_config.cfg_parallel_size, sequence_parallel_size=parallel_config.sequence_parallel_size, ulysses_degree=parallel_config.ulysses_degree, ring_degree=parallel_config.ring_degree, tensor_parallel_size=parallel_config.tensor_parallel_size, pipeline_parallel_size=parallel_config.pipeline_parallel_size, ) load_config = LoadConfig() model_loader = DiffusersPipelineLoader(load_config) time_before_load = time.perf_counter() with DeviceMemoryProfiler() as m: self.pipeline = model_loader.load_model( od_config=self.od_config, load_device=str(self.device), ) time_after_load = time.perf_counter() logger.info( "Model loading took %.4f GiB and %.6f seconds", m.consumed_memory / GiB_bytes, time_after_load - time_before_load, ) logger.info(f"Worker {self.rank}: Model loaded successfully.") # Setup cache backend based on type (both backends use enable()/reset() interface) self.cache_backend = get_cache_backend(self.od_config.cache_backend, self.od_config.cache_config) if self.cache_backend is not None: self.cache_backend.enable(self.pipeline) def generate(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput: """ Generate output for the given requests. Args: requests: List of diffusion requests Returns: DiffusionOutput with generated results """ return self.execute_model(requests, self.od_config) @torch.inference_mode() def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput: """ Execute a forward pass. """ assert self.pipeline is not None if not reqs or len(reqs) == 0: raise ValueError("Cannot execute model with empty request list") # TODO: dealing with first req for now req = reqs[0] if req.generator is None and req.seed is not None: req.generator = torch.Generator(device=self.device).manual_seed(req.seed) # Refresh cache context if needed if self.cache_backend is not None and self.cache_backend.is_enabled(): self.cache_backend.refresh(self.pipeline, req.num_inference_steps) with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config): output = self.pipeline.forward(req) return output def shutdown(self) -> None: destroy_distributed_env() class WorkerProc: """Wrapper that runs one Worker in a separate process.""" def __init__( self, od_config: OmniDiffusionConfig, gpu_id: int, broadcast_handle, ): self.od_config = od_config # Inter-process Communication self.context = zmq.Context(io_threads=2) # Initialize MessageQueue reader from handle (unified for generation & RPC) self.mq = MessageQueue.create_from_handle(broadcast_handle, gpu_id) self.result_mq = None self.result_mq_handle = None # Setup result sender (only for rank 0 for now, or whoever needs to reply) # Assuming only rank 0 replies to scheduler as per original logic if gpu_id == 0: # Create MessageQueue for results (1 writer -> 1 reader) # We assume the reader (SyncScheduler) will act as rank 0 self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0]) self.result_mq_handle = self.result_mq.export_handle() logger.info(f"Worker {gpu_id} created result MessageQueue") assert od_config.master_port is not None self.worker = self._create_worker(gpu_id, od_config) self.gpu_id = gpu_id self._running = True def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> GPUWorker: """Create a worker instance. Override in subclasses for different worker types.""" return GPUWorker( local_rank=gpu_id, rank=gpu_id, od_config=od_config, ) def return_result(self, output: DiffusionOutput): """ replies to client, only on rank 0 """ if self.result_mq is not None: self.result_mq.enqueue(output) def recv_message(self): """ Receive unified messages (RPC requests, shutdown) from broadcast queue. Uses indefinite=True to block until a message arrives. """ return self.mq.dequeue(indefinite=True) def execute_rpc(self, rpc_request: dict) -> tuple[object | None, bool]: """Execute an RPC request and indicate whether to reply.""" method = rpc_request["method"] args = rpc_request.get("args", ()) kwargs = rpc_request.get("kwargs", {}) output_rank = rpc_request.get("output_rank") exec_all_ranks = rpc_request.get("exec_all_ranks", False) should_execute = exec_all_ranks or output_rank is None or output_rank == self.gpu_id should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None if not should_execute: return None, False try: if isinstance(method, str): func = getattr(self.worker, method) result = func(*args, **kwargs) else: result = method(self.worker, *args, **kwargs) return result, should_reply except Exception as e: logger.error(f"Error executing RPC: {e}", exc_info=True) return {"status": "error", "error": str(e)}, should_reply # TODO: queueing, cancellation def worker_busy_loop(self) -> None: """Main busy loop for Multiprocessing Workers""" logger.info(f"Worker {self.gpu_id} ready to receive requests via shared memory") while self._running: # Receive unified message (generation request, RPC request, or shutdown) msg = None try: msg = self.recv_message() except Exception as e: logger.error( f"Error receiving message in worker loop: {e}", exc_info=True, ) continue if msg is None or len(msg) == 0: logger.warning("Worker %s: Received empty payload, ignoring", self.gpu_id) continue # Route message based on type if isinstance(msg, dict) and msg.get("type") == "rpc": # Handle RPC request try: result, should_reply = self.execute_rpc(msg) if should_reply: self.return_result(result) except Exception as e: logger.error(f"Error processing RPC: {e}", exc_info=True) if self.result_mq is not None: self.return_result({"status": "error", "error": str(e)}) elif isinstance(msg, dict) and msg.get("type") == "shutdown": # Handle shutdown message logger.info("Worker %s: Received shutdown message", self.gpu_id) self._running = False continue else: # Handle generation request (OmniDiffusionRequest list) try: output = self.worker.execute_model(msg, self.od_config) except Exception as e: logger.error( f"Error executing forward in event loop: {e}", exc_info=True, ) output = DiffusionOutput(error=str(e)) try: self.return_result(output) except zmq.ZMQError as e: # Reply failed; log and keep loop alive to accept future requests logger.error(f"ZMQ error sending reply: {e}") continue logger.info("event loop terminated.") try: self.worker.shutdown() except Exception as exc: # pragma: no cover - best effort cleanup logger.warning("Worker %s: Shutdown encountered an error: %s", self.gpu_id, exc) # if self.result_sender is not None: # self.result_sender.close() self.context.term() @staticmethod def worker_main( rank: int, od_config: OmniDiffusionConfig, pipe_writer: mp.connection.Connection, broadcast_handle, ) -> None: """Worker initialization and execution loops.""" worker_proc = WorkerProc( od_config, gpu_id=rank, broadcast_handle=broadcast_handle, ) logger.info(f"Worker {rank}: Scheduler loop started.") pipe_writer.send( { "status": "ready", "result_handle": worker_proc.result_mq_handle if rank == 0 else None, } ) worker_proc.worker_busy_loop() logger.info(f"Worker {rank}: Shutdown complete.")