[Core] Cleanup engine pause/sleep logic

tests/v1/distributed/test_async_llm_dp.py

@@ -3,8 +3,10 @@
 
 import asyncio
 import os
+import time
 from contextlib import ExitStack
 from dataclasses import dataclass
+from typing import Any
 
 import pytest
 
@@ -187,24 +189,33 @@ def log_engine_initialized(self):
 # =============================================================================
 # DP Pause/Resume Tests
 # =============================================================================
+# When expert_parallel=False: uses non-MoE model (DP replicas as separate engines).
+# When expert_parallel=True: uses MoE model + EP (DPEngineCoreProc, sync pause path).
 
 DP_PAUSE_MODEL = "hmellor/tiny-random-LlamaForCausalLM"
+DP_PAUSE_MODEL_MOE = "ibm-research/PowerMoE-3b"
 DP_PAUSE_PROMPT = "This is a test of data parallel pause"
 
 
+def _get_dp_pause_engine_args(expert_parallel: bool) -> AsyncEngineArgs:
+    """Engine args for DP pause tests: MoE+EP when expert_parallel else small Llama."""
+    model = DP_PAUSE_MODEL_MOE if expert_parallel else DP_PAUSE_MODEL
+    return AsyncEngineArgs(
+        model=model,
+        enforce_eager=True,
+        tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
+        data_parallel_size=DP_SIZE,
+        data_parallel_backend="mp",
+        enable_expert_parallel=expert_parallel,
+    )
+
+
 @pytest.mark.asyncio
-async def test_dp_pause_resume_basic():
+@pytest.mark.parametrize("expert_parallel", [False, True])
+async def test_dp_pause_resume_basic(expert_parallel: bool):
     """Pausing from the client (one call) pauses all DP ranks; resume clears it."""
-    if current_platform.is_rocm():
-        pytest.skip("DP pause tests use mp backend only")
     with ExitStack() as after:
-        engine_args = AsyncEngineArgs(
-            model=DP_PAUSE_MODEL,
-            enforce_eager=True,
-            tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
-            data_parallel_size=DP_SIZE,
-            data_parallel_backend="mp",
-        )
+        engine_args = _get_dp_pause_engine_args(expert_parallel)
         engine = AsyncLLM.from_engine_args(engine_args)
         after.callback(engine.shutdown)
 
@@ -226,18 +237,11 @@ async def test_dp_pause_resume_basic():
 
 
 @pytest.mark.asyncio
-async def test_dp_pause_abort():
+@pytest.mark.parametrize("expert_parallel", [False, True])
+async def test_dp_pause_abort(expert_parallel: bool):
     """Pause with abort from one client aborts in-flight requests on all DP ranks."""
-    if current_platform.is_rocm():
-        pytest.skip("DP pause tests use mp backend only")
     with ExitStack() as after:
-        engine_args = AsyncEngineArgs(
-            model=DP_PAUSE_MODEL,
-            enforce_eager=True,
-            tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
-            data_parallel_size=DP_SIZE,
-            data_parallel_backend="mp",
-        )
+        engine_args = _get_dp_pause_engine_args(expert_parallel)
         engine = AsyncLLM.from_engine_args(engine_args)
         after.callback(engine.shutdown)
 
@@ -286,41 +290,111 @@ async def gen(rid: str):
 
 
 @pytest.mark.asyncio
-async def test_dp_pause_keep_then_resume():
-    """Pause with keep queues new requests; resume allows them to run."""
-    if current_platform.is_rocm():
-        pytest.skip("DP pause tests use mp backend only")
+@pytest.mark.parametrize("expert_parallel", [False, True])
+async def test_dp_pause_keep_then_resume(expert_parallel: bool):
+    """Start generation, pause after a few tokens (keep mode), resume; verify gap."""
+
+    pause_duration = 2.0
+    min_tokens_before_pause = 3
+
     with ExitStack() as after:
-        engine_args = AsyncEngineArgs(
-            model=DP_PAUSE_MODEL,
-            enforce_eager=True,
-            tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
-            data_parallel_size=DP_SIZE,
-            data_parallel_backend="mp",
-        )
+        engine_args = _get_dp_pause_engine_args(expert_parallel)
         engine = AsyncLLM.from_engine_args(engine_args)
         after.callback(engine.shutdown)
 
-        await engine.pause_generation(mode="keep")
-        assert await engine.is_paused()
-
-        request_done = asyncio.Event()
+        sampling_params = SamplingParams(max_tokens=15, ignore_eos=True)
+        token_times: list[tuple[int, float]] = []
+        pause_token_idx = 0
 
-        async def gen():
-            async for out in engine.generate(
-                request_id="queued-keep",
+        async def generator_task():
+            nonlocal pause_token_idx
+            out = None
+            async for output in engine.generate(
+                request_id="keep-resume-req",
                 prompt=DP_PAUSE_PROMPT,
-                sampling_params=SamplingParams(max_tokens=5),
+                sampling_params=sampling_params,
             ):
-                pass
-            request_done.set()
+                token_count = len(output.outputs[0].token_ids)
+                token_times.append((token_count, time.monotonic()))
+                out = output
             return out
 
-        task = asyncio.create_task(gen())
-        await asyncio.sleep(0.2)
-        assert not request_done.is_set()
+        async def controller_task():
+            nonlocal pause_token_idx
+            while len(token_times) < min_tokens_before_pause:
+                await asyncio.sleep(0.01)
+            await engine.pause_generation(mode="keep")
+            await asyncio.sleep(pause_duration)
+            pause_token_idx = len(token_times)
+            await engine.resume_generation()
+
+        gen_task = asyncio.create_task(generator_task())
+        ctrl_task = asyncio.create_task(controller_task())
+        final_output, _ = await asyncio.gather(gen_task, ctrl_task)
+
+        assert final_output is not None and final_output.finished
+        assert await engine.is_paused() is False
+        assert pause_token_idx >= min_tokens_before_pause
+        if pause_token_idx > 0 and pause_token_idx < len(token_times):
+            pause_gap = (
+                token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1]
+            )
+            assert pause_gap >= pause_duration * 0.8, (
+                f"Expected gap ~{pause_duration}s after pause, got {pause_gap:.3f}s"
+            )
 
+
+@pytest.mark.asyncio
+async def test_dp_pause_keep_race_staggered_engines():
+    """Race: send pause(keep) to engine 0, then add two requests,
+    then pause(keep) to engine 1. Ensures no deadlock when pause
+    requests are staggered and requests arrive in between."""
+    if DP_SIZE != 2:
+        pytest.skip("test_dp_pause_keep_race_staggered_engines requires DP_SIZE=2")
+
+    with ExitStack() as after:
+        engine_args = _get_dp_pause_engine_args(expert_parallel=True)
+        engine = AsyncLLM.from_engine_args(engine_args)
+        after.callback(engine.shutdown)
+
+        client = engine.engine_core
+
+        original_call_utility = client.call_utility_async
+        mid_pause_tasks: list[asyncio.Task] = []
+
+        async def staggered_pause_keep(method: str, *args) -> Any:
+            if method != "pause_scheduler" or not args or args[0] != "keep":
+                return await original_call_utility(method, *args)
+            # Send pause(keep) to engine 0 first
+            await client._call_utility_async(
+                method, *args, engine=client.core_engines[0]
+            )
+            # In the middle: send two requests (race window)
+            sp = SamplingParams(max_tokens=5, ignore_eos=True)
+
+            async def consume_gen(req_id: str) -> None:
+                async for _ in engine.generate(
+                    request_id=req_id,
+                    prompt=DP_PAUSE_PROMPT,
+                    sampling_params=sp,
+                ):
+                    pass
+
+            t1 = asyncio.create_task(consume_gen("race-1"))
+            t2 = asyncio.create_task(consume_gen("race-2"))
+            mid_pause_tasks.extend([t1, t2])
+            await asyncio.sleep(3)
+            # Then send pause(keep) to engine 1
+            result = await client._call_utility_async(
+                method, *args, engine=client.core_engines[1]
+            )
+            return result
+
+        client.call_utility_async = staggered_pause_keep
+
+        await engine.pause_generation(mode="keep")
+        assert await engine.is_paused()
         await engine.resume_generation()
-        final = await asyncio.wait_for(task, timeout=10.0)
-        assert final.finished
         assert not await engine.is_paused()
+        # Let the two requests we sent mid-pause complete
+        await asyncio.gather(*mid_pause_tasks)

tests/v1/engine/test_engine_core_client.py

@@ -280,20 +280,15 @@ def echo_dc_nested(
 
 
 def future_echo(self, value: Any, num_wait_loops: int = 2) -> Future:
-    """Utility that returns a Future completed by a per_step_hook after
-    num_wait_loops engine steps (tests deferred utility path).
+    """Utility that returns a Future completed once the engine is idle
+    (tests deferred utility path).
     """
     future: Future = Future()
-    remaining = [num_wait_loops]
 
-    def _step(engine: EngineCore) -> bool:
-        remaining[0] -= 1
-        if remaining[0] <= 0:
-            future.set_result(value)
-            return True  # remove hook
-        return False
+    def idle(engine: EngineCore):
+        future.set_result(value)
 
-    self.per_step_hooks.add(_step)
+    self._idle_state_callbacks.append(idle)
     return future
 
 
@@ -832,8 +827,8 @@ async def test_engine_core_client_future_utility_async(
     monkeypatch: pytest.MonkeyPatch,
     subprocess_future_echo_patch,
 ):
-    """Test that a utility returning a Future (completed by a per_step_hook
-    after N steps) completes when the future is done (engine uses add_done_callback).
+    """Test that a utility returning a Future completes when the future is done
+    (engine uses add_done_callback).
     """
     with monkeypatch.context() as m:
         m.setattr(EngineCore, "future_echo", future_echo, raising=False)

vllm/engine/protocol.py

@@ -148,7 +148,7 @@ async def reset_prefix_cache(
         ...
 
     @abstractmethod
-    async def sleep(self, level: int = 1) -> None:
+    async def sleep(self, level: int = 1, mode: "PauseMode" = "abort") -> None:
         """Sleep the engine"""
         ...