[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>

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)