anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Zeyu Huang | 黃澤宇CursorHongsheng Liu
 authored and GitHub committed 2026-07-22 17:07:09 +08:00
1 parent 39e83e27f0
commit f43d761c26
6 files changed
+58 -153
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_worker_wrapper_base.py
+3 -29
View File
@@ -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)