[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/e2e/offline_inference/test_qwen3_omni.py

@@ -3,6 +3,9 @@
 
 """
 E2E offline tests for Omni model with video input and audio output.
+
+Abort / sleep-admission lives in ``test_qwen3_omni_colocate_async.py`` so it
+does not overlap this module's OmniRunners.
 """
 
 import os

tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Isolated AsyncOmni abort / sleep-admission regression for Qwen3-Omni.
+
+This lives in its own module so the L3 merge job can run it *after*
+``test_qwen3_omni.py`` (offline) and the online serving file. Those modules keep
+module-scoped OmniRunner / OmniServer instances until teardown; overlapping a
+CuMem sleep-mode engine with them kills StageEngineCoreProc at handshake
+(``Failed core proc(s): {}``).
+"""
+
+import asyncio
+import os
+from contextlib import ExitStack
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+from vllm import SamplingParams
+from vllm.inputs import TokensPrompt
+
+from tests.helpers.clean import wait_for_gpu_memory_to_clear
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.platforms import current_omni_platform
+
+MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+
+
+def _output_token_ids(output) -> list[int]:
+    if not getattr(output, "outputs", None):
+        return []
+    completion = output.outputs[0]
+    token_ids = list(completion.token_ids or [])
+    if token_ids:
+        return token_ids
+    return list(getattr(completion, "cumulative_token_ids", None) or [])
+
+
+def _finish_reason(output) -> str | None:
+    if not getattr(output, "outputs", None):
+        return None
+    return getattr(output.outputs[0], "finish_reason", None)
+
+
+def _colocate_async_deploy() -> str:
+    """Thinker-only Instruct deploy with sleep mode on a single GPU.
+
+    Abort / sleep admission is a thinker control-plane contract. The 3-stage
+    CI yaml plus CuMem on every stage fails engine-core startup in the L3
+    merge job: Stage 1 resolves ``devices: "1"`` while only GPU 0 is visible.
+    """
+    return modify_stage_config(
+        get_deploy_config_path("qwen3_omni_moe_thinking.yaml"),
+        updates={
+            "stages": {
+                0: {
+                    "enable_sleep_mode": True,
+                    "enforce_eager": True,
+                    "trust_remote_code": True,
+                    "tensor_parallel_size": 1,
+                    "devices": "0",
+                    "gpu_memory_utilization": 0.9,
+                    "max_num_seqs": 1,
+                    "max_model_len": 8192,
+                    "max_num_batched_tokens": 8192,
+                    "enable_prefix_caching": False,
+                    "skip_mm_profiling": True,
+                },
+            },
+        },
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
+@pytest.mark.asyncio
+async def test_colocate_async_abort_tokens_and_sleep_admission() -> None:
+    """**required** regression. Control-plane APIs (``abort`` / ``sleep`` /
+    ``wake_up`` / ``resume_generation``) are not in ``send_omni_request``.
+
+    On ``origin/main`` this test fails:
+    - abort: ``generate()`` never yields ``finish_reason="abort"`` with the
+      tokens produced so far, so resume has an empty prefix (or the stream
+      hangs after frontend state is dropped).
+    - sleep: ``generate()`` is admitted into EngineCore while weights/KV are
+      offloaded, so the task errors (asleep / corrupted ADD frame) instead of
+      waiting for ``resume_generation()``.
+
+    On this PR both legs pass: abort returns a resumeable prefix, and sleep
+    holds ``generate()`` until wake + resume.
+    """
+    prompt = "What color is the sky? Write a long, detailed explanation."
+    device_count = current_omni_platform.device_count()
+    if device_count > 0:
+        # Fail closed if a prior module-scoped runner is still on the GPU.
+        # ``clean_gpu_memory_between_tests`` only logs a note on timeout.
+        wait_for_gpu_memory_to_clear(
+            devices=list(range(device_count)),
+            threshold_ratio=0.15,
+            timeout_s=120,
+        )
+
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            deploy_config=_colocate_async_deploy(),
+            enable_sleep_mode=True,
+        )
+        after.callback(engine.shutdown)
+
+        request_id = "qwen3-abort-partial"
+        outputs: list = []
+
+        async def _generate(max_tokens: int, req_id: str, gen_prompt) -> None:
+            async for output in engine.generate(
+                prompt=gen_prompt,
+                request_id=req_id,
+                sampling_params=SamplingParams(temperature=0.0, max_tokens=max_tokens),
+                output_modalities=["text"],
+            ):
+                outputs.append(output)
+
+        abort_task = asyncio.create_task(_generate(256, request_id, prompt))
+        prefix: list[int] = []
+        prompt_token_ids: list[int] = []
+        for _ in range(600):
+            if abort_task.done():
+                break
+            if outputs:
+                latest = outputs[-1]
+                prefix = _output_token_ids(latest)
+                prompt_token_ids = list(getattr(latest, "prompt_token_ids", None) or [])
+                if prefix:
+                    break
+            await asyncio.sleep(0.1)
+
+        assert prefix, "generate produced no tokens before abort"
+        assert not abort_task.done(), "generate finished before abort; raise max_tokens"
+        await engine.abort(request_id)
+        await asyncio.wait_for(abort_task, timeout=90)
+
+        final = outputs[-1]
+        assert final.finished
+        assert _finish_reason(final) == "abort"
+        abort_prefix = _output_token_ids(final)
+        assert abort_prefix, "abort dropped the generated prefix (main-branch behavior)"
+
+        resume_prompt: str | TokensPrompt = (
+            TokensPrompt(prompt_token_ids=prompt_token_ids + abort_prefix) if prompt_token_ids else prompt
+        )
+        outputs.clear()
+        async for output in engine.generate(
+            prompt=resume_prompt,
+            request_id="qwen3-abort-resume",
+            sampling_params=SamplingParams(temperature=0.0, max_tokens=8),
+            output_modalities=["text"],
+        ):
+            outputs.append(output)
+        assert outputs
+        assert any(_output_token_ids(out) for out in outputs)
+
+        # Trainer order is pause → abort → sleep. Pause here after the
+        # resume generate so EngineCore is idle before CuMem offload.
+        await engine.pause_generation(mode="abort", clear_cache=True)
+        await engine.sleep(level=1)
+        outputs.clear()
+        sleep_task = asyncio.create_task(_generate(8, "qwen3-sleep-admission", prompt))
+        await asyncio.sleep(1.0)
+        assert not sleep_task.done(), "generate() ran while EngineCore was sleeping (main-branch admission race)"
+        await engine.wake_up()
+        await asyncio.sleep(0.5)
+        assert not sleep_task.done(), "generate() resumed before resume_generation()"
+        await engine.resume_generation()
+        await asyncio.wait_for(sleep_task, timeout=180)
+        assert outputs
+        assert _finish_reason(outputs[-1]) != "abort"
+    await asyncio.sleep(5)

tests/engine/test_async_omni_engine_abort_ack.py

@@ -67,7 +67,8 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture):
         assert not task.done()
 
         rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True))
-        await task
+        result = await task
+        assert result == []
 
     try:
         asyncio.run(_run())

tests/engine/test_orchestrator.py

@@ -256,8 +256,32 @@ class FakeOutputProcessor:
         )
 
     def abort_requests(self, request_ids, internal: bool = False):
-        self.abort_calls.append(request_ids)
-        return request_ids
+        aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal)
+        return aborted_ids
+
+    def abort_requests_collecting_outputs(self, request_ids, *, internal: bool = False, commit_state: bool = True):
+        """Mirror MultimodalOutputProcessor collecting for AR abort-prefix tests."""
+        del internal, commit_state
+        ids = list(request_ids)
+        self.abort_calls.append(ids)
+        outputs: list[RequestOutput] = []
+        for rid in ids:
+            seeded = next(
+                (ro for ro in self.request_outputs if getattr(ro, "request_id", None) == rid),
+                None,
+            )
+            token_ids = list(seeded.outputs[0].token_ids) if seeded is not None and seeded.outputs else [1, 2]
+            # Intentionally stamp an internal-looking id on the RequestOutput so
+            # StagePool must re-key by the orchestrator id it passed in.
+            outputs.append(
+                _build_request_output(
+                    f"engine-internal-{rid}",
+                    token_ids=token_ids,
+                    finished=True,
+                    finish_reason="abort",
+                )
+            )
+        return ids, outputs
 
     def update_scheduler_stats(self, _scheduler_stats) -> None:
         return None
@@ -320,6 +344,7 @@ def _build_request_output(
     prompt_token_ids: list[int] | None = None,
     finished: bool = True,
     text: str = "test",
+    finish_reason: str | None = None,
 ) -> RequestOutput:
     completion = CompletionOutput(
         index=0,
@@ -327,7 +352,7 @@ def _build_request_output(
         token_ids=list(token_ids or [1, 2]),
         cumulative_logprob=0.0,
         logprobs=None,
-        finish_reason="stop" if finished else None,
+        finish_reason=(finish_reason if finish_reason is not None else ("stop" if finished else None)),
         stop_reason=None,
     )
     return RequestOutput(
@@ -983,7 +1008,11 @@ async def test_run_abort(orchestrator_factory) -> None:
 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),
+        FakeStageClient(
+            stage_type="llm",
+            final_output=True,
+            next_inputs=[{"prompt_token_ids": [7, 8, 9]}],
+        ),
     ]
     processors = [
         FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[1], finished=True)]),
@@ -1002,6 +1031,11 @@ async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> N
             final_stage_id=1,
         )
         await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
+        # Abort outputs come from the final AR stage's output processor, and
+        # StagePool.abort_requests only collects for requests with a live
+        # replica binding. Forward off stage-0 first so stage-1 is bound.
+        stages[0].push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0))
+        await _wait_for(lambda: len(stages[1].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)
@@ -1011,11 +1045,53 @@ async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> N
         assert result.error is None
         assert result.rpc_correlation_key == ("abort", "abort-1")
         assert stages[0].abort_calls == [["req-ack"]]
+        assert stages[1].abort_calls == [["req-ack"]]
         assert "req-ack" not in orchestrator_fixture.orchestrator.request_states
+        assert result.abort_outputs is not None
+        assert len(result.abort_outputs) == 1
+        abort_msg = result.abort_outputs[0]
+        assert abort_msg.request_id == "req-ack"
+        assert abort_msg.finished is True
+        assert abort_msg.stage_id == 1
+        assert list(abort_msg.engine_outputs.outputs[0].token_ids) == [2]
+        assert abort_msg.engine_outputs.outputs[0].finish_reason == "abort"
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
 
 
+@pytest.mark.asyncio
+async def test_abort_marks_only_last_final_stage_output_finished() -> None:
+    """A request with two final AR stages must keep earlier abort outputs unfinished."""
+    stages = [
+        FakeStageClient(stage_type="llm", final_output=True),
+        FakeStageClient(stage_type="llm", final_output=True),
+    ]
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[1], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[2], finished=True)]),
+    ]
+    pools = _build_stage_pools([[stages[0]], [stages[1]]], output_processors=processors)
+    pools[0]._request_bindings["req-two"] = 0
+    pools[1]._request_bindings["req-two"] = 0
+
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=asyncio.Queue(),
+        rpc_async_queue=asyncio.Queue(),
+        stage_pools=pools,
+    )
+    orchestrator.request_states["req-two"] = OrchestratorRequestState(
+        request_id="req-two",
+        final_stage_id=1,
+        final_output_stage_ids={0, 1},
+    )
+
+    outputs = await orchestrator._abort_request_ids(["req-two"])
+    assert [(msg.stage_id, msg.finished) for msg in outputs] == [(0, False), (1, True)]
+    assert list(outputs[0].engine_outputs.outputs[0].token_ids) == [1]
+    assert list(outputs[1].engine_outputs.outputs[0].token_ids) == [2]
+
+
 @pytest.mark.asyncio
 async def test_handle_abort_emits_error_result_on_failure() -> None:
     rpc_queue: asyncio.Queue = asyncio.Queue()

tests/engine/test_output_processor.py

@@ -700,3 +700,62 @@ def test_num_cache_creation_tokens_reaches_direct_request_output():
 
     assert direct.num_cached_tokens == 3
     assert direct.num_cache_creation_tokens == 2
+
+
+def _abort_processor_with_parent(child_ids: list[str]):
+    processor = object.__new__(MultimodalOutputProcessor)
+    processor.request_states = {}
+    processor.external_req_ids = {}
+    processor.parent_requests = {}
+    processor._native_text_metrics_by_request = {}
+    processor.lora_states = SimpleNamespace(request_finished=lambda *_args, **_kwargs: None)
+    parent = SimpleNamespace(request_id="parent", child_requests=set(child_ids))
+    processor.parent_requests["parent"] = parent
+    for child_id in child_ids:
+        req_state = SimpleNamespace(
+            parent_req=parent,
+            lora_name=None,
+            output_kind=RequestOutputKind.CUMULATIVE,
+            detokenizer=SimpleNamespace(output_token_ids=[7, 8]),
+            queue=None,
+            external_req_id="parent",
+            make_request_output=MagicMock(return_value=SimpleNamespace(request_id=child_id)),
+        )
+        processor.request_states[child_id] = req_state
+    return processor, parent
+
+
+def test_abort_last_parallel_child_drops_parent_request():
+    processor, parent = _abort_processor_with_parent(["0_parent", "1_parent"])
+
+    aborted, _outputs = processor.abort_requests_collecting_outputs(["0_parent"], internal=True)
+    assert aborted == ["0_parent"]
+    assert "parent" in processor.parent_requests
+    assert parent.child_requests == {"1_parent"}
+
+    aborted, _outputs = processor.abort_requests_collecting_outputs(["1_parent"], internal=True)
+    assert aborted == ["1_parent"]
+    assert "parent" not in processor.parent_requests
+    assert not parent.child_requests
+
+
+def test_abort_snapshot_leaves_state_until_commit():
+    processor, parent = _abort_processor_with_parent(["0_parent"])
+    processor.external_req_ids["parent"] = ["0_parent"]
+
+    aborted, outputs = processor.abort_requests_collecting_outputs(
+        ["parent"],
+        internal=False,
+        commit_state=False,
+    )
+    assert aborted == ["0_parent"]
+    assert outputs
+    assert "0_parent" in processor.request_states
+    assert processor.external_req_ids["parent"] == ["0_parent"]
+    assert "parent" in processor.parent_requests
+    assert parent.child_requests == {"0_parent"}
+
+    processor.commit_aborted_request_state(["parent"], internal=False)
+    assert "0_parent" not in processor.request_states
+    assert "parent" not in processor.parent_requests
+    assert "parent" not in processor.external_req_ids

tests/engine/test_stage_pool_collective_rpc.py

@@ -64,8 +64,61 @@ def test_collective_rpc_control_helper_honors_timeout():
             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"
+        with pytest.raises(asyncio.TimeoutError):
+            await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,))
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_control_method_reraises_worker_error():
+    async def run() -> None:
+        pool, _client = _make_pool(wake_up_async=AsyncMock(side_effect=RuntimeError("worker died")))
+        with pytest.raises(RuntimeError, match="worker died"):
+            await pool.collective_rpc(0, "wake_up")
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_non_control_method_still_returns_error_dict():
+    async def run() -> None:
+        pool, _client = _make_pool(collective_rpc_async=AsyncMock(side_effect=RuntimeError("probe failed")))
+        result = await pool.collective_rpc(0, "reset_prefix_cache")
         assert result["supported"] is False
+        assert "probe failed" in result["error"]
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
+    async def run() -> None:
+        class RecordingOutputProcessor:
+            def __init__(self) -> None:
+                self.collected = False
+                self.committed = False
+
+            def abort_requests_collecting_outputs(self, request_ids, *, internal=False, commit_state=True):
+                del internal
+                self.collected = True
+                assert commit_state is False
+                return list(request_ids), [SimpleNamespace(request_id=request_ids[0])]
+
+            def commit_aborted_request_state(self, request_ids, *, internal=False):
+                del request_ids, internal
+                self.committed = True
+
+        abort = AsyncMock(side_effect=RuntimeError("engine abort failed"))
+        output_processor = RecordingOutputProcessor()
+        client = SimpleNamespace(stage_type="llm", abort_requests_async=abort)
+        pool = StagePool(0, [client], output_processor=output_processor)  # type: ignore[arg-type]
+        pool._request_bindings["req-1"] = 0
+
+        with pytest.raises(RuntimeError, match="engine abort failed"):
+            await pool.abort_requests(["req-1"])
+        assert output_processor.collected is True
+        assert output_processor.committed is False
+        abort.assert_awaited_once()
 
     asyncio.run(run())

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

tests/worker/test_worker_base.py

@@ -199,7 +199,14 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags):
     worker = object.__new__(OmniGPUWorkerBase)
     worker.rank = 0
     worker.device = "cuda:0"
-    monkeypatch.setattr(base, "current_omni_platform", _fake_platform())
+    order: list[str] = []
+    platform = SimpleNamespace(
+        get_current_memory_usage=lambda device: 0,
+        empty_cache=lambda: order.append("empty_cache"),
+        synchronize=lambda: order.append("sync"),
+        collect=lambda: None,
+    )
+    monkeypatch.setattr(base, "current_omni_platform", platform)
 
     calls = {}
 
@@ -209,12 +216,17 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags):
             return cls()
 
         def sleep(self, offload_tags):
+            order.append("allocator_sleep")
             calls["offload_tags"] = offload_tags
 
     monkeypatch.setattr(cumem_mod, "CuMemAllocator", FakeAllocator)
 
     assert OmniGPUWorkerBase.sleep(worker, level=level) is True
     assert calls["offload_tags"] == expected_tags
+    assert order[0] == "sync"
+    assert "allocator_sleep" in order
+    assert order.index("sync") < order.index("allocator_sleep")
+    assert "empty_cache" not in order
 
 
 def test_wake_up_forwards_tags_to_allocator(monkeypatch):