705 lines
27 KiB
Python
705 lines
27 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
Diffusion Worker for vLLM-Omni.
|
|
|
|
Handles GPU infrastructure initialization and delegates model operations
|
|
to DiffusionModelRunner.
|
|
"""
|
|
|
|
import gc
|
|
import multiprocessing as mp
|
|
import os
|
|
from collections.abc import Iterable
|
|
from contextlib import AbstractContextManager, nullcontext
|
|
from typing import Any
|
|
|
|
import torch
|
|
import zmq
|
|
from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
|
|
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.import_utils import resolve_obj_by_qualname
|
|
from vllm.utils.mem_utils import GiB_bytes
|
|
from vllm.v1.worker.workspace import init_workspace_manager
|
|
|
|
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.ipc import pack_diffusion_output_shm
|
|
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
|
|
from vllm_omni.diffusion.profiler import CurrentProfiler
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
|
from vllm_omni.lora.request import LoRARequest
|
|
from vllm_omni.platforms import current_omni_platform
|
|
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class DiffusionWorker:
|
|
"""
|
|
A worker that manages GPU infrastructure and delegates to the model runner.
|
|
|
|
This class handles infrastructure initialization only:
|
|
- Device setup (CUDA device selection)
|
|
- Distributed environment (NCCL, model parallel)
|
|
- Memory management (sleep/wake)
|
|
|
|
All model-related operations (loading, compilation, execution) are
|
|
delegated to DiffusionModelRunner.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
local_rank: int,
|
|
rank: int,
|
|
od_config: OmniDiffusionConfig,
|
|
skip_load_model: bool = False,
|
|
):
|
|
self.local_rank = local_rank
|
|
self.rank = rank
|
|
self.od_config = od_config
|
|
self.device: torch.device | None = None
|
|
self.vllm_config: VllmConfig | None = None
|
|
self.model_runner: DiffusionModelRunner | None = None
|
|
self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
|
|
self.lora_manager: DiffusionLoRAManager | None = None
|
|
self.init_device()
|
|
# Create model runner
|
|
self.model_runner = DiffusionModelRunner(
|
|
vllm_config=self.vllm_config,
|
|
od_config=self.od_config,
|
|
device=self.device,
|
|
)
|
|
if not skip_load_model:
|
|
self.load_model(load_format=self.od_config.diffusion_load_format)
|
|
self.init_lora_manager()
|
|
logger.info(f"Worker {self.rank}: Initialization complete.")
|
|
|
|
def init_device(self) -> None:
|
|
"""Initialize the device and distributed environment."""
|
|
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)
|
|
|
|
# Setup device
|
|
self.device = current_omni_platform.get_torch_device(rank)
|
|
current_omni_platform.set_device(self.device)
|
|
|
|
# Create vllm_config for parallel configuration
|
|
vllm_config = VllmConfig(compilation_config=CompilationConfig())
|
|
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
|
|
vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel
|
|
self.vllm_config = vllm_config
|
|
|
|
# Initialize distributed environment
|
|
with (
|
|
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)
|
|
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,
|
|
fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1,
|
|
hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1,
|
|
enable_expert_parallel=parallel_config.enable_expert_parallel,
|
|
)
|
|
init_workspace_manager(self.device)
|
|
|
|
def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None) -> None:
|
|
"""Load the diffusion model using DiffusionModelRunner."""
|
|
with (
|
|
set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config),
|
|
set_current_vllm_config(self.vllm_config),
|
|
):
|
|
self.model_runner.load_model(
|
|
memory_pool_context_fn=self._maybe_get_memory_pool_context,
|
|
load_format=load_format,
|
|
custom_pipeline_name=custom_pipeline_name,
|
|
)
|
|
process_memory = get_process_gpu_memory(self.local_rank)
|
|
if process_memory is not None:
|
|
logger.info(
|
|
"Worker %d: Process-scoped GPU memory after model loading: %.2f GiB.",
|
|
self.rank,
|
|
process_memory / GiB_bytes,
|
|
)
|
|
|
|
# When load_format is "dummy", pipeline will init with custom pipeline later
|
|
if load_format != "dummy":
|
|
assert self.model_runner.pipeline is not None
|
|
|
|
def init_lora_manager(self) -> None:
|
|
"""Initialize the LoRA manager for this worker."""
|
|
if self.model_runner.pipeline is None:
|
|
return
|
|
self.lora_manager = DiffusionLoRAManager(
|
|
pipeline=self.model_runner.pipeline,
|
|
device=self.device,
|
|
dtype=self.od_config.dtype,
|
|
max_cached_adapters=self.od_config.max_cpu_loras,
|
|
lora_path=self.od_config.lora_path,
|
|
lora_scale=self.od_config.lora_scale,
|
|
)
|
|
|
|
def generate(self, request: OmniDiffusionRequest) -> DiffusionOutput:
|
|
"""Generate output for the given requests."""
|
|
return self.execute_model(request, self.od_config)
|
|
|
|
@classmethod
|
|
def start_profile(cls, trace_path_template: str) -> str:
|
|
"""Start profiling for this GPU worker."""
|
|
return CurrentProfiler.start(trace_path_template)
|
|
|
|
@classmethod
|
|
def stop_profile(cls) -> dict | None:
|
|
"""Stop profiling and return the result dictionary."""
|
|
return CurrentProfiler.stop()
|
|
|
|
def execute_model(self, req: OmniDiffusionRequest, od_config: OmniDiffusionConfig) -> DiffusionOutput:
|
|
"""Execute a forward pass by delegating to the model runner."""
|
|
assert self.model_runner is not None, "Model runner not initialized"
|
|
if self.lora_manager is not None:
|
|
try:
|
|
self.lora_manager.set_active_adapter(req.sampling_params.lora_request, req.sampling_params.lora_scale)
|
|
except Exception as exc:
|
|
if req.sampling_params.lora_request is not None:
|
|
raise
|
|
logger.warning("LoRA activation skipped: %s", exc)
|
|
return self.model_runner.execute_model(req)
|
|
|
|
def load_weights(self, weights) -> set[str]:
|
|
"""Load weights by delegating to the model runner."""
|
|
assert self.model_runner is not None, "Model runner not initialized"
|
|
return self.model_runner.load_weights(weights)
|
|
|
|
def remove_lora(self, adapter_id: int) -> bool:
|
|
return self.lora_manager.remove_adapter(adapter_id)
|
|
|
|
def add_lora(self, lora_request: LoRARequest) -> bool:
|
|
# NOTE (Alex): We have not implemented the API routing
|
|
# for the frontend server yet.
|
|
return self.lora_manager.add_adapter(lora_request)
|
|
|
|
def list_loras(self) -> list[int]:
|
|
return self.lora_manager.list_adapters()
|
|
|
|
def pin_lora(self, adapter_id: int) -> bool:
|
|
return self.lora_manager.pin_adapter(adapter_id)
|
|
|
|
def sleep(self, level: int = 1) -> bool:
|
|
"""
|
|
Put the worker to sleep, offloading model weights.
|
|
|
|
Args:
|
|
level: Sleep level. Level 1 offloads weights, level 2 also saves buffers.
|
|
"""
|
|
from vllm.device_allocator.cumem import CuMemAllocator
|
|
|
|
process_memory_before_sleep = get_process_gpu_memory(self.local_rank)
|
|
free_bytes_before_sleep = None
|
|
if process_memory_before_sleep is None:
|
|
free_bytes_before_sleep = current_omni_platform.get_free_memory()
|
|
|
|
# Save the buffers before level 2 sleep
|
|
if level == 2 and self.model_runner is not None:
|
|
model = self.model_runner.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())
|
|
process_memory_after_sleep = get_process_gpu_memory(self.local_rank)
|
|
if process_memory_before_sleep is not None and process_memory_after_sleep is not None:
|
|
freed_bytes = process_memory_before_sleep - process_memory_after_sleep
|
|
used_bytes = process_memory_after_sleep
|
|
accounting_scope = "process-scoped"
|
|
else:
|
|
free_bytes_after_sleep = current_omni_platform.get_free_memory()
|
|
assert free_bytes_before_sleep is not None
|
|
device_id = self.device.index if self.device.index is not None else 0
|
|
total = current_omni_platform.get_device_total_memory(device_id)
|
|
freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
|
|
used_bytes = total - free_bytes_after_sleep
|
|
accounting_scope = "device-scoped fallback"
|
|
assert freed_bytes >= 0, "Memory usage increased after sleeping."
|
|
logger.info(
|
|
"Sleep mode (%s) freed %.2f GiB memory, %.2f GiB memory is still in use.",
|
|
accounting_scope,
|
|
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) and self.model_runner is not None:
|
|
model = self.model_runner.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:
|
|
"""Get memory pool context for sleep mode support."""
|
|
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:
|
|
"""Shutdown the worker and cleanup distributed environment."""
|
|
destroy_distributed_env()
|
|
|
|
|
|
class CustomPipelineWorkerExtension:
|
|
def re_init_pipeline(self, custom_pipeline_args: dict[str, Any]) -> None:
|
|
"""
|
|
Re-initialize the pipeline with custom arguments.
|
|
|
|
Args:
|
|
custom_pipeline_args: Dictionary of arguments for custom pipeline initialization
|
|
"""
|
|
|
|
# Clean up old pipeline
|
|
if self.model_runner.pipeline is not None:
|
|
del self.model_runner.pipeline
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
# Get custom pipeline class name
|
|
custom_pipeline_name = custom_pipeline_args["pipeline_class"]
|
|
|
|
# Use the DiffusionWorker's load_model method which handles the forward context
|
|
self.load_model(
|
|
load_format="custom_pipeline",
|
|
custom_pipeline_name=custom_pipeline_name,
|
|
)
|
|
self.init_lora_manager()
|
|
|
|
|
|
class WorkerProc:
|
|
"""Wrapper that runs one Worker in a separate process."""
|
|
|
|
def __init__(
|
|
self,
|
|
od_config: OmniDiffusionConfig,
|
|
gpu_id: int,
|
|
broadcast_handle,
|
|
worker_extension_cls: str | None = None,
|
|
custom_pipeline_args: dict[str, Any] | None = None,
|
|
):
|
|
self.od_config = od_config
|
|
self.gpu_id = gpu_id
|
|
|
|
# Inter-process Communication
|
|
self.context = zmq.Context(io_threads=2)
|
|
|
|
# Initialize MessageQueue reader from handle
|
|
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)
|
|
if gpu_id == 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
|
|
|
|
# Create worker using WorkerWrapperBase for extension support
|
|
self.worker = self._create_worker(gpu_id, od_config, worker_extension_cls, custom_pipeline_args)
|
|
self._running = True
|
|
|
|
def _create_worker(
|
|
self,
|
|
gpu_id: int,
|
|
od_config: OmniDiffusionConfig,
|
|
worker_extension_cls: str | None,
|
|
custom_pipeline_args: dict[str, Any] | None = None,
|
|
) -> DiffusionWorker:
|
|
"""Create a worker instance. Override in subclasses for different worker types."""
|
|
wrapper = WorkerWrapperBase(
|
|
gpu_id=gpu_id,
|
|
od_config=od_config,
|
|
worker_extension_cls=worker_extension_cls,
|
|
custom_pipeline_args=custom_pipeline_args,
|
|
)
|
|
return wrapper
|
|
|
|
def return_result(self, output: DiffusionOutput):
|
|
"""Reply to client, only on rank 0."""
|
|
if self.result_mq is not None:
|
|
try:
|
|
pack_diffusion_output_shm(output)
|
|
except Exception as e:
|
|
logger.warning("SHM pack failed, falling back to raw enqueue: %s", e)
|
|
self.result_mq.enqueue(output)
|
|
|
|
def recv_message(self):
|
|
"""Receive messages from broadcast queue."""
|
|
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:
|
|
# Use execute_method from WorkerWrapperBase for consistent method resolution
|
|
result = self.worker.execute_method(method, *args, **kwargs)
|
|
return result, should_reply
|
|
except Exception as e:
|
|
logger.error(f"Error executing RPC: {e}", exc_info=True)
|
|
raise e
|
|
|
|
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:
|
|
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":
|
|
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(DiffusionOutput(error=str(e)))
|
|
|
|
elif isinstance(msg, dict) and msg.get("type") == "shutdown":
|
|
logger.info("Worker %s: Received shutdown message", self.gpu_id)
|
|
self._running = False
|
|
continue
|
|
|
|
else:
|
|
# Handle generation request
|
|
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:
|
|
logger.error(f"ZMQ error sending reply: {e}")
|
|
continue
|
|
|
|
logger.info("event loop terminated.")
|
|
try:
|
|
self.worker.shutdown()
|
|
except Exception as exc:
|
|
logger.warning("Worker %s: Shutdown encountered an error: %s", self.gpu_id, exc)
|
|
self.context.term()
|
|
|
|
@staticmethod
|
|
def worker_main(
|
|
rank: int,
|
|
od_config: OmniDiffusionConfig,
|
|
pipe_writer: mp.connection.Connection,
|
|
broadcast_handle,
|
|
worker_extension_cls: str | None = None,
|
|
custom_pipeline_args: dict[str, Any] | None = None,
|
|
) -> None:
|
|
"""Worker initialization and execution loops."""
|
|
from vllm_omni.plugins import load_omni_general_plugins
|
|
|
|
load_omni_general_plugins()
|
|
worker_proc = WorkerProc(
|
|
od_config,
|
|
gpu_id=rank,
|
|
broadcast_handle=broadcast_handle,
|
|
worker_extension_cls=worker_extension_cls,
|
|
custom_pipeline_args=custom_pipeline_args,
|
|
)
|
|
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.")
|
|
|
|
|
|
class WorkerWrapperBase:
|
|
"""
|
|
Wrapper base class that creates DiffusionWorker with optional worker_extension_cls support.
|
|
This enables dynamic inheritance for DiffusionWorker to extend with custom functionality.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
gpu_id: int,
|
|
od_config: OmniDiffusionConfig,
|
|
base_worker_class: type = DiffusionWorker,
|
|
worker_extension_cls: str | None = None,
|
|
custom_pipeline_args: dict[str, Any] | None = None,
|
|
):
|
|
"""
|
|
Initialize WorkerWrapperBase with support for worker extensions.
|
|
|
|
Args:
|
|
gpu_id: GPU device ID
|
|
od_config: OmniDiffusionConfig configuration
|
|
worker_extension_cls: Optional qualified name of worker extension class
|
|
custom_pipeline_args: Optional arguments for custom pipeline initialization
|
|
"""
|
|
self.gpu_id = gpu_id
|
|
self.od_config = od_config
|
|
self.base_worker_class = base_worker_class
|
|
self.worker_extension_cls = worker_extension_cls
|
|
self.custom_pipeline_args = custom_pipeline_args
|
|
|
|
# Prepare worker class with extension support
|
|
worker_class = self._prepare_worker_class()
|
|
|
|
# Create the actual worker instance
|
|
# When custom_pipeline_args is provided, skip initial model loading
|
|
# since re_init_pipeline will handle it. This avoids allocating memory
|
|
# through CuMemAllocator twice, which causes assertion failures in
|
|
# sleep mode.
|
|
self.worker = worker_class(
|
|
local_rank=gpu_id,
|
|
rank=gpu_id,
|
|
od_config=od_config,
|
|
skip_load_model=(self.custom_pipeline_args is not None),
|
|
)
|
|
|
|
# Re-initialize pipeline with custom pipeline if provided
|
|
if self.custom_pipeline_args is not None:
|
|
self.worker.re_init_pipeline(self.custom_pipeline_args)
|
|
|
|
def _prepare_worker_class(self) -> type:
|
|
"""
|
|
Prepare the worker class with optional extension.
|
|
Dynamically extends GPUWorker with worker_extension_cls if provided.
|
|
|
|
Returns:
|
|
The worker class (potentially extended)
|
|
"""
|
|
worker_class = self.base_worker_class
|
|
|
|
# If custom_pipeline_args is provided, use CustomPipelineWorkerExtension
|
|
if self.custom_pipeline_args is not None:
|
|
# Set worker_extension_cls to CustomPipelineWorkerExtension if not already set
|
|
if self.worker_extension_cls is None:
|
|
self.worker_extension_cls = CustomPipelineWorkerExtension
|
|
|
|
if self.worker_extension_cls:
|
|
if isinstance(self.worker_extension_cls, str):
|
|
worker_extension_cls = resolve_obj_by_qualname(self.worker_extension_cls)
|
|
else:
|
|
worker_extension_cls = self.worker_extension_cls
|
|
extended_calls = []
|
|
|
|
if worker_extension_cls not in worker_class.__bases__:
|
|
# Check for conflicts between worker and extension
|
|
for attr in dir(worker_extension_cls):
|
|
if attr.startswith("__"):
|
|
continue
|
|
if hasattr(worker_class, attr):
|
|
logger.warning(
|
|
f"Worker class {worker_class} already has attribute "
|
|
f"{attr}, which may conflict with worker extension "
|
|
f"class {worker_extension_cls}."
|
|
)
|
|
if callable(getattr(worker_extension_cls, attr)):
|
|
extended_calls.append(attr)
|
|
|
|
# Dynamically inherit the worker extension class
|
|
class_name = f"{worker_class.__name__}With{worker_extension_cls.__name__}"
|
|
worker_class = type(class_name, (worker_extension_cls, worker_class), {})
|
|
logger.info(
|
|
"Created extended worker class %s from %s for extended calls %s",
|
|
class_name,
|
|
worker_extension_cls,
|
|
extended_calls,
|
|
)
|
|
|
|
return worker_class
|
|
|
|
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.worker.generate(requests)
|
|
|
|
def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
|
|
"""
|
|
Execute a forward pass.
|
|
|
|
Args:
|
|
reqs: List of diffusion requests
|
|
od_config: OmniDiffusionConfig configuration
|
|
|
|
Returns:
|
|
DiffusionOutput with generated results
|
|
"""
|
|
return self.worker.execute_model(reqs, od_config)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
"""
|
|
Load model weights.
|
|
|
|
Args:
|
|
weights: Iterable of (name, tensor) tuples
|
|
|
|
Returns:
|
|
Set of loaded weight names
|
|
"""
|
|
return self.worker.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.
|
|
|
|
Returns:
|
|
True on success
|
|
"""
|
|
return self.worker.sleep(level)
|
|
|
|
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.
|
|
|
|
Returns:
|
|
True on success
|
|
"""
|
|
return self.worker.wake_up(tags)
|
|
|
|
def shutdown(self) -> None:
|
|
"""Shutdown the worker and cleanup resources."""
|
|
return self.worker.shutdown()
|
|
|
|
def execute_method(self, method: str | bytes, *args, **kwargs) -> Any:
|
|
"""
|
|
Execute a method on the worker.
|
|
|
|
Args:
|
|
method: Method name (str) or serialized callable (bytes)
|
|
|
|
Returns:
|
|
Result of the method execution (type depends on the method)
|
|
|
|
Raises:
|
|
Exception: If method execution fails
|
|
"""
|
|
try:
|
|
# Method resolution order:
|
|
# 1. If method is defined in this class, it will be called directly
|
|
# 2. Otherwise, since we define `__getattr__` and redirect attribute
|
|
# query to `self.worker`, the method will be called on the worker
|
|
assert isinstance(method, str), "Method must be str"
|
|
func = getattr(self.worker, method)
|
|
return func(*args, **kwargs)
|
|
|
|
except Exception as e:
|
|
msg = f"Error executing method {method!r}. This might cause issues in distributed execution."
|
|
logger.exception(msg)
|
|
raise e
|
|
|
|
def __getattr__(self, attr: str):
|
|
"""Delegate attribute access to the wrapped worker."""
|
|
return getattr(self.worker, attr)
|