diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md index 97d1db499..9634adad3 100644 --- a/docs/design/module/dit_module.md +++ b/docs/design/module/dit_module.md @@ -267,7 +267,7 @@ class WorkerProc: self.result_mq = MessageQueue(n_reader=1, ...) # Initialize GPU worker - self.worker = GPUWorker(local_rank=gpu_id, rank=gpu_id, od_config=od_config) + self.worker = self._create_worker(gpu_id, od_config, ...) ``` **Initialization Steps**: @@ -289,13 +289,18 @@ class WorkerProc: #### 3.2 GPU Worker ```python -class GPUWorker: - def init_device_and_model(self): +class DiffusionWorker: + def init_device(self): # Set distributed environment variables os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) - # Initialize PyTorch distributed + # Select the device and initialize distributed execution + self.device = current_omni_platform.get_torch_device(rank) + current_omni_platform.set_device(self.device) + self.vllm_config = _create_diffusion_worker_vllm_config( + self.device, self.od_config + ) init_distributed_environment(world_size, rank) parallel_config = self.od_config.parallel_config initialize_model_parallel( @@ -306,16 +311,18 @@ class GPUWorker: pipeline_parallel_size=parallel_config.pipeline_parallel_size, ) - # Load model - model_loader = DiffusersPipelineLoader(load_config, self.od_config) - self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}") + def load_model(self, load_format="default", custom_pipeline_name=None): + self.model_runner.load_model( + load_format=load_format, + custom_pipeline_name=custom_pipeline_name, + ) - # Setup cache backend - from vllm_omni.diffusion.cache.selector import get_cache_backend - self.cache_backend = get_cache_backend(od_config.cache_backend, od_config.cache_config) - - if self.cache_backend is not None: - self.cache_backend.enable(self.pipeline) + def init_lora_manager(self): + self.lora_manager = DiffusionLoRAManager( + pipeline=self.model_runner.pipeline, + device=self.device, + dtype=self.od_config.dtype, + ) ``` **Key Features**: @@ -329,17 +336,17 @@ class GPUWorker: #### 3.3 Worker Busy Loop ```python -def worker_busy_loop(self): +def _worker_busy_loop(self): while self._running: # 1. Receive unified message (generation request, RPC request, or shutdown) - msg = self.recv_message() + msg = self.mq.dequeue(timeout=1.0) # 2. Route message based on type if isinstance(msg, dict) and msg.get("type") == "rpc": # Handle RPC request - result, should_reply = self.execute_rpc(msg) + result, should_reply = self._execute_rpc(msg) if should_reply: - self.return_result(result) + self._return_result(result) elif isinstance(msg, dict) and msg.get("type") == "shutdown": # Handle shutdown message @@ -348,7 +355,7 @@ def worker_busy_loop(self): else: # Handle generation request (OmniDiffusionRequest list) output = self.worker.execute_model(msg, self.od_config) - self.return_result(output) + self._return_result(output) ``` **Execution Flow**: @@ -920,8 +927,8 @@ def initialize_model_parallel( └─> DiffusionEngine submits scheduled request to executor.add_req(req) 4. Worker Execution - └─> WorkerProc.worker_busy_loop() - └─> GPUWorker.execute_model(reqs) + └─> WorkerProc._worker_busy_loop() + └─> DiffusionWorker.execute_model(req) └─> Pipeline.forward(req) ├─> encode_prompt() ├─> prepare_latents() diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index fb45cd41d..6a49d0909 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -5,7 +5,6 @@ Unit tests for DiffusionWorker class. This module tests the DiffusionWorker implementation: -- load_weights: Loading model weights - sleep: Putting worker into sleep mode (levels 1 and 2) - wake_up: Waking worker from sleep mode """ @@ -59,38 +58,6 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config): return worker -class TestDiffusionWorkerLoadWeights: - """Test DiffusionWorker.load_weights method.""" - - def test_load_weights_calls_pipeline(self, mocker: MockerFixture, mock_gpu_worker): - """Test that load_weights delegates to model_runner.load_weights.""" - # Setup mock weights - mock_weights = [ - ("layer1.weight", torch.randn(10, 10)), - ("layer2.weight", torch.randn(20, 20)), - ] - expected_loaded = {"layer1.weight", "layer2.weight"} - - # Configure model_runner mock - mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=expected_loaded) - - # Call load_weights - result = mock_gpu_worker.load_weights(mock_weights) - - # Verify model_runner.load_weights was called with the weights - mock_gpu_worker.model_runner.load_weights.assert_called_once_with(mock_weights) - assert result == expected_loaded - - def test_load_weights_empty_iterable(self, mocker: MockerFixture, mock_gpu_worker): - """Test load_weights with empty weights iterable.""" - mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=set()) - - result = mock_gpu_worker.load_weights([]) - - mock_gpu_worker.model_runner.load_weights.assert_called_once_with([]) - assert result == set() - - def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods(): from types import SimpleNamespace diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 9b8e08a74..5f5d90f64 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -526,7 +526,7 @@ class TestWorkerProcRpcRankStatus: monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object) - result, should_reply = proc.execute_rpc( + result, should_reply = proc._execute_rpc( { "method": "remove_lora", "args": (), @@ -549,7 +549,7 @@ class TestWorkerProcRpcRankStatus: proc.worker.execute_method = Mock(side_effect=RuntimeError("local boom")) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) - result, should_reply = proc.execute_rpc( + result, should_reply = proc._execute_rpc( { "method": "add_lora", "args": (), @@ -575,7 +575,7 @@ class TestWorkerProcRpcRankStatus: proc = self._make_worker_proc() with pytest.raises(ValueError, match="collect_rank_status requires exec_all_ranks=True"): - proc.execute_rpc( + proc._execute_rpc( { "method": "ping", "args": (), @@ -594,7 +594,7 @@ class TestWorkerProcRpcRankStatus: proc.worker.execute_method = Mock(side_effect=original) with pytest.raises(ValueError) as excinfo: - proc.execute_rpc( + proc._execute_rpc( { "method": "bad", "args": (), diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index e9c4d613f..170faec80 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -161,19 +161,6 @@ class TestWorkerWrapperBaseExtension: class TestWorkerWrapperBaseDelegation: """Test WorkerWrapperBase delegation to wrapped worker.""" - def test_generate_delegation(self, mocker: MockerFixture, mock_od_config): - """Test that generate() delegates to worker.generate().""" - mocker.patch.object(DiffusionWorker, "__init__", return_value=None) - wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - mock_output = mocker.Mock() - wrapper.worker.generate = mocker.Mock(return_value=mock_output) - - mock_requests = [mocker.Mock()] - result = wrapper.generate(mock_requests) - - wrapper.worker.generate.assert_called_once_with(mock_requests) - assert result == mock_output - def test_execute_model_delegation(self, mocker: MockerFixture, mock_od_config): """Test that execute_model() delegates to worker.execute_model().""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) @@ -181,25 +168,12 @@ class TestWorkerWrapperBaseDelegation: mock_output = mocker.Mock() wrapper.worker.execute_model = mocker.Mock(return_value=mock_output) - mock_reqs = [mocker.Mock()] - result = wrapper.execute_model(mock_reqs, mock_od_config) + mock_req = mocker.Mock() + result = wrapper.execute_model(mock_req, mock_od_config) - wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config, kv_prefetch_job=None) + wrapper.worker.execute_model.assert_called_once_with(mock_req, mock_od_config, kv_prefetch_job=None) assert result == mock_output - def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config): - """Test that load_weights() delegates to worker.load_weights().""" - mocker.patch.object(DiffusionWorker, "__init__", return_value=None) - wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - expected_result = {"weight1", "weight2"} - wrapper.worker.load_weights = mocker.Mock(return_value=expected_result) - - mock_weights = [("weight1", mocker.Mock()), ("weight2", mocker.Mock())] - result = wrapper.load_weights(mock_weights) - - wrapper.worker.load_weights.assert_called_once_with(mock_weights) - assert result == expected_result - def test_sleep_delegation(self, mocker: MockerFixture, mock_od_config): """Test that sleep() delegates to worker.sleep().""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 1500ccd3f..df3ee5910 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -13,7 +13,7 @@ import multiprocessing as mp import os import signal import traceback -from collections.abc import Iterable, Iterator +from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import dataclass from typing import Any @@ -387,10 +387,6 @@ class DiffusionWorker: 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) - def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: """Start or stop profiling for this GPU worker. @@ -506,11 +502,6 @@ class DiffusionWorker: raise logger.warning("LoRA activation skipped: %s", exc) - 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) @@ -525,7 +516,7 @@ class DiffusionWorker: def pin_lora(self, adapter_id: int) -> bool: return self.lora_manager.pin_adapter(adapter_id) - def sleep(self, level: int = 1) -> bool: + def sleep(self, level: int = 1) -> int: """ Put the worker to sleep, offloading model weights. @@ -598,7 +589,7 @@ class DiffusionWorker: logger.info(f"[Worker {self.rank}] Wake-up complete.") return True - def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK: + def handle_sleep_task(self, task: OmniSleepTask | dict) -> OmniACK | None: from vllm_omni.platforms import current_omni_platform try: @@ -655,7 +646,7 @@ class DiffusionWorker: pass return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e)) - def handle_wake_task(self, task: OmniWakeTask) -> OmniACK: + def handle_wake_task(self, task: OmniWakeTask | dict) -> OmniACK | None: from vllm_omni.platforms import current_omni_platform try: @@ -798,7 +789,7 @@ class WorkerProc: od_config: OmniDiffusionConfig, worker_extension_cls: str | None, custom_pipeline_args: dict[str, Any] | None = None, - ) -> DiffusionWorker: + ) -> "WorkerWrapperBase": """Create a worker instance. Override in subclasses for different worker types.""" worker_cls_path = current_omni_platform.get_diffusion_worker_cls() base_worker_class = resolve_obj_by_qualname(worker_cls_path) @@ -811,7 +802,7 @@ class WorkerProc: ) return wrapper - def return_result(self, output: Any): + def _return_result(self, output: Any) -> None: """Reply to client, only on rank 0.""" if self.result_mq is not None: if isinstance(output, OmniACK): @@ -824,10 +815,6 @@ class WorkerProc: logger.warning("SHM pack failed for model output: %s", e) self.result_mq.enqueue(output) - def recv_message(self): - """Receive messages from broadcast queue.""" - return self.mq.dequeue(indefinite=True) - def _gather_rpc_rank_statuses(self, status: dict[str, Any]) -> list[dict[str, Any]]: if not torch.distributed.is_initialized(): return [status] @@ -840,7 +827,7 @@ class WorkerProc: logger.warning("RPC rank status gather returned missing entries for ranks: %s", missing_ranks) return [s for s in statuses if s is not None] - def execute_rpc(self, rpc_request: dict) -> tuple[object | None, bool]: + 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", ()) @@ -905,7 +892,7 @@ class WorkerProc: raise rpc_exception return result, should_reply - def worker_busy_loop(self) -> None: + 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") @@ -930,21 +917,21 @@ class WorkerProc: if isinstance(msg, dict) and msg.get("type") == "sleep": task = OmniSleepTask(level=msg.get("level", 2), task_id=msg.get("task_id", "local")) ack = self.worker.handle_sleep_task(task) - self.return_result(ack) + self._return_result(ack) elif isinstance(msg, dict) and msg.get("type") == "wake_up": task = OmniWakeTask(tags=msg.get("tags"), task_id=msg.get("task_id", "local")) ack = self.worker.handle_wake_task(task) - self.return_result(ack) + self._return_result(ack) # Route message based on type elif isinstance(msg, dict) and msg.get("type") == "rpc": try: - result, should_reply = self.execute_rpc(msg) + result, should_reply = self._execute_rpc(msg) if should_reply: - self.return_result(result) + 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)}) + self._return_result({"status": "error", "error": str(e)}) elif isinstance(msg, dict) and msg.get("type") == "shutdown": logger.info("Worker %s: Received shutdown message", self.gpu_id) @@ -963,7 +950,7 @@ class WorkerProc: output = DiffusionOutput.from_exception(e) try: - self.return_result(output) + self._return_result(output) except zmq.ZMQError as e: logger.error(f"ZMQ error sending reply: {e}") continue @@ -1022,7 +1009,7 @@ class WorkerProc: "result_handle": worker_proc.result_mq_handle if rank == 0 else None, } ) - worker_proc.worker_busy_loop() + worker_proc._worker_busy_loop() except SystemExit: logger.info("Worker %d: Shutdown signal received, starting cleanup.", rank) raise @@ -1139,21 +1126,9 @@ class WorkerWrapperBase: 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], + req: OmniDiffusionRequest, od_config: OmniDiffusionConfig, kv_prefetch_job: KVPrefetchJob | None = None, ) -> DiffusionOutput: @@ -1161,32 +1136,20 @@ class WorkerWrapperBase: Execute a forward pass. Args: - reqs: List of diffusion requests + req: Diffusion request. od_config: OmniDiffusionConfig configuration kv_prefetch_job: Optional next-request KV prefetch descriptor. Returns: DiffusionOutput with generated results """ - return self.worker.execute_model(reqs, od_config, kv_prefetch_job=kv_prefetch_job) + return self.worker.execute_model(req, od_config, kv_prefetch_job=kv_prefetch_job) def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: """Execute one diffusion step.""" return self.worker.execute_stepwise(scheduler_output) - 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: + def sleep(self, level: int = 1) -> int: """ Put the worker to sleep. The worker should not process any requests. The caller should guarantee that no requests are being processed @@ -1194,11 +1157,11 @@ class WorkerWrapperBase: Args: level: The sleep level. Level 1 sleep will offload the model - weights and discard the kv cache. - Currently only support level 1. + weights and discard the kv cache. Level 2 also saves buffers. + Currently only support level 1 and level 2. Returns: - True on success + Bytes held by the allocator before sleeping. """ return self.worker.sleep(level) @@ -1219,10 +1182,10 @@ class WorkerWrapperBase: """ return self.worker.wake_up(tags) - def handle_sleep_task(self, task): + def handle_sleep_task(self, task: OmniSleepTask | dict) -> OmniACK | None: return self.worker.handle_sleep_task(task) - def handle_wake_task(self, task): + def handle_wake_task(self, task: OmniWakeTask | dict) -> OmniACK | None: return self.worker.handle_wake_task(task) def shutdown(self) -> None: diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index 8d85fae70..b49388fee 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/vllm_omni/diffusion/worker/utils.py @@ -163,12 +163,6 @@ class StepRequestState: return self.denoise_completed return self.chunk_index >= self.total_chunks - @property - def new_request(self) -> bool: - # TODO: this is only an approximation for current stepwise mode. - # A real "new request" signal should eventually come from scheduler/runner state transitions. - return self.step_index == 0 or self.timesteps is None - class BaseRunnerOutput(ABC): @abstractmethod