From d15b7ca3e12ce969a37547f7480470b5b755dbf3 Mon Sep 17 00:00:00 2001 From: Nguyen Kha Nhat Long Date: Mon, 24 Aug 2026 16:21:32 +0800 Subject: [PATCH] [Feature]: Support pause / resume and sleep / wake for AR stages in AsyncOmni (#6084) Signed-off-by: knlnguyen1802 Co-authored-by: Cursor --- .../test_async_omni_engine_abort_ack.py | 156 +++++++++ tests/engine/test_orchestrator.py | 65 ++++ tests/engine/test_rpc_result_router.py | 14 + .../engine/test_stage_pool_collective_rpc.py | 71 ++++ tests/entrypoints/test_async_omni.py | 52 +++ .../test_async_omni_pause_sleep_routing.py | 310 ++++++++++++++++++ tests/entrypoints/test_omni_sleep_mode.py | 8 + vllm_omni/engine/async_omni_engine.py | 46 ++- vllm_omni/engine/messages.py | 14 + vllm_omni/engine/orchestrator.py | 35 +- vllm_omni/engine/stage_pool.py | 19 ++ vllm_omni/entrypoints/async_omni.py | 303 +++++++++++++---- 12 files changed, 1024 insertions(+), 69 deletions(-) create mode 100644 tests/engine/test_async_omni_engine_abort_ack.py create mode 100644 tests/engine/test_stage_pool_collective_rpc.py create mode 100644 tests/entrypoints/test_async_omni_pause_sleep_routing.py diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py new file mode 100644 index 000000000..db095f9bd --- /dev/null +++ b/tests/engine/test_async_omni_engine_abort_ack.py @@ -0,0 +1,156 @@ +"""Unit tests for acknowledged abort on AsyncOmniEngine.""" + +from __future__ import annotations + +import asyncio +import queue +import threading +from types import SimpleNamespace + +import pytest +from pytest_mock import MockerFixture + +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.messages import AbortRequestMessage, AbortResultMessage +from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_engine(request_q: queue.Queue, rpc_q: queue.Queue) -> AsyncOmniEngine: + engine = object.__new__(AsyncOmniEngine) + engine.request_queue = SimpleNamespace(sync_q=request_q) + engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) + engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) + return engine + + +def test_abort_result_message_correlation_key(): + msg = AbortResultMessage(rpc_id="abc", success=True) + assert msg.rpc_correlation_key == ("abort", "abc") + + +def test_sync_abort_is_fire_and_forget(): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + + engine.abort(["req-1"]) + + msg = request_q.get_nowait() + assert isinstance(msg, AbortRequestMessage) + assert msg.request_ids == ["req-1"] + assert msg.rpc_id is None + assert rpc_q.empty() + engine._correlated_rpc_client.close() + + +def test_abort_async_waits_for_ack(mocker: MockerFixture): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_omni_engine.uuid.uuid4", + return_value=SimpleNamespace(hex="abort-rpc-1"), + ) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-a", "req-b"], timeout=1)) + msg = await asyncio.to_thread(request_q.get, True, 1) + assert isinstance(msg, AbortRequestMessage) + assert msg.request_ids == ["req-a", "req-b"] + assert msg.rpc_id == "abort-rpc-1" + assert not task.done() + + rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True)) + await task + + try: + asyncio.run(_run()) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_raises_on_orchestrator_error(mocker: MockerFixture): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_omni_engine.uuid.uuid4", + return_value=SimpleNamespace(hex="abort-rpc-err"), + ) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-x"], timeout=1)) + await asyncio.to_thread(request_q.get, True, 1) + rpc_q.put( + AbortResultMessage( + rpc_id="abort-rpc-err", + success=False, + error="stage abort failed", + ) + ) + with pytest.raises(RuntimeError, match="stage abort failed"): + await task + + try: + asyncio.run(_run()) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_times_out_without_result(mocker: MockerFixture): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_omni_engine.uuid.uuid4", + return_value=SimpleNamespace(hex="abort-rpc-timeout"), + ) + + async def _run() -> None: + with pytest.raises(TimeoutError, match="abort timed out"): + await engine.abort_async(["req-timeout"], timeout=0.05) + + try: + asyncio.run(_run()) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture): + class SignallingQueue(queue.Queue): + def __init__(self) -> None: + super().__init__(maxsize=1) + self.put_attempted = threading.Event() + + def put(self, item, block=True, timeout=None): + self.put_attempted.set() + return super().put(item, block=block, timeout=timeout) + + request_q = SignallingQueue() + request_q.put("queue-is-full") + request_q.put_attempted.clear() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_omni_engine.uuid.uuid4", + return_value=SimpleNamespace(hex="blocked-abort"), + ) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-blocked"], timeout=1)) + assert await asyncio.to_thread(request_q.put_attempted.wait, 1) + assert not task.done() + + assert request_q.get(timeout=1) == "queue-is-full" + msg = await asyncio.to_thread(request_q.get, True, 1) + assert isinstance(msg, AbortRequestMessage) + assert msg.rpc_id == "blocked-abort" + rpc_q.put(AbortResultMessage(rpc_id="blocked-abort", success=True)) + await task + + try: + asyncio.run(_run()) + finally: + engine._correlated_rpc_client.close() diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 3d63943a7..981331d45 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -19,6 +19,7 @@ from vllm.v1.metrics.stats import IterationStats from vllm_omni.engine.messages import ( AbortRequestMessage, + AbortResultMessage, AddCompanionRequestMessage, CollectiveRPCRequestMessage, CollectiveRPCResultMessage, @@ -805,10 +806,74 @@ async def test_run_abort(orchestrator_factory) -> None: assert stages[0].abort_calls == [["req-abort"]] assert stages[1].abort_calls == [] assert "req-abort" not in orchestrator_fixture.orchestrator.request_states + # Fire-and-forget abort must not emit an RPC result. + assert orchestrator_fixture.queues[2].sync_q.empty() finally: await _shutdown_orchestrator(orchestrator_fixture) +@pytest.mark.asyncio +async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> None: + stages = [ + FakeStageClient(stage_type="llm", final_output=False), + FakeStageClient(stage_type="llm", final_output=True), + ] + processors = [ + FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[1], finished=True)]), + FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[2], finished=True)]), + ] + orchestrator_fixture = orchestrator_factory(stages, output_processors=processors) + request = SimpleNamespace(request_id="req-ack", prompt_token_ids=[1, 2, 3]) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-ack", + prompt=request, + original_prompt={"prompt": "cancel with ack"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + await _wait_for(lambda: len(stages[0].add_request_calls) == 1) + + orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["req-ack"], rpc_id="abort-1")) + result = await _get_rpc_message(orchestrator_fixture) + assert isinstance(result, AbortResultMessage) + assert result.rpc_id == "abort-1" + assert result.success is True + assert result.error is None + assert result.rpc_correlation_key == ("abort", "abort-1") + assert stages[0].abort_calls == [["req-ack"]] + assert "req-ack" not in orchestrator_fixture.orchestrator.request_states + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_handle_abort_emits_error_result_on_failure() -> None: + rpc_queue: asyncio.Queue = asyncio.Queue() + orchestrator = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=rpc_queue, + stage_pools=[], + ) + + async def boom(_request_ids, *, abort=False): + del _request_ids, abort + raise RuntimeError("cleanup exploded") + + orchestrator._cleanup_request_ids = boom # type: ignore[method-assign] + + await orchestrator._handle_abort(AbortRequestMessage(request_ids=["req-fail"], rpc_id="abort-err")) + + result = rpc_queue.get_nowait() + assert isinstance(result, AbortResultMessage) + assert result.rpc_id == "abort-err" + assert result.success is False + assert "cleanup exploded" in (result.error or "") + + def _duplex_open_message( session_id: str, *, diff --git a/tests/engine/test_rpc_result_router.py b/tests/engine/test_rpc_result_router.py index ec983392b..1b3159617 100644 --- a/tests/engine/test_rpc_result_router.py +++ b/tests/engine/test_rpc_result_router.py @@ -4,6 +4,7 @@ from typing import Literal import pytest from vllm_omni.engine.messages import ( + AbortResultMessage, CollectiveRPCResultMessage, EngineQueueMessage, ErrorMessage, @@ -93,6 +94,19 @@ def test_rpc_result_router_does_not_broadcast_uncorrelated_nonfatal_errors(): router.close() +def test_rpc_result_router_routes_abort_results_by_correlation_id(): + source: queue.Queue = queue.Queue() + router = RpcResultRouter(source) + waiter = router.register(("abort", "abort-42")) + + source.put(AbortResultMessage(rpc_id="abort-42", success=True)) + + result = waiter.get(timeout=1) + assert isinstance(result, AbortResultMessage) + assert result.rpc_id == "abort-42" + router.close() + + def test_rpc_result_router_close_unblocks_waiters_and_stops_consumer(): source: queue.Queue = queue.Queue() router = RpcResultRouter(source) diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py new file mode 100644 index 000000000..b252cf0b8 --- /dev/null +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -0,0 +1,71 @@ +"""Unit tests for StagePool.collective_rpc EngineCore control dispatch.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from vllm_omni.engine.stage_pool import StagePool + +pytestmark = [pytest.mark.core_model] + + +def _make_pool(*, stage_type: str = "llm", **client_methods: AsyncMock) -> tuple[StagePool, SimpleNamespace]: + client = SimpleNamespace(stage_type=stage_type, **client_methods) + if "collective_rpc_async" not in client_methods: + client.collective_rpc_async = AsyncMock(return_value={"via": "collective"}) + pool = StagePool(0, [client]) # type: ignore[arg-type] + return pool, client + + +@pytest.mark.cpu +def test_collective_rpc_normalizes_none_args_on_control_helper(): + async def run() -> None: + pause = AsyncMock(return_value="paused") + pool, client = _make_pool(pause_scheduler_async=pause) + + result = await pool.collective_rpc(0, "pause_scheduler", args=None, kwargs={"mode": "abort"}) + + assert result == "paused" + pause.assert_awaited_once_with(mode="abort") + client.collective_rpc_async.assert_not_awaited() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_collective_rpc_unrelated_async_helper_uses_collective_path(): + async def run() -> None: + other = AsyncMock(return_value="should-not-run") + pool, client = _make_pool(reset_prefix_cache_async=other) + + result = await pool.collective_rpc(0, "reset_prefix_cache", timeout=1.5, args=("x",), kwargs={"k": 1}) + + assert result == {"via": "collective"} + other.assert_not_awaited() + client.collective_rpc_async.assert_awaited_once_with( + method="reset_prefix_cache", + timeout=1.5, + args=("x",), + kwargs={"k": 1}, + ) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_collective_rpc_control_helper_honors_timeout(): + async def run() -> None: + async def slow_sleep(*_args, **_kwargs): + await asyncio.sleep(1.0) + return "slept" + + pool, _client = _make_pool(sleep_async=slow_sleep) + result = await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,)) + assert result != "slept" + assert result["supported"] is False + + asyncio.run(run()) diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index a672c32e9..a830d1567 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -185,6 +185,58 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi asyncio.run(run()) +@pytest.mark.cpu +def test_abort_pops_request_states_only_after_ack(): + async def run(): + release = asyncio.Event() + seen_during_wait: list[int] = [] + + async def slow_abort_async(request_ids): + del request_ids + seen_during_wait.append(len(omni.request_states)) + await release.wait() + + omni = get_async_omni_instance(fake_abort_request=slow_abort_async) + omni.request_states["req-1-aaaa"] = SimpleNamespace( + request_id="req-1-aaaa", + external_request_id="req-1", + input_stream_task=None, + ) + + task = asyncio.create_task(omni.abort("req-1")) + await asyncio.sleep(0) + assert not task.done() + assert "req-1-aaaa" in omni.request_states + assert seen_during_wait == [1] + + release.set() + await task + assert "req-1-aaaa" not in omni.request_states + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_abort_propagates_engine_errors_without_popping_state(): + async def run(): + async def failing_abort_async(request_ids): + del request_ids + raise RuntimeError("orchestrator abort failed") + + omni = get_async_omni_instance(fake_abort_request=failing_abort_async) + omni.request_states["req-1-bbbb"] = SimpleNamespace( + request_id="req-1-bbbb", + external_request_id="req-1", + input_stream_task=None, + ) + + with pytest.raises(RuntimeError, match="orchestrator abort failed"): + await omni.abort("req-1") + assert "req-1-bbbb" in omni.request_states + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_accepts_request_after_repeated_cancellations(): async def run_test(): diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py new file mode 100644 index 000000000..5518d0b8b --- /dev/null +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -0,0 +1,310 @@ +"""Unit tests for AR EngineCore vs diffusion worker pause/sleep routing.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, call + +import pytest + +from vllm_omni.diffusion.data import CuMemTag, OmniACK +from vllm_omni.entrypoints.async_omni import AsyncOmni + +pytestmark = [pytest.mark.core_model] + + +def _make_omni(*, stage_types: list[str]) -> AsyncOmni: + omni = object.__new__(AsyncOmni) + omni._pause_cond = asyncio.Condition() + omni._paused = False + omni._hold_admission_until_resume = False + omni._sleeping_tags = set() + omni._level2_sleeping = False + omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock()) + omni._final_output_handler = lambda: None + + stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types] + omni.engine = SimpleNamespace( + stage_clients=stage_clients, + stage_vllm_configs=[None] * len(stage_types), + collective_rpc_async=AsyncMock(return_value=[True]), + ) + omni.collective_rpc = AsyncMock(return_value=[True]) + return omni + + +@pytest.mark.cpu +def test_split_stage_ids_by_type(): + omni = _make_omni(stage_types=["llm", "diffusion", "llm"]) + ar_ids, diff_ids = omni._split_stage_ids_by_type() + assert ar_ids == [0, 2] + assert diff_ids == [1] + + +@pytest.mark.cpu +def test_split_stage_ids_by_type_rejects_out_of_range(): + omni = _make_omni(stage_types=["llm", "diffusion"]) + with pytest.raises(ValueError, match=r"Invalid stage_ids \[2\].*0\.\.1"): + omni._split_stage_ids_by_type([0, 2]) + with pytest.raises(ValueError, match=r"Invalid stage_ids \[-1\]"): + omni._split_stage_ids_by_type([-1]) + + +@pytest.mark.cpu +def test_pause_generation_routes_ar_via_collective_rpc(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_mm_cache = AsyncMock() + omni.reset_encoder_cache = AsyncMock() + + await omni.pause_generation(mode="abort", clear_cache=True) + + assert omni._paused is True + assert omni._hold_admission_until_resume is True + omni.collective_rpc.assert_awaited_once_with( + method="pause_scheduler", + args=(), + kwargs={"mode": "abort", "clear_cache": True}, + stage_ids=[0], + ) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_pause_generation_still_rpcs_when_already_paused(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_mm_cache = AsyncMock() + omni.reset_encoder_cache = AsyncMock() + omni._paused = True + + await omni.pause_generation(mode="abort", clear_cache=True, stage_ids=[0]) + + omni.collective_rpc.assert_awaited_once_with( + method="pause_scheduler", + args=(), + kwargs={"mode": "abort", "clear_cache": True}, + stage_ids=[0], + ) + omni.reset_prefix_cache.assert_awaited_once() + omni.reset_mm_cache.assert_awaited_once() + omni.reset_encoder_cache.assert_awaited_once() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_resume_generation_resumes_ar_then_clears_frontend_pause(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + omni._paused = True + + await omni.resume_generation() + + omni.collective_rpc.assert_awaited_once_with( + method="resume_scheduler", + args=(), + kwargs=None, + stage_ids=[0], + ) + assert omni._paused is False + assert omni._hold_admission_until_resume is False + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_sleep_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0) + omni._sleep_diffusion = AsyncMock(return_value=[diffusion_ack]) + + acks = await omni.sleep(stage_ids=[0, 1], level=1, mode="abort") + + omni.collective_rpc.assert_awaited_once_with( + method="sleep", + args=(1, "abort"), + kwargs=None, + stage_ids=[0], + ) + omni._sleep_diffusion.assert_awaited_once_with([1], 1) + assert {ack.stage_id for ack in acks} == {0, 1} + assert any(ack.metadata.get("path") == "engine_core" for ack in acks if ack.stage_id == 0) + assert CuMemTag.WEIGHTS.value in omni._sleeping_tags + assert CuMemTag.KV_CACHE.value in omni._sleeping_tags + # Sleep gates frontend admission for the trainer resume contract. + assert omni._paused is True + assert omni._hold_admission_until_resume is True + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_sleep_blocks_admission_before_engine_core_rpc(): + """Problem 3: _paused must be set before awaiting sleep RPC.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + paused_at_rpc: list[bool] = [] + + async def rpc_side_effect(**kwargs): + if kwargs.get("method") == "sleep": + paused_at_rpc.append(omni._paused) + return [True] + + omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect) + + assert omni._paused is False + await omni.sleep(level=1, mode="abort") + + assert paused_at_rpc == [True] + assert omni._paused is True + assert omni._hold_admission_until_resume is True + # Sleep must not also call pause_scheduler (EngineCore.sleep pauses). + assert all(c.kwargs.get("method") != "pause_scheduler" for c in omni.collective_rpc.await_args_list) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + omni._paused = True + omni._hold_admission_until_resume = True + omni._sleeping_tags = {CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value} + diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0) + omni._wake_diffusion = AsyncMock(return_value=[diffusion_ack]) + + acks = await omni.wake_up(stage_ids=[0, 1]) + + omni.collective_rpc.assert_awaited_once() + wake_kwargs = omni.collective_rpc.await_args.kwargs + assert wake_kwargs["method"] == "wake_up" + assert wake_kwargs["stage_ids"] == [0] + assert set(wake_kwargs["kwargs"]["tags"]) == { + CuMemTag.WEIGHTS.value, + CuMemTag.KV_CACHE.value, + } + omni._wake_diffusion.assert_awaited_once() + assert {ack.stage_id for ack in acks} == {0, 1} + assert not omni._sleeping_tags + # Mixed/AR wake restores memory but does not resume frontend admission. + assert omni._paused is True + assert omni._hold_admission_until_resume is True + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_wake_up_does_not_resume_frontend_admission(): + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + await omni.sleep(level=1, mode="abort") + assert omni._paused is True + assert omni._hold_admission_until_resume is True + + await omni.wake_up() + + assert omni._paused is True + assert omni._hold_admission_until_resume is True + assert not omni._sleeping_tags + # Explicit resume is required after AR sleep/wake. + await omni.resume_generation() + assert omni._paused is False + assert omni._hold_admission_until_resume is False + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_sleep_level1_wake_without_tags_clears_all_sleeping_tags(): + """sleep(level=1) tracks WEIGHTS+KV; untagged wake_up must clear both.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + + await omni.sleep(level=1, mode="abort") + assert CuMemTag.WEIGHTS.value in omni._sleeping_tags + assert CuMemTag.KV_CACHE.value in omni._sleeping_tags + + await omni.wake_up() + assert not omni._sleeping_tags + assert omni._paused is True + assert omni._hold_admission_until_resume is True + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_sleep_diffusion_only_skips_engine_core_collective_rpc(): + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)]) + + await omni.sleep(level=1) + + # AR EngineCore sleep path must not run for diffusion-only. + assert call(method="sleep", args=(1, "abort"), kwargs=None, stage_ids=[0]) not in ( + omni.collective_rpc.await_args_list + ) + omni._sleep_diffusion.assert_awaited_once_with([0], 1) + assert omni._paused is True + assert omni._hold_admission_until_resume is False + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_wake_up_restores_admission_for_diffusion_only(): + """Pure diffusion must keep sleep → wake → generate (no resume).""" + + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)]) + omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)]) + + await omni.sleep(level=1) + assert omni._paused is True + + async def wait_for_admission() -> None: + async with omni._pause_cond: + await omni._pause_cond.wait_for(lambda: not omni._paused) + + waiter = asyncio.create_task(wait_for_admission()) + await omni.wake_up() + await asyncio.wait_for(waiter, timeout=1.0) + assert omni._paused is False + assert omni._hold_admission_until_resume is False + assert not omni._sleeping_tags + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_pause_then_sleep_wake_keeps_admission_paused_for_diffusion(): + """Explicit pause_generation still requires resume after diffusion wake.""" + + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_mm_cache = AsyncMock() + omni.reset_encoder_cache = AsyncMock() + omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)]) + omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)]) + + await omni.pause_generation() + await omni.sleep(level=1) + await omni.wake_up() + + assert omni._paused is True + assert omni._hold_admission_until_resume is True + await omni.resume_generation() + assert omni._paused is False + + asyncio.run(run()) diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py index 0377e1a9a..e71a13b1f 100644 --- a/tests/entrypoints/test_omni_sleep_mode.py +++ b/tests/entrypoints/test_omni_sleep_mode.py @@ -114,6 +114,10 @@ async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None: await engine.wake_up(stage_ids=stage_ids) except Exception as e: logger.warning("ensure_awake failed (stage_ids=%s): %s", stage_ids, e) + try: + await engine.resume_generation(stage_ids=stage_ids) + except Exception as e: + logger.warning("ensure_resume failed (stage_ids=%s): %s", stage_ids, e) def _build_llm_stages() -> tuple[list[dict], list[dict]]: @@ -240,6 +244,9 @@ class TestOmniLlmSleepMode: try: await llm_engine.sleep(stage_ids=[0], level=1) await llm_engine.wake_up(stage_ids=[0], tags=["weights"]) + # wake_up restores weights but keeps the AR admission hold. + # Resume so generate() can reach the sleeping-tags rejection. + await llm_engine.resume_generation(stage_ids=[0]) with pytest.raises(RuntimeError, match="partially or fully asleep"): async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)): pass @@ -490,6 +497,7 @@ async def test_multistage_llm_diffusion_sleep_wake(tp_size: int): assert len(acks) == 2 await engine.wake_up(stage_ids=[0, 1]) + await engine.resume_generation(stage_ids=[0, 1]) async for _ in engine.generate("verify", sampling_params_list=[SamplingParams(), sp]): pass finally: diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index af366b2c5..9d7c66e1d 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -50,6 +50,7 @@ from vllm_omni.engine.async_engine_utils import ( ) from vllm_omni.engine.messages import ( AbortRequestMessage, + AbortResultMessage, AddCompanionRequestMessage, CollectiveRPCRequestMessage, CollectiveRPCResultMessage, @@ -1736,14 +1737,51 @@ class AsyncOmniEngine: return self.stage_metadata[stage_id] def abort(self, request_ids: list[str]) -> None: - """Send abort message to the Orchestrator.""" + """Fire-and-forget abort: enqueue and return without waiting. + + Prefer :meth:`abort_async` when the caller needs acknowledgment that + stage aborts, binding release, and orchestrator request cleanup finished. + """ if self.request_queue is None: raise RuntimeError("request_queue is not initialized") self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) - async def abort_async(self, request_ids: list[str]) -> None: - """Async abort API.""" - self.abort(request_ids) + async def abort_async( + self, + request_ids: list[str], + timeout: float | None = None, + ) -> None: + """Abort requests and wait for orchestrator acknowledgment. + + Unlike :meth:`abort`, this generates an ``rpc_id``, correlates the + :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and + raises if the orchestrator reports failure or times out. + """ + if self.request_queue is None: + raise RuntimeError("request_queue is not initialized") + transport = self._correlated_rpc_client + if transport is None: + raise RuntimeError("correlated RPC client is not initialized") + + rpc_id = uuid.uuid4().hex + msg = AbortRequestMessage(request_ids=request_ids, rpc_id=rpc_id) + + def _wait() -> AbortResultMessage: + result_msg = transport.execute( + ("abort", rpc_id), + msg, + timeout=timeout, + timeout_message=f"abort timed out after {timeout} seconds", + block_on_submit=True, + ) + if not isinstance(result_msg, AbortResultMessage): + raise RuntimeError(f"unexpected abort result type: {type(result_msg).__name__}") + return result_msg + + loop = asyncio.get_running_loop() + result_msg = await loop.run_in_executor(None, _wait) + if not result_msg.success: + raise RuntimeError(result_msg.error or "abort failed") def submit_interaction( self, diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 80bf42b0e..65ec53b10 100644 --- a/vllm_omni/engine/messages.py +++ b/vllm_omni/engine/messages.py @@ -42,6 +42,20 @@ class AddCompanionRequestMessage(EngineQueueMessage, kw_only=True): class AbortRequestMessage(EngineQueueMessage, kw_only=True): type: Literal["abort"] = "abort" request_ids: list[str] + # When set, Orchestrator emits AbortResultMessage on rpc_async_queue so + # callers can await acknowledgment. Omit for fire-and-forget abort(). + rpc_id: str | None = None + + +class AbortResultMessage(EngineQueueMessage, kw_only=True): + type: Literal["abort_result"] = "abort_result" + rpc_id: str + success: bool + error: str | None = None + + @property + def rpc_correlation_key(self) -> tuple[str, str]: + return ("abort", self.rpc_id) class InteractionMessage(EngineQueueMessage, kw_only=True): diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 0fb3fb14c..f3f47263b 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -36,6 +36,7 @@ from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.membership_controller import MembershipController from vllm_omni.engine.messages import ( AbortRequestMessage, + AbortResultMessage, AddCompanionRequestMessage, CollectiveRPCRequestMessage, CollectiveRPCResultMessage, @@ -817,13 +818,35 @@ class Orchestrator: ) async def _handle_abort(self, msg: AbortRequestMessage) -> None: - """Handle an abort message from the main thread.""" + """Handle an abort message from the main thread. + + When ``msg.rpc_id`` is set, emit :class:`AbortResultMessage` after + stage aborts, binding release, and request cleanup complete so the + caller can await acknowledgment. Fire-and-forget aborts omit rpc_id. + """ request_ids = msg.request_ids - # _cleanup_request_ids is CFG-aware: it expands aborted parents to - # their companions and fails a deferred parent whose companion is - # aborted before its output arrived. - await self._cleanup_request_ids(list(request_ids), abort=True) - logger.info("[Orchestrator] Aborted request(s) %s", request_ids) + error: str | None = None + try: + # _cleanup_request_ids is CFG-aware: it expands aborted parents to + # their companions and fails a deferred parent whose companion is + # aborted before its output arrived. + await self._cleanup_request_ids(list(request_ids), abort=True) + logger.info("[Orchestrator] Aborted request(s) %s", request_ids) + except Exception as exc: + error = str(exc) + logger.exception("[Orchestrator] Abort failed for request(s) %s", request_ids) + if msg.rpc_id is None: + raise + if msg.rpc_id is not None: + # Always emit a result when rpc_id is set so CorrelatedRpcClient + # waiters are unblocked even on failure. + await self.rpc_async_queue.put( + AbortResultMessage( + rpc_id=msg.rpc_id, + success=error is None, + error=error, + ) + ) async def _handle_interaction(self, msg: InteractionMessage) -> None: """Handle a midway interaction for an active streaming diffusion request.""" diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 7d1bfc60a..d30f88a23 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -94,6 +94,16 @@ class StagePool: DISPATCH_WAIT_TIMEOUT_S: float = 10.0 DISPATCH_RETRY_INTERVAL_S: float = 0.1 + # Only these EngineCore helpers may skip collective_rpc_async. A generic + # ``{method}_async`` on AsyncMPClient must not silently drop timeout. + _ENGINE_CORE_CONTROL_ASYNC_METHODS = frozenset( + { + "pause_scheduler", + "resume_scheduler", + "sleep", + "wake_up", + } + ) def __init__( self, @@ -1214,6 +1224,7 @@ class StagePool: kwargs: dict[str, Any] | None = None, ) -> dict[str, Any] | Any: """Dispatch a stage-scoped control-plane RPC to one physical route.""" + args = tuple(args or ()) kwargs = dict(kwargs or {}) client = self.clients[replica_id] if client is None: @@ -1222,6 +1233,14 @@ class StagePool: "error": f"stage {self.stage_id} replica {replica_id} is not attached", } try: + if self.stage_type != "diffusion" and method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS: + client_method = getattr(client, f"{method}_async", None) + if callable(client_method): + result = client_method(*args, **kwargs) + if timeout is not None: + return await asyncio.wait_for(result, timeout=timeout) + return await result + return await client.collective_rpc_async( method=method, timeout=timeout, diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 33905b254..da0169467 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -146,6 +146,11 @@ class AsyncOmni(EngineClient, OmniBase): OmniBase.__init__(self, model=model, **kwargs) self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False + # True after pause_generation() or AR EngineCore sleep; wake_up must + # not reopen generate() until resume_generation(). Diffusion-only + # sleep uses _paused as a temporary admission gate and clears it + # on wake so sleep → wake → generate keeps working. + self._hold_admission_until_resume: bool = False self._sleeping_tags: set[str] = set() self._level2_sleeping: bool = False self._duplex_request_client: DuplexRequestClient | None = None @@ -1009,6 +1014,30 @@ class AsyncOmni(EngineClient, OmniBase): return results + async def _engine_core_rpc( + self, + method: str, + *, + stage_ids: list[int], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> list[Any]: + """Call an AR EngineCore helper via collective_rpc (orchestrator loop). + + StagePool resolves ``{method}_async`` on the AR client when present + (vLLM AsyncMPClient convention). Raises if any replica reports failure. + """ + results = await self.collective_rpc( + method=method, + args=args, + kwargs=kwargs, + stage_ids=stage_ids, + ) + for result in results: + if isinstance(result, dict) and result.get("error"): + raise RuntimeError(f"{method} failed: {result['error']}") + return results + @staticmethod def _coerce_stage_bool(result: Any) -> bool: """Reduce a stage RPC result to a boolean. @@ -1079,7 +1108,12 @@ class AsyncOmni(EngineClient, OmniBase): await self._abort(internal_req_ids) async def _abort(self, request_ids: list[str]) -> None: - """Submit request IDs to be aborted to the engine.""" + """Abort request IDs via the engine and clean frontend state after ack. + + Waits for orchestrator abort acknowledgment before popping + ``request_states`` so generate() cleanup stays consistent with + backend binding/request teardown. Orchestrator abort errors propagate. + """ await self.engine.abort_async(request_ids) for rid in request_ids: state = self.request_states.pop(rid, None) @@ -1089,23 +1123,77 @@ class AsyncOmni(EngineClient, OmniBase): if self.log_stats: logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) + def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]: + """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC).""" + n_stages = len(self.engine.stage_clients) + if stage_ids is None: + stage_ids = list(range(n_stages)) + else: + invalid = [sid for sid in stage_ids if not isinstance(sid, int) or sid < 0 or sid >= n_stages] + if invalid: + raise ValueError( + f"Invalid stage_ids {invalid}; valid range is 0..{n_stages - 1}" + if n_stages + else f"Invalid stage_ids {invalid}; this engine has no stages" + ) + ar_stage_ids: list[int] = [] + diffusion_stage_ids: list[int] = [] + for sid in stage_ids: + client = self.engine.stage_clients[sid] + if getattr(client, "stage_type", "llm") == "diffusion": + diffusion_stage_ids.append(sid) + else: + ar_stage_ids.append(sid) + return ar_stage_ids, diffusion_stage_ids + async def pause_generation( self, *, mode: PauseMode = "abort", wait_for_inflight_requests: bool = False, clear_cache: bool = True, + stage_ids: list[int] | None = None, ) -> None: - """Pause generation.""" + """Pause generation, mirroring vLLM AsyncLLM.pause_generation. + + 1. Stop frontend admission (``_paused``). + 2. For AR/LLM stages, call EngineCore.pause_scheduler via the + Orchestrator loop (abort/wait/keep + optional cache clear). + 3. Diffusion stages have no EngineCore scheduler — only frontend + admission is paused for them. + + Note: ``sleep()`` already pauses the AR scheduler internally (same as + vLLM EngineCore.sleep). Call this API when you need pause *without* + freeing GPU memory (e.g. weight sync). + """ + if wait_for_inflight_requests: + mode = "wait" + async with self._pause_cond: - if self._paused: - return + # Keep running EngineCore pause + cache clear even when frontend + # admission is already paused (sleep or a prior pause_generation). self._paused = True + self._hold_admission_until_resume = True - # TODO: Implement request draining if wait_for_inflight_requests + ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + if ar_stage_ids: + logger.info( + "[%s] Pausing AR stage(s) %s via EngineCore.pause_scheduler(mode=%s)", + self._name, + ar_stage_ids, + mode, + ) + # Same API name as vLLM AsyncMPClient.pause_scheduler_async; routed + # through collective_rpc so it runs on the orchestrator event loop. + await self._engine_core_rpc( + "pause_scheduler", + stage_ids=ar_stage_ids, + kwargs={"mode": mode, "clear_cache": clear_cache}, + ) + # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler + # already clears AR-side caches when clear_cache=True. if clear_cache: - # Clear caches for all stages. await self.reset_prefix_cache( reset_running_requests=not wait_for_inflight_requests, reset_connector=True, @@ -1113,14 +1201,20 @@ class AsyncOmni(EngineClient, OmniBase): await self.reset_mm_cache() await self.reset_encoder_cache() - async def resume_generation(self) -> None: - """Resume generation.""" + async def resume_generation(self, stage_ids: list[int] | None = None) -> None: + """Resume generation after :meth:`pause_generation`.""" + ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + if ar_stage_ids: + logger.info("[%s] Resuming AR stage(s) %s via EngineCore", self._name, ar_stage_ids) + await self._engine_core_rpc("resume_scheduler", stage_ids=ar_stage_ids) + async with self._pause_cond: self._paused = False + self._hold_admission_until_resume = False self._pause_cond.notify_all() async def is_paused(self) -> bool: - """Check if paused.""" + """Check if frontend admission is paused.""" async with self._pause_cond: return self._paused @@ -1178,33 +1272,65 @@ class AsyncOmni(EngineClient, OmniBase): async def sleep( self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort" ) -> list[OmniACK]: + """Put stages to sleep. + + AR/LLM stages use EngineCore.sleep (pause scheduler, wait idle, then + offload/discard memory) — matching vLLM AsyncLLM.sleep. + + Diffusion stages keep the existing worker-level handle_sleep_task RPC + because StageDiffusionProc does not expose EngineCore.pause_scheduler. + + Frontend admission is blocked at the start of this call (``_paused``) + so pipelined :meth:`generate` cannot race into stages while sleep is + in flight. This does **not** invoke EngineCore.pause_scheduler again + (sleep already pauses the AR scheduler). + + For AR / mixed engines, ``wake_up`` does **not** clear ``_paused``; + callers must :meth:`resume_generation` when ready (typical trainer + order: pause → abort → sleep → train → wake → resume). Diffusion-only + engines have no EngineCore pause to hold, so ``wake_up`` restores + admission and ``sleep → wake → generate`` keeps working. + """ + # Block admission before any sleep RPC so generate() waits on + # _pause_cond during the drain/offload window. EngineCore.sleep will + # pause the AR scheduler itself; avoid a second pause_scheduler here. + async with self._pause_cond: + self._paused = True + self._final_output_handler() - if stage_ids is None: - stage_ids = list(range(len(self.engine.stage_clients))) - total_workers = 0 - for sid in stage_ids: - client = self.engine.stage_clients[sid] - # During the Diffusion phase, regardless of the TP amount, - # currently only a summary ACK is reported at Rank 0. - if getattr(client, "stage_type", "") == "diffusion": - total_workers += 1 - else: - config = self.engine.stage_vllm_configs[sid] - actual_tp = config.parallel_config.tensor_parallel_size if config else 1 - total_workers += actual_tp + ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + final_acks: list[OmniACK] = [] + if ar_stage_ids: + self._hold_admission_until_resume = True + logger.info( + "[%s] Sleeping AR stage(s) %s via EngineCore.sleep(level=%s, mode=%s)", + self._name, + ar_stage_ids, + level, + mode, + ) + await self._engine_core_rpc( + "sleep", + stage_ids=ar_stage_ids, + args=(level, mode), + ) + # EngineCore.sleep has no OmniACK handshake; emit stage-level SUCCESS + # markers so callers/tests that count ACKs keep a stable API. + task_id = f"engine_core-sleep-{uuid.uuid4().hex[:8]}" + final_acks.extend( + OmniACK( + task_id=task_id, + status="SUCCESS", + stage_id=sid, + rank=0, + metadata={"path": "engine_core", "level": level, "mode": mode}, + ) + for sid in ar_stage_ids + ) + + if diffusion_stage_ids: + final_acks.extend(await self._sleep_diffusion(diffusion_stage_ids, level)) - task_id = str(uuid.uuid4()) - self.event_resolver.watch_task(task_id, expected_count=total_workers) - logger.info(f"[{self._name}] Sleep initiated (Task: {task_id}). Awaiting {total_workers} ACKs...") - task = OmniSleepTask(level=level, task_id=task_id) - rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids) - final_acks = [] - 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) if not hasattr(self, "_sleeping_tags"): self._sleeping_tags = set() self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value]) @@ -1212,7 +1338,36 @@ class AsyncOmni(EngineClient, OmniBase): self._level2_sleeping = True return final_acks + async def _sleep_diffusion(self, stage_ids: list[int], level: int) -> list[OmniACK]: + """Worker-level sleep RPC for diffusion stages only.""" + # Diffusion reports one summary ACK at rank 0 regardless of TP. + total_workers = len(stage_ids) + task_id = str(uuid.uuid4()) + self.event_resolver.watch_task(task_id, expected_count=total_workers) + 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) + 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) + return final_acks + async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]: + """Wake stages after sleep. + + AR/LLM stages use EngineCore.wake_up (restore memory, auto-resume + scheduler). Diffusion stages keep the worker-level wake RPC. + + Does **not** clear the frontend ``_paused`` admission gate when + :meth:`pause_generation` ran or AR stages were slept — call + :meth:`resume_generation` when the trainer is ready to admit new + requests. Diffusion-only ``sleep`` uses ``_paused`` only as a race + guard; this method restores admission after a successful wake. + """ self._final_output_handler() if getattr(self, "_level2_sleeping", False): @@ -1231,30 +1386,30 @@ class AsyncOmni(EngineClient, OmniBase): logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.") return [] - if stage_ids is None: - stage_ids = list(range(len(self.engine.stage_clients))) - total_workers = 0 - for sid in stage_ids: - client = self.engine.stage_clients[sid] - if getattr(client, "stage_type", "") == "diffusion": - total_workers += 1 - else: - config = self.engine.stage_vllm_configs[sid] - total_workers += config.parallel_config.tensor_parallel_size if config else 1 - task_id = str(uuid.uuid4()) - self.event_resolver.watch_task(task_id, expected_count=total_workers) - logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...") - 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 = [] - 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) - current_omni_platform.synchronize() - await asyncio.sleep(0.1) + ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + + final_acks: list[OmniACK] = [] + if ar_stage_ids: + logger.info("[%s] Waking AR stage(s) %s via EngineCore", self._name, ar_stage_ids) + await self._engine_core_rpc( + "wake_up", + stage_ids=ar_stage_ids, + kwargs={"tags": requested_tags}, + ) + task_id = f"engine_core-wake-{uuid.uuid4().hex[:8]}" + final_acks.extend( + OmniACK( + task_id=task_id, + status="SUCCESS", + stage_id=sid, + rank=0, + metadata={"path": "engine_core", "tags": list(requested_tags)}, + ) + for sid in ar_stage_ids + ) + + if diffusion_stage_ids: + final_acks.extend(await self._wake_diffusion(diffusion_stage_ids, requested_tags)) for t in requested_tags: if hasattr(self, "_sleeping_tags"): @@ -1263,7 +1418,37 @@ class AsyncOmni(EngineClient, OmniBase): # wake support (e.g. tags=["kv_cache"] only) is added in the future. if not getattr(self, "_sleeping_tags", None): self._level2_sleeping = False - logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.") + logger.info( + "[%s] Wake-up complete for stage(s) %s.", + self._name, + ar_stage_ids + diffusion_stage_ids, + ) + # Diffusion-only sleep uses `_paused` as a race guard. Restore + # generate() admission after memory is back. AR/mixed sleep and + # pause_generation keep the trainer hold until resume_generation. + if not getattr(self, "_hold_admission_until_resume", False): + async with self._pause_cond: + self._paused = False + self._pause_cond.notify_all() + return final_acks + + async def _wake_diffusion(self, stage_ids: list[int], requested_tags: list[str]) -> list[OmniACK]: + """Worker-level wake RPC for diffusion stages only.""" + total_workers = len(stage_ids) + task_id = str(uuid.uuid4()) + self.event_resolver.watch_task(task_id, expected_count=total_workers) + 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) + current_omni_platform.synchronize() + await asyncio.sleep(0.1) return final_acks async def is_sleeping(self) -> bool: