diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 9a740519f..a6575d5cc 100644 --- a/docs/features/sleep_mode.md +++ b/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. | diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py index b1c5f1189..7374483a4 100644 --- a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py +++ b/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 diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index e4b229f90..7da8f1396 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/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()) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 63f13d725..34fd76f22 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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 ) + self._record_stage_sleep(ar_stage_ids, sleep_tags) + if level == 2: + self._level2_sleeping = True if diffusion_stage_ids: - final_acks.extend(await self._sleep_diffusion(diffusion_stage_ids, level)) + 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) - self._record_stage_sleep( - ar_stage_ids + diffusion_stage_ids, - [CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value], - ) - if level == 2: - self._level2_sleeping = True 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) + 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: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 59c3e939b..28e4f68bd 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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") - acks = await engine_client.sleep(stage_ids=request.stage_ids, level=request.level) - for sid in request.stage_ids: - sleeping_set.add(sid) + try: + acks = await engine_client.sleep(stage_ids=request.stage_ids, level=request.level) + 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") - acks = await engine_client.wake_up(stage_ids=request.stage_ids) + 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)