[BugFix]Fix diffusion profile RPC None results

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -730,6 +730,39 @@ def test_initialize_client_requires_replica_id(self):
                 batch_size=1,
             )
 
+    @pytest.mark.asyncio
+    async def test_collective_rpc_async_returns_none_result(self, monkeypatch):
+        client = self._make_client()
+        client._owns_process = False
+        client._proc = None
+        client._encoder.encode.return_value = b"encoded-rpc"
+
+        async def _unexpected_poll(*_, **__):
+            raise AssertionError("collective_rpc_async should not keep polling after a None rpc_result arrives")
+
+        client._response_poller = SimpleNamespace(poll=_unexpected_poll)
+
+        rpc_id = "rpc-none"
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.stage_diffusion_client.uuid.uuid4",
+            lambda: SimpleNamespace(hex=rpc_id),
+        )
+
+        def _drain() -> None:
+            client._rpc_results[rpc_id] = None
+
+        client._drain_responses = _drain
+
+        result = await client.collective_rpc_async(
+            method="profile",
+            timeout=0.01,
+            args=(False, None),
+        )
+
+        assert result is None
+        client._request_socket.send.assert_called_once_with(b"encoded-rpc")
+        assert rpc_id not in client._pending_rpcs
+
 
 # ───────── monitor thread & death sentinel integration tests ─────────
 

vllm_omni/diffusion/stage_diffusion_client.py

@@ -37,6 +37,7 @@
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
 logger = init_logger(__name__)
+_MISSING_RPC_RESULT = object()
 
 
 def create_diffusion_client(
@@ -506,8 +507,8 @@ async def collective_rpc_async(
         try:
             while True:
                 self._drain_responses()
-                result = self._rpc_results.pop(rpc_id, None)
-                if result is not None:
+                result = self._rpc_results.pop(rpc_id, _MISSING_RPC_RESULT)
+                if result is not _MISSING_RPC_RESULT:
                     return result
                 proc = self._proc_manager.proc
                 if self._engine_dead or not proc.is_alive():

vllm_omni/entrypoints/openai/api_server.py

@@ -3234,9 +3234,9 @@ async def start_profile(raw_request: Request, request: ProfileRequest | None = N
         stages = request.stages if request else None
         logger.info("Starting profiler for stages: %s", stages if stages else "all")
         engine_client = raw_request.app.state.engine_client
-        result = await engine_client.start_profile(stages=stages)
+        await engine_client.start_profile(stages=stages)
         logger.info("Profiler started.")
-        return JSONResponse(content=result)
+        return JSONResponse(content={"status": "SUCCESS"})
     except Exception as e:
         logger.exception("Failed to start profiler: %s", e)
         raise HTTPException(
@@ -3260,9 +3260,9 @@ async def stop_profile(raw_request: Request, request: ProfileRequest | None = No
         stages = request.stages if request else None
         logger.info("Stopping profiler for stages: %s", stages if stages else "all")
         engine_client = raw_request.app.state.engine_client
-        result = await engine_client.stop_profile(stages=stages)
+        await engine_client.stop_profile(stages=stages)
         logger.info("Profiler stopped.")
-        return JSONResponse(content=result)
+        return JSONResponse(content={"status": "SUCCESS"})
     except Exception as e:
         logger.exception("Failed to stop profiler: %s", e)
         raise HTTPException(