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

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

docs/features/sleep_mode.md

@@ -251,8 +251,8 @@ curl -X POST http://localhost:8000/v1/images/generations \
 
 | Method | Arguments | Return Type | Description |
 | :--- | :--- | :--- | :--- |
-| **sleep** | `stage_ids: List[int], level: int` | `List[OmniACK]` | Triggers hibernation for specified stages. |
-| **wake_up** | `stage_ids: List[int]` | `List[OmniACK]` | Reloads weights and re-maps memory. |
+| **sleep** | `stage_ids: List[int], level: int` | `List[OmniACK]` | Triggers hibernation for specified stages. Raises `RuntimeError` if a stage fails; the stages stay marked as sleeping, so `wake_up` brings them back, unless a level-2 sleep already discarded weights on some replica. |
+| **wake_up** | `stage_ids: List[int]` | `List[OmniACK]` | Reloads weights and re-maps memory. Raises `RuntimeError` if a stage fails; that stage stays sleeping and the call can be retried. |
 | **pause_generation** | `mode: str, stage_ids: List[int]` | `None` | Stops admission; with `mode="keep"` also stops diffusion schedulers and returns after their ACK. |
 | **resume_generation** | `stage_ids: List[int]` | `None` | Reopens paused schedulers, then admission. |
 
@@ -261,7 +261,7 @@ curl -X POST http://localhost:8000/v1/images/generations \
 | Field | Type | Description |
 | :--- | :--- | :--- |
 | **task_id** | `str` | Unique identifier for the operation. |
-| **status** | `str` | `SUCCESS` or `ERROR`. |
+| **status** | `str` | `SUCCESS` or `ERROR`. `sleep` and `wake_up` raise instead of returning an `ERROR` ACK, and `/v1/omni/sleep` and `/v1/omni/wakeup` return HTTP 500 with the error. |
 | **stage_id** | `int` | The ID of the stage that responded. |
 | **rank** | `int` | The rank ID within the Tensor Parallel group. |
 | **freed_bytes** | `int` | Actual amount of physical VRAM reclaimed. |

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

vllm_omni/entrypoints/async_omni.py

@@ -940,6 +940,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
             await self._clear_frontend_mm_cache()
 
         self._final_output_handler()
+        sleep_tags = [CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value]
         final_acks: list[OmniACK] = []
         if ar_stage_ids:
             self._hold_admission_until_resume = True
@@ -968,16 +969,23 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
                 )
                 for sid in ar_stage_ids
             )
-
-        if diffusion_stage_ids:
-            final_acks.extend(await self._sleep_diffusion(diffusion_stage_ids, level))
-
-        self._record_stage_sleep(
-            ar_stage_ids + diffusion_stage_ids,
-            [CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value],
-        )
+            self._record_stage_sleep(ar_stage_ids, sleep_tags)
             if level == 2:
                 self._level2_sleeping = True
+
+        if diffusion_stage_ids:
+            acks: list[OmniACK] = []
+            try:
+                acks = await self._sleep_diffusion(diffusion_stage_ids, level)
+            finally:
+                # A failed sleep may have released part of a stage, so wake_up must still reach it.
+                self._record_stage_sleep(diffusion_stage_ids, sleep_tags)
+                # Level 2 has discarded weights only where a replica finished its sleep.
+                if level == 2 and any(self._diffusion_ack_error(ack) is None for ack in acks):
+                    self._level2_sleeping = True
+            self._raise_on_diffusion_errors("handle_sleep_task", acks)
+            final_acks.extend(acks)
+
         return final_acks
 
     async def _sleep_diffusion(self, stage_ids: list[int], level: int) -> list[OmniACK]:
@@ -989,15 +997,45 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         logger.info("[%s] Sleep (diffusion) initiated (Task: %s).", self._name, task_id)
         task = OmniSleepTask(level=level, task_id=task_id)
         rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids)
+        return await self._resolve_diffusion_acks(rpc_results)
+
+    async def _resolve_diffusion_acks(self, rpc_results: list[Any]) -> list[OmniACK]:
+        """Resolve the ACKs of a diffusion worker RPC."""
         final_acks: list[OmniACK] = []
         for stage_res in rpc_results:
             worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
             for ack in worker_acks:
-                if ack is not None:
+                if ack is None:
+                    continue
+                # StagePool's result for a failed RPC has no task_id to resolve.
+                if not isinstance(ack, dict) or "task_id" in ack:
                     await self.event_resolver.resolve(ack)
                 final_acks.append(ack)
         return final_acks
 
+    @classmethod
+    def _raise_on_diffusion_errors(cls, method: str, acks: list[OmniACK]) -> None:
+        errors = [error for ack in acks if (error := cls._diffusion_ack_error(ack))]
+        if errors:
+            raise RuntimeError(f"{method} failed: {'; '.join(errors)}")
+
+    @staticmethod
+    def _diffusion_ack_error(ack: OmniACK | dict[str, Any]) -> str | None:
+        # In-process stages return OmniACK, subprocess stages return its dict form,
+        # and StagePool returns {"supported": False, "error": ...} when the RPC failed.
+        if isinstance(ack, dict):
+            status, error_msg, rpc_error = ack.get("status"), ack.get("error_msg"), ack.get("error")
+            # A subprocess RPC error arrives as {"error": True, "reason": ...}.
+            if rpc_error is True:
+                rpc_error = ack.get("reason") or "RPC failed"
+            if not rpc_error and (ack.get("todo") or ack.get("supported") is False):
+                rpc_error = str(ack)
+        else:
+            status, error_msg, rpc_error = getattr(ack, "status", None), getattr(ack, "error_msg", None), None
+        if status == "ERROR":
+            return error_msg or "worker reported ERROR"
+        return rpc_error
+
     async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]:
         """Wake stages after sleep.
 
@@ -1052,11 +1090,14 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
                 )
                 for sid in ar_stage_ids
             )
+            self._clear_stage_sleep(ar_stage_ids, requested_tags)
 
+        for sid in diffusion_stage_ids:
+            final_acks.extend(await self._wake_diffusion([sid], requested_tags))
+            self._clear_stage_sleep([sid], requested_tags)
         if diffusion_stage_ids:
-            final_acks.extend(await self._wake_diffusion(diffusion_stage_ids, requested_tags))
+            await asyncio.sleep(0.1)
 
-        self._clear_stage_sleep(target_stage_ids, requested_tags)
         # Only clear the level-2 flag once all tags are warm, in case partial
         # wake support (e.g. tags=["kv_cache"] only) is added in the future.
         if not getattr(self, "_sleeping_tags", None):
@@ -1083,22 +1124,19 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         logger.info("[%s] Wake-up (diffusion) initiated (Task: %s).", self._name, task_id)
         task = OmniWakeTask(tags=requested_tags, task_id=task_id)
         rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
-        final_acks: list[OmniACK] = []
-        for stage_res in rpc_results:
-            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
-            for ack in worker_acks:
-                if ack is not None:
-                    await self.event_resolver.resolve(ack)
-                    final_acks.append(ack)
-        await asyncio.sleep(0.1)
-        return final_acks
+        acks = await self._resolve_diffusion_acks(rpc_results)
+        self._raise_on_diffusion_errors("handle_wake_task", acks)
+        return acks
 
-    async def is_sleeping(self) -> bool:
-        """Return whether all stages are sleeping.
+    async def is_sleeping(self, stage_ids: list[int] | None = None) -> bool:
+        """Return whether all stages are sleeping, or with ``stage_ids``,
+        whether any of those stages is recorded as sleeping.
 
         TODO(AsyncOmni): query the orchestrator once all stage backends expose
         a real sleeping-state RPC. For now we track the requested state locally.
         """
+        if stage_ids is not None:
+            return bool(self._sleeping_tags_for_stages(stage_ids))
         return bool(getattr(self, "_sleeping_tags", None))
 
     async def add_lora(self, lora_request: LoRARequest) -> bool:

vllm_omni/entrypoints/openai/api_server.py

@@ -3018,9 +3018,13 @@ async def omni_sleep(request: OmniSleepRequest, raw_request: Request):
     sleeping_set = raw_request.app.state.sleeping_stages
     if not hasattr(engine_client, "sleep"):
         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)
-    for sid in request.stage_ids:
-        sleeping_set.add(sid)
+    except RuntimeError as e:
+        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to sleep: {e}") from e
+    finally:
+        # Mirror the engine: a failed sleep may still leave stages for wakeup to reach.
+        sleeping_set.update([sid for sid in request.stage_ids if await engine_client.is_sleeping(stage_ids=[sid])])
     return {
         "status": "SUCCESS",
         "acks": [dataclasses.asdict(a) if dataclasses.is_dataclass(a) and not isinstance(a, type) else a for a in acks],
@@ -3036,7 +3040,12 @@ async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request):
         return {"status": "SKIPPED", "reason": "Target stages are not sleeping."}
     if not hasattr(engine_client, "wake_up"):
         raise HTTPException(status_code=501, detail="Engine does not support wake_up")
+    try:
         acks = await engine_client.wake_up(stage_ids=request.stage_ids)
+    except NotImplementedError:
+        raise
+    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:
         if sid in sleeping_set:
             sleeping_set.remove(sid)