[Bugfix] Validate sleep stage_ids before blocking admission (#8501)

Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com>

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}

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:

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()

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: