# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing as mp import os import time from collections.abc import Iterable from contextlib import AbstractContextManager, nullcontext import torch import zmq from vllm.config import LoadConfig, VllmConfig 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.compile import regionally_compile from vllm_omni.diffusion.data import ( DiffusionOutput, OmniDiffusionConfig, ) 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._sleep_saved_buffers: dict[str, torch.Tensor] = {} 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_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_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 self._maybe_get_memory_pool_context(tag="weights"): 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.") if not self.od_config.enforce_eager: try: self.pipeline.transformer = regionally_compile( self.pipeline.transformer, dynamic=True, ) logger.info(f"Worker {self.rank}: Model compiled with torch.compile.") except Exception as e: logger.warning(f"Worker {self.rank}: torch.compile failed with error: {e}. Using eager mode.") # 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 load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return self.pipeline.load_weights(weights) def sleep(self, level: int = 1) -> bool: """ Put the worker to sleep. The worker should not process any requests. The caller should guarantee that no requests are being processed during the sleep period, before `wake_up` is called. Args: level: The sleep level. Level 1 sleep will offload the model weights and discard the kv cache. Currently only support level 1. """ from vllm.device_allocator.cumem import CuMemAllocator free_bytes_before_sleep = torch.cuda.mem_get_info()[0] # Save the buffers before level 2 sleep if level == 2: model = self.pipeline self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()} allocator = CuMemAllocator.get_instance() allocator.sleep(offload_tags=("weights",) if level == 1 else tuple()) free_bytes_after_sleep, total = torch.cuda.mem_get_info() freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep used_bytes = total - free_bytes_after_sleep assert freed_bytes >= 0, "Memory usage increased after sleeping." logger.info( "Sleep mode freed %.2f GiB memory, %.2f GiB memory is still in use.", freed_bytes / GiB_bytes, used_bytes / GiB_bytes, ) return True def wake_up(self, tags: list[str] | None = None) -> bool: """ Wake up the worker from sleep mode. See the sleep function method for more details. Args: tags: An optional list of tags to reallocate the worker memory for specific memory allocations. Values must be in `("weights")`. If None, all memory is reallocated. wake_up should be called with all tags (or None) before the worker is used again. """ from vllm.device_allocator.cumem import CuMemAllocator allocator = CuMemAllocator.get_instance() allocator.wake_up(tags) # Restore the buffers after level 2 sleep if len(self._sleep_saved_buffers): model = self.pipeline for name, buffer in model.named_buffers(): if name in self._sleep_saved_buffers: buffer.data.copy_(self._sleep_saved_buffers[name].data) self._sleep_saved_buffers = {} return True def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: if self.od_config.enable_sleep_mode: from vllm.device_allocator.cumem import CuMemAllocator allocator = CuMemAllocator.get_instance() if tag == "weights": assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process." return allocator.use_memory_pool(tag=tag) else: return nullcontext() 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.")