[Bugfix] Fix dynamic function call on collective_rpc of DiffusionWorker (#2217)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Samit <285365963@qq.com>

tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py

@@ -3,21 +3,54 @@
 
 from __future__ import annotations
 
+from contextlib import ExitStack
+
 import pytest
 
 from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
     vLLMOmniColocateWorkerExtensionForTest,
 )
+from tests.utils import hardware_test
 from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
+from vllm_omni.entrypoints.async_omni import AsyncOmni
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+CUSTOM_PIPELINE_CLASS = (
+    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
+)
+WORKER_EXTENSION_CLASS = (
+    "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
+)
+MODEL = "tiny-random/Qwen-Image"
+
+pytestmark = [pytest.mark.core_model]
 
 
+@pytest.mark.cpu
 def test_worker_extension_inheritance():
     assert issubclass(vLLMOmniColocateWorkerExtensionForTest, CustomPipelineWorkerExtension)
 
 
+@pytest.mark.cpu
 def test_worker_extension_test_function():
     assert (
         vLLMOmniColocateWorkerExtensionForTest.test_extension_name() == "vllm-omni-colocate-worker-extension-for-test"
     )
+
+
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_collective_rpc_test_extension_name():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+        )
+        after.callback(engine.shutdown)
+
+        result = await engine.collective_rpc(method="test_extension_name")
+
+    assert isinstance(result, list), f"Expected list, got {type(result)}"
+    assert result == [["vllm-omni-colocate-worker-extension-for-test"]]

vllm_omni/diffusion/stage_diffusion_client.py

@@ -186,23 +186,17 @@ class StageDiffusionClient:
                 return await asyncio.wait_for(result, timeout=timeout)
             return await result
 
-        if method in {"sleep", "wake_up"}:
-            loop = asyncio.get_running_loop()
-            return await loop.run_in_executor(
-                self._engine._executor,
-                self._engine.engine.collective_rpc,
-                method,
-                timeout,
-                args,
-                kwargs,
-                None,
-            )
-
-        return {
-            "supported": False,
-            "todo": True,
-            "reason": f"Diffusion stage collective_rpc method {method} is not implemented yet",
-        }
+        # Fall back to collective RPC for other methods
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            self._engine._executor,
+            self._engine.engine.collective_rpc,
+            method,
+            timeout,
+            args,
+            kwargs,
+            None,
+        )
 
     def shutdown(self) -> None:
         for task in self._tasks.values():