[Refactor]: clean up diffusion worker interfaces (#5217)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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()

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
 

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": (),

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)

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:

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