[Feature]: Support pause / resume and sleep / wake for AR stages in AsyncOmni (#6084)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

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

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,
     *,

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)

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