diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 70df012389b..1b015c02a88 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/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 ───────── diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 8edf78bfb36..1e443b973d2 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/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(): diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 35b9dc20fac..ebdfd221273 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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(