[Bugfix] Fix colocate-async sleep admission race and deliver AR abort tokens (#6367)

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

tests/entrypoints/test_async_omni.py

@@ -180,7 +180,6 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
             assert re.fullmatch(rf"{re.escape(cancel_prefix)}-[0-9a-f]+", rid)
         assert len(aborted_ids) == expected_cancel_count
         assert len(set(aborted_ids)) == expected_cancel_count
-        assert len(omni.request_states) == len(req_ids) - expected_cancel_count
 
         for t in tasks:
             t.cancel()
@@ -190,7 +189,9 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
 
 
 @pytest.mark.cpu
-def test_abort_pops_request_states_only_after_ack():
+def test_abort_keeps_request_states_until_generate_cleanup():
+    """Frontend abort must not drop request_states; generate() owns cleanup."""
+
     async def run():
         release = asyncio.Event()
         seen_during_wait: list[int] = []
@@ -215,7 +216,7 @@ def test_abort_pops_request_states_only_after_ack():
 
         release.set()
         await task
-        assert "req-1-aaaa" not in omni.request_states
+        assert "req-1-aaaa" in omni.request_states
 
     asyncio.run(run())
 
@@ -241,6 +242,93 @@ def test_abort_propagates_engine_errors_without_popping_state():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_abort_enqueues_prefix_tokens_from_engine():
+    """AR abort must deliver the generated prefix on the request queue (CPU)."""
+
+    async def run():
+        from vllm.outputs import CompletionOutput
+
+        from vllm_omni.engine.messages import OutputMessage
+
+        queue: asyncio.Queue = asyncio.Queue()
+
+        async def abort_with_prefix(request_ids):
+            rid = request_ids[0]
+            engine_output = OmniRequestOutput(
+                request_id=rid,
+                finished=True,
+                stage_id=0,
+                final_output_type="text",
+                outputs=[
+                    CompletionOutput(
+                        index=0,
+                        text="partial",
+                        token_ids=[7, 8, 9],
+                        cumulative_logprob=None,
+                        logprobs=None,
+                        finish_reason="abort",
+                        stop_reason=None,
+                    )
+                ],
+            )
+            return [
+                OutputMessage(
+                    request_id=rid,
+                    stage_id=0,
+                    replica_id=None,
+                    engine_outputs=engine_output,
+                    metrics=None,
+                    finished=True,
+                    stage_submit_ts=None,
+                )
+            ]
+
+        omni = get_async_omni_instance(fake_abort_request=abort_with_prefix)
+        omni.request_states["req-1-cccc"] = SimpleNamespace(
+            request_id="req-1-cccc",
+            external_request_id="req-1",
+            input_stream_task=None,
+            queue=queue,
+        )
+        await omni._abort(["req-1-cccc"])
+        msg = queue.get_nowait()
+        assert msg.finished is True
+        assert list(msg.engine_outputs.outputs[0].token_ids) == [7, 8, 9]
+        assert msg.engine_outputs.outputs[0].finish_reason == "abort"
+        assert "req-1-cccc" in omni.request_states
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_enqueues_synthetic_finished_when_engine_returns_empty():
+    """Empty abort_async must still unblock generate() with finish_reason=abort."""
+
+    async def run():
+        queue: asyncio.Queue = asyncio.Queue()
+
+        async def empty_abort_async(request_ids):
+            del request_ids
+            return []
+
+        omni = get_async_omni_instance(fake_abort_request=empty_abort_async)
+        omni.request_states["req-1-dddd"] = SimpleNamespace(
+            request_id="req-1-dddd",
+            external_request_id="req-1",
+            input_stream_task=None,
+            queue=queue,
+        )
+        await omni._abort(["req-1-dddd"])
+        msg = queue.get_nowait()
+        assert msg.finished is True
+        assert msg.request_id == "req-1-dddd"
+        assert msg.engine_outputs.outputs[0].finish_reason == "abort"
+        assert "req-1-dddd" in omni.request_states
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_generate_accepts_request_after_repeated_cancellations():
     async def run_test():

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -18,11 +18,14 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
     omni = object.__new__(AsyncOmni)
     omni._pause_cond = asyncio.Condition()
     omni._paused = False
+    omni._admitting = 0
     omni._hold_admission_until_resume = False
     omni._sleeping_tags = set()
+    omni._stage_sleeping_tags = {}
     omni._level2_sleeping = False
     omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock())
     omni._final_output_handler = lambda: None
+    omni.reset_mm_cache = AsyncMock()
 
     stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types]
     omni.engine = SimpleNamespace(
@@ -171,6 +174,33 @@ def test_sleep_blocks_admission_before_engine_core_rpc():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_sleep_waits_for_in_flight_generate_admission():
+    """Sleep must not offload while generate() is still in add_request."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        omni._admitting = 1
+        rpc_started = asyncio.Event()
+
+        async def rpc_side_effect(**kwargs):
+            if kwargs.get("method") == "sleep":
+                rpc_started.set()
+            return [True]
+
+        omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+        sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
+        await asyncio.sleep(0.05)
+        assert not sleep_task.done()
+        assert not rpc_started.is_set()
+        await omni._release_generate_admission()
+        await asyncio.wait_for(sleep_task, timeout=1)
+        assert rpc_started.is_set()
+        assert omni._admitting == 0
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
     async def run() -> None:
@@ -260,6 +290,80 @@ def test_sleep_diffusion_only_skips_engine_core_collective_rpc():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_streaming_wait_for_first_chunk_does_not_block_sleep():
+    """Waiting for the next client chunk must not hold an admission slot."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        rpc_started = asyncio.Event()
+        first_chunk = asyncio.Event()
+
+        async def rpc_side_effect(**kwargs):
+            if kwargs.get("method") == "sleep":
+                rpc_started.set()
+            return [True]
+
+        omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+
+        async def wait_then_submit() -> None:
+            await first_chunk.wait()
+            await omni._submit_with_admission(asyncio.sleep(0))
+
+        wait_task = asyncio.create_task(wait_then_submit())
+        sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
+        await asyncio.sleep(0.05)
+        assert not wait_task.done()
+        assert omni._admitting == 0
+        await asyncio.wait_for(sleep_task, timeout=1)
+        assert rpc_started.is_set()
+
+        wait_task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await wait_task
+        assert omni._admitting == 0
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_streaming_in_flight_add_blocks_sleep():
+    """Sleep must wait while a streaming ADD/update holds an admission slot."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        add_started = asyncio.Event()
+        add_release = asyncio.Event()
+        rpc_started = asyncio.Event()
+
+        async def rpc_side_effect(**kwargs):
+            if kwargs.get("method") == "sleep":
+                rpc_started.set()
+            return [True]
+
+        omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+
+        async def slow_add() -> None:
+            add_started.set()
+            await add_release.wait()
+
+        add_task = asyncio.create_task(omni._submit_with_admission(slow_add()))
+        await add_started.wait()
+        sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
+        await asyncio.sleep(0.05)
+        assert not sleep_task.done()
+        assert not rpc_started.is_set()
+        assert omni._admitting == 1
+
+        add_release.set()
+        await asyncio.wait_for(add_task, timeout=1)
+        await asyncio.wait_for(sleep_task, timeout=1)
+        assert rpc_started.is_set()
+        assert omni._admitting == 0
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_wake_up_restores_admission_for_diffusion_only():
     """Pure diffusion must keep sleep → wake → generate (no resume)."""
@@ -308,3 +412,30 @@ def test_pause_then_sleep_wake_keeps_admission_paused_for_diffusion():
         assert omni._paused is False
 
     asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_partial_wake_does_not_skip_remaining_sleeping_stage():
+    """sleep(stage_ids=[0]) then wake([0]) must not skip a later wake([1])."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "llm"])
+        await omni.sleep(stage_ids=[0, 1], level=1, mode="abort")
+        assert omni._stage_sleeping_tags.keys() == {0, 1}
+
+        await omni.wake_up(stage_ids=[0])
+        assert 0 not in omni._stage_sleeping_tags
+        assert 1 in omni._stage_sleeping_tags
+        assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
+
+        omni.collective_rpc.reset_mock()
+        acks = await omni.wake_up(stage_ids=[1])
+        assert acks
+        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"] == [1]
+        assert not omni._sleeping_tags
+        assert not omni._stage_sleeping_tags
+
+    asyncio.run(run())

tests/entrypoints/test_omni_sleep_mode.py

@@ -109,7 +109,7 @@ def get_dynamic_devices(stage_idx: int, num_stages: int, tp_size: int) -> str:
 
 
 async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None:
-    """Best-effort full wake so the next shared-engine test starts active."""
+    """Best-effort full wake + resume so the next shared-engine test starts active."""
     try:
         await engine.wake_up(stage_ids=stage_ids)
     except Exception as e:
@@ -532,6 +532,7 @@ async def test_pure_diffusion_sleep_wake():
     try:
         await engine.sleep(level=1)
         await engine.wake_up()
+        await engine.resume_generation()
         async for _ in engine.generate(
             "test",
             sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256),

tests/entrypoints/test_stage_utils.py

@@ -216,6 +216,20 @@ def test_resolve_stage_physical_devices_uses_visible_baseline(monkeypatch: pytes
     assert resolve_stage_physical_devices(1, "1", visible_baseline="3,4,5") == "4"
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_resolve_stage_physical_devices_captured_unset_ignores_live_env(monkeypatch: pytest.MonkeyPatch):
+    """Parallel init captures an unset CUDA_VISIBLE_DEVICES as None.
+
+    A sibling thread may then narrow the live env for spawn. The captured
+    None must not fall back to that live value, or logical device 0 would
+    map onto the other stage's GPU.
+    """
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7")
+    assert resolve_stage_physical_devices(0, "0", visible_baseline=None) == "0"
+    assert resolve_stage_physical_devices(1, "1", visible_baseline=None) == "1"
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_map_device_list_raises_on_non_numeric():