[Core][Diffusion] Add pause_generation(mode="keep") for diffusion stages (#7685)

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

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -25,6 +25,7 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
     omni._paused = False
     omni._admitting = 0
     omni._hold_admission_until_resume = False
+    omni._paused_stage_ids = set()
     omni._sleeping_tags = set()
     omni._stage_sleeping_tags = {}
     omni._level2_sleeping = False
@@ -460,3 +461,225 @@ def test_partial_wake_does_not_skip_remaining_sleeping_stage():
         assert not omni._stage_sleeping_tags
 
     asyncio.run(run())
+
+
+# ───────────────────── diffusion pause_generation(mode="keep") ─────────────────────
+
+
+def _rpc_methods(omni: AsyncOmni) -> list[tuple[str, list[int] | None]]:
+    return [(call.kwargs["method"], call.kwargs.get("stage_ids")) for call in omni.collective_rpc.await_args_list]
+
+
+@pytest.mark.cpu
+def test_diffusion_keep_routes_to_backend():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+
+        await omni.pause_generation(mode="keep", clear_cache=False)
+
+        omni.collective_rpc.assert_awaited_once_with(
+            method="pause_scheduler",
+            args=(),
+            kwargs={"mode": "keep"},
+            stage_ids=[0],
+        )
+        assert await omni.is_paused() is True
+        assert omni._hold_admission_until_resume is True
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "pause_kwargs",
+    [
+        {},
+        {"mode": "abort"},
+        {"mode": "wait"},
+        {"mode": "keep", "wait_for_inflight_requests": True},
+    ],
+)
+def test_diffusion_non_keep_modes_pause_frontend_only(pause_kwargs):
+    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()
+
+        await omni.pause_generation(**pause_kwargs)
+        assert await omni.is_paused() is True
+        omni.collective_rpc.assert_not_awaited()
+
+        await omni.resume_generation()
+        omni.collective_rpc.assert_not_awaited()
+        assert await omni.is_paused() is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_mixed_engine_keep_sends_engine_core_and_diffusion_pauses():
+    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="keep", clear_cache=True)
+
+        assert omni.collective_rpc.await_args_list[0].kwargs == {
+            "method": "pause_scheduler",
+            "args": (),
+            "kwargs": {"mode": "keep", "clear_cache": True},
+            "stage_ids": [0],
+        }
+        assert omni.collective_rpc.await_args_list[1].kwargs == {
+            "method": "pause_scheduler",
+            "args": (),
+            "kwargs": {"mode": "keep"},
+            "stage_ids": [1],
+        }
+        # The frontend cache reset still runs after the backend pauses.
+        omni.reset_prefix_cache.assert_awaited_once()
+
+        await omni.resume_generation()
+        assert _rpc_methods(omni)[2:] == [("resume_scheduler", [0]), ("resume_scheduler", [1])]
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_keep_targets_only_requested_diffusion_stages():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion", "diffusion"])
+
+        await omni.pause_generation(mode="keep", clear_cache=False, stage_ids=[1])
+        assert _rpc_methods(omni) == [("pause_scheduler", [1])]
+
+        await omni.resume_generation(stage_ids=[0])
+        assert _rpc_methods(omni) == [("pause_scheduler", [1])]
+        # Stage 1's scheduler is still closed: admission must stay closed too,
+        # otherwise new requests queue on a stage that will not schedule them.
+        assert await omni.is_paused() is True
+
+        await omni.resume_generation(stage_ids=[1])
+        assert _rpc_methods(omni) == [("pause_scheduler", [1]), ("resume_scheduler", [1])]
+        assert await omni.is_paused() is False
+
+        await omni.resume_generation()
+        assert len(_rpc_methods(omni)) == 2
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_keep_then_sleep_wake_resume_keeps_admission_paused_until_resume():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.reset_prefix_cache = AsyncMock(return_value=True)
+        omni.reset_encoder_cache = AsyncMock()
+        omni.collective_rpc = AsyncMock(
+            side_effect=lambda **kw: [OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)]
+            if kw["method"] in {"handle_sleep_task", "handle_wake_task"}
+            else [None]
+        )
+
+        await omni.pause_generation(mode="keep", clear_cache=False)
+        await omni.sleep(level=1)
+        await omni.wake_up()
+        assert await omni.is_paused() is True
+
+        await omni.resume_generation()
+        assert await omni.is_paused() is False
+        assert [m for m, _ in _rpc_methods(omni)] == [
+            "pause_scheduler",
+            "handle_sleep_task",
+            "handle_wake_task",
+            "resume_scheduler",
+        ]
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "failure",
+    [
+        AsyncMock(side_effect=NotImplementedError("request-level execution only")),
+        AsyncMock(return_value=[{"supported": False, "error": "TimeoutError: outputs did not drain"}]),
+    ],
+)
+def test_keep_failure_surfaces_and_leaves_explicit_resume_possible(failure):
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.collective_rpc = failure
+
+        with pytest.raises((NotImplementedError, RuntimeError)):
+            await omni.pause_generation(mode="keep", clear_cache=False)
+
+        assert await omni.is_paused() is True
+        assert omni._paused_stage_ids == {0}
+
+        omni.collective_rpc = AsyncMock(return_value=[None])
+        await omni.resume_generation()
+        omni.collective_rpc.assert_awaited_once_with(
+            method="resume_scheduler",
+            args=(),
+            kwargs=None,
+            stage_ids=[0],
+        )
+        assert await omni.is_paused() is False
+        assert omni._paused_stage_ids == set()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_resume_failure_keeps_admission_paused_and_target_recorded():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        await omni.pause_generation(mode="keep", clear_cache=False)
+
+        omni.collective_rpc = AsyncMock(return_value=[{"error": "replica gone"}])
+        with pytest.raises(RuntimeError, match="resume_scheduler failed"):
+            await omni.resume_generation()
+
+        assert await omni.is_paused() is True
+        assert omni._paused_stage_ids == {0}
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("first, second", [([1], [0]), ([0], [1])])
+def test_mixed_engine_partial_resume_keeps_admission_paused(first, second):
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        await omni.pause_generation(mode="keep", clear_cache=False)
+
+        await omni.resume_generation(stage_ids=first)
+        assert await omni.is_paused() is True
+
+        await omni.resume_generation(stage_ids=second)
+        assert await omni.is_paused() is False
+        assert _rpc_methods(omni)[2:] == [("resume_scheduler", first), ("resume_scheduler", second)]
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_two_ar_stages_partial_resume_keeps_admission_paused():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "llm"])
+        omni.reset_prefix_cache = AsyncMock(return_value=True)
+        omni.reset_mm_cache = AsyncMock()
+        omni.reset_encoder_cache = AsyncMock()
+        await omni.pause_generation(mode="abort")
+
+        await omni.resume_generation(stage_ids=[0])
+        assert await omni.is_paused() is True
+
+        await omni.resume_generation(stage_ids=[1])
+        assert await omni.is_paused() is False
+
+    asyncio.run(run())