[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: