From a0e626caafa14a1d165d615aa1ae97aeb59a2d84 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:27:53 +0800 Subject: [PATCH] [Bugfix] Validate sleep stage_ids before blocking admission (#8501) Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- .../openai_api/test_omni_sleep_wakeup.py | 18 +++++++++++++-- .../test_async_omni_pause_sleep_routing.py | 22 +++++++++++++++++++ vllm_omni/entrypoints/async_omni.py | 8 +++++-- vllm_omni/entrypoints/openai/api_server.py | 4 ++++ 4 files changed, 48 insertions(+), 4 deletions(-) diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py index 2da64d274..45c14720a 100644 --- a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py +++ b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py @@ -125,9 +125,10 @@ def test_sleep_client_error_does_not_mark_stages(sleep_capable_engine, mocker): app = _make_app(sleep_capable_engine) client = TestClient(app) - with pytest.raises(ValueError, match="unknown stage id 9"): - client.post("/v1/omni/sleep", json={"stage_ids": [9], "level": 1}) + response = client.post("/v1/omni/sleep", json={"stage_ids": [9], "level": 1}) + assert response.status_code == 400 + assert "unknown stage id 9" in response.json()["detail"] assert app.state.sleeping_stages == set() @@ -213,6 +214,19 @@ def test_wakeup_failure_returns_error_and_keeps_sleeping_set(sleep_capable_engin assert app.state.sleeping_stages == {0} +def test_wakeup_client_error_returns_400_and_keeps_sleeping_set(sleep_capable_engine, mocker): + sleep_capable_engine.wake_up = mocker.AsyncMock(side_effect=ValueError("unknown stage id 9")) + app = _make_app(sleep_capable_engine) + app.state.sleeping_stages = {0, 9} + client = TestClient(app) + + response = client.post("/v1/omni/wakeup", json={"stage_ids": [9]}) + + assert response.status_code == 400 + assert "unknown stage id 9" in response.json()["detail"] + assert app.state.sleeping_stages == {0, 9} + + def test_wakeup_engine_not_support(sleep_incapable_engine): app = _make_app(sleep_incapable_engine) app.state.sleeping_stages = {0} diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index bc4665dcc..febd29a85 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -353,6 +353,28 @@ def test_sleep_waits_for_in_flight_generate_admission(): asyncio.run(run()) +@pytest.mark.cpu +def test_sleep_rejects_invalid_stage_ids_without_blocking_admission(): + """Invalid stage_ids must raise before _paused is set. + + Raising afterwards leaves no sleeping tags behind, so wake_up() early-returns + "already warm" and nothing ever clears _paused: generate() hangs forever. + """ + + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + + with pytest.raises(ValueError, match=r"Invalid stage_ids \[99\]"): + await omni.sleep(stage_ids=[99], level=1, mode="abort") + + assert omni._paused is False + assert omni._hold_admission_until_resume is False + omni._clear_frontend_mm_cache.assert_not_awaited() + omni.collective_rpc.assert_not_awaited() + + asyncio.run(run()) + + @pytest.mark.cpu def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc(): async def run() -> None: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 01c8ee588..097c7aca8 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -959,6 +959,12 @@ class AsyncOmni(AsyncOmniBase, EngineClient): engines have no EngineCore pause to hold, so ``wake_up`` restores admission and ``sleep → wake → generate`` keeps working. """ + # Validate before touching admission state. Raising after + # ``_paused = True`` would wedge generate() for good: nothing was + # slept, so wake_up() has no tags to restore and never reaches the + # block that clears ``_paused``. + ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + # Block admission before any sleep RPC so generate() waits on # _pause_cond during the drain/offload window. Wait until generate() # coroutines that already passed the pause check have submitted (or @@ -966,8 +972,6 @@ class AsyncOmni(AsyncOmniBase, EngineClient): async with self._pause_cond: self._paused = True await self._pause_cond.wait_for(lambda: getattr(self, "_admitting", 0) == 0) - - ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) # EngineCore.sleep resets receiver caches itself; only clear P0 here. if 0 in ar_stage_ids: await self._clear_frontend_mm_cache() diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 9cbc13894..88bcac8cb 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -3062,6 +3062,8 @@ async def omni_sleep(request: OmniSleepRequest, raw_request: Request): raise HTTPException(status_code=501, detail="Engine does not support sleep") try: acks = await engine_client.sleep(stage_ids=request.stage_ids, level=request.level) + except ValueError as e: + raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) from e except RuntimeError as e: raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to sleep: {e}") from e finally: @@ -3086,6 +3088,8 @@ async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request): acks = await engine_client.wake_up(stage_ids=request.stage_ids) except NotImplementedError: raise + except ValueError as e: + raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) from e except RuntimeError as e: raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to wake up: {e}") from e for sid in request.stage_ids: