[Bugfix] Raise when a diffusion stage fails sleep or wake_up (#7811)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>

tests/entrypoints/openai_api/test_omni_sleep_wakeup.py

@@ -32,6 +32,7 @@ def sleep_capable_engine(mocker):
     engine = mocker.MagicMock()
     engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)])
     engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)])
+    engine.is_sleeping = mocker.AsyncMock(return_value=True)
     return engine
 
 
@@ -90,6 +91,46 @@ def test_sleep_updates_sleeping_set(sleep_capable_engine):
     assert app.state.sleeping_stages == {5, 0, 1}
 
 
+def test_sleep_failure_keeps_stages_for_wakeup(sleep_capable_engine, mocker):
+    sleep_capable_engine.sleep = mocker.AsyncMock(side_effect=RuntimeError("handle_sleep_task failed: out of memory"))
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 1})
+
+    assert response.status_code == 500
+    assert "handle_sleep_task failed: out of memory" in response.json()["detail"]
+    # The stage may be partly asleep, so /v1/omni/wakeup must still reach it.
+    assert app.state.sleeping_stages == {0}
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+    assert response.status_code == 200
+    sleep_capable_engine.wake_up.assert_awaited_once_with(stage_ids=[0])
+
+
+def test_sleep_failure_marks_only_stages_the_engine_recorded(sleep_capable_engine, mocker):
+    sleep_capable_engine.sleep = mocker.AsyncMock(side_effect=RuntimeError("sleep failed: out of memory"))
+    sleep_capable_engine.is_sleeping = mocker.AsyncMock(side_effect=lambda stage_ids: stage_ids == [1])
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 1})
+
+    assert response.status_code == 500
+    assert app.state.sleeping_stages == {1}
+
+
+def test_sleep_client_error_does_not_mark_stages(sleep_capable_engine, mocker):
+    sleep_capable_engine.sleep = mocker.AsyncMock(side_effect=ValueError("unknown stage id 9"))
+    sleep_capable_engine.is_sleeping = mocker.AsyncMock(return_value=False)
+    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})
+
+    assert app.state.sleeping_stages == set()
+
+
 def test_sleep_engine_not_support(sleep_incapable_engine):
     app = _make_app(sleep_incapable_engine)
     client = TestClient(app)
@@ -159,6 +200,19 @@ def test_wakeup_empty_stage_ids(sleep_capable_engine):
     sleep_capable_engine.wake_up.assert_not_awaited()
 
 
+def test_wakeup_failure_returns_error_and_keeps_sleeping_set(sleep_capable_engine, mocker):
+    sleep_capable_engine.wake_up = mocker.AsyncMock(side_effect=RuntimeError("handle_wake_task failed: out of memory"))
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0}
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+
+    assert response.status_code == 500
+    assert "handle_wake_task failed: out of memory" in response.json()["detail"]
+    assert app.state.sleeping_stages == {0}
+
+
 def test_wakeup_engine_not_support(sleep_incapable_engine):
     app = _make_app(sleep_incapable_engine)
     app.state.sleeping_stages = {0}
@@ -224,6 +278,7 @@ def pure_diffusion_engine(mocker):
     engine.stage_configs = [{"stage_type": "diffusion"}]
     engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)])
     engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)])
+    engine.is_sleeping = mocker.AsyncMock(return_value=True)
     # Remove attributes that would make _get_vllm_config return a config
     del engine.get_vllm_config
     del engine.vllm_config

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -826,3 +826,241 @@ def test_pause_and_sleep_do_not_reset_unselected_stages_or_reset_twice(operation
         assert omni._clear_frontend_mm_cache.await_count == int(0 in stage_ids)
 
     asyncio.run(run())
+
+
+_SLEEP_TAGS = {CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value}
+
+# In-process stages return OmniACK, subprocess stages return its dict form,
+# and StagePool returns an error dict when the RPC itself failed.
+_FAILED_DIFFUSION_RESULTS = [
+    pytest.param(OmniACK(task_id="t", status="ERROR", error_msg="out of memory"), id="ack"),
+    pytest.param({"task_id": "t", "status": "ERROR", "error_msg": "out of memory"}, id="ack-dict"),
+    pytest.param({"supported": False, "error": "out of memory"}, id="rpc-error"),
+]
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("result", _FAILED_DIFFUSION_RESULTS)
+def test_sleep_raises_when_diffusion_stage_fails(result):
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.collective_rpc = AsyncMock(return_value=[result])
+
+        with pytest.raises(RuntimeError, match="handle_sleep_task failed: out of memory"):
+            await omni.sleep(level=1)
+
+        # The stage may be partly asleep, so it stays on record for wake_up.
+        assert omni._sleeping_tags == _SLEEP_TAGS
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("result", _FAILED_DIFFUSION_RESULTS)
+def test_wake_up_raises_when_diffusion_stage_fails(result):
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        await omni.sleep(level=1)
+        omni.collective_rpc = AsyncMock(return_value=[result])
+
+        with pytest.raises(RuntimeError, match="handle_wake_task failed: out of memory"):
+            await omni.wake_up()
+
+        assert omni._sleeping_tags == _SLEEP_TAGS
+        assert omni._paused is True
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_records_all_stages_when_diffusion_stage_fails():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        omni._sleep_diffusion = AsyncMock(side_effect=RuntimeError("handle_sleep_task failed"))
+
+        with pytest.raises(RuntimeError):
+            await omni.sleep(level=1)
+
+        assert omni._stage_sleeping_tags == {0: _SLEEP_TAGS, 1: _SLEEP_TAGS}
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_clears_ar_stage_when_diffusion_stage_fails():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        await omni.sleep(level=1)
+        omni._wake_diffusion = AsyncMock(side_effect=RuntimeError("handle_wake_task failed"))
+
+        with pytest.raises(RuntimeError):
+            await omni.wake_up()
+
+        assert omni._stage_sleeping_tags == {1: _SLEEP_TAGS}
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_level2_is_recorded_when_diffusion_stage_fails():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        omni._sleep_diffusion = AsyncMock(side_effect=RuntimeError("handle_sleep_task failed"))
+
+        with pytest.raises(RuntimeError):
+            await omni.sleep(level=2)
+
+        with pytest.raises(NotImplementedError):
+            await omni.wake_up(stage_ids=[0])
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_reaches_diffusion_stage_after_failed_sleep():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.collective_rpc = AsyncMock(
+            return_value=[
+                [
+                    OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0),
+                    OmniACK(task_id="t", status="ERROR", error_msg="out of memory"),
+                ]
+            ]
+        )
+        with pytest.raises(RuntimeError, match="out of memory"):
+            await omni.sleep(level=1)
+
+        omni.collective_rpc = AsyncMock(return_value=[OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)])
+        await omni.wake_up()
+
+        omni.collective_rpc.assert_awaited_once()
+        assert not omni._sleeping_tags
+        assert omni._paused is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_clears_diffusion_stages_that_woke_before_a_failure():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion", "diffusion"])
+        await omni.sleep(level=1)
+        omni.collective_rpc = AsyncMock(
+            side_effect=[
+                [OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)],
+                [OmniACK(task_id="t", status="ERROR", error_msg="out of memory")],
+            ]
+        )
+
+        with pytest.raises(RuntimeError, match="out of memory"):
+            await omni.wake_up()
+
+        assert omni._stage_sleeping_tags == {1: _SLEEP_TAGS}
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_settles_once_for_several_diffusion_stages(mocker):
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion", "diffusion"])
+        await omni.sleep(level=1)
+        settle = mocker.patch("vllm_omni.entrypoints.async_omni.asyncio.sleep", new=AsyncMock())
+
+        await omni.wake_up()
+
+        assert omni.collective_rpc.await_count == 3
+        settle.assert_awaited_once_with(0.1)
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("result", [{"supported": False}, {"todo": "not supported yet"}], ids=["unsupported", "todo"])
+def test_sleep_raises_on_rpc_result_without_error(result):
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.collective_rpc = AsyncMock(return_value=[result])
+
+        with pytest.raises(RuntimeError, match="handle_sleep_task failed"):
+            await omni.sleep(level=1)
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_rpc_failure_result_is_not_resolved_as_ack():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.event_resolver.resolve = AsyncMock()
+        omni.collective_rpc = AsyncMock(return_value=[{"supported": False, "error": "out of memory"}])
+
+        with pytest.raises(RuntimeError, match="out of memory"):
+            await omni.sleep(level=1)
+
+        omni.event_resolver.resolve.assert_not_awaited()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_raises_with_reason_of_rpc_error_envelope():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.collective_rpc = AsyncMock(return_value=[{"error": True, "reason": "worker crashed"}])
+
+        with pytest.raises(RuntimeError, match="handle_sleep_task failed: worker crashed"):
+            await omni.sleep(level=1)
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_failed_level2_sleep_can_still_be_woken():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.collective_rpc = AsyncMock(return_value=[OmniACK(task_id="t", status="ERROR", error_msg="out of memory")])
+        with pytest.raises(RuntimeError, match="out of memory"):
+            await omni.sleep(level=2)
+
+        omni.collective_rpc = AsyncMock(return_value=[OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)])
+        await omni.wake_up()
+
+        omni.collective_rpc.assert_awaited_once()
+        assert omni._paused is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_partial_level2_sleep_still_blocks_wake_up():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.collective_rpc = AsyncMock(
+            return_value=[
+                [
+                    OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0),
+                    OmniACK(task_id="t", status="ERROR", error_msg="out of memory"),
+                ]
+            ]
+        )
+        with pytest.raises(RuntimeError, match="out of memory"):
+            await omni.sleep(level=2)
+
+        with pytest.raises(NotImplementedError):
+            await omni.wake_up()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_is_sleeping_for_given_stages():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        await omni.sleep(stage_ids=[1], level=1)
+
+        assert await omni.is_sleeping(stage_ids=[1])
+        assert not await omni.is_sleeping(stage_ids=[0])
+
+    asyncio.run(run())