[CI/Build] Speed up omni sleep-mode entrypoint tests with shared engines (#5713)

Signed-off-by: wangyu <410167048@qq.com>

tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Omni sleep/wakeup JSON validation (live ``Qwen/Qwen-Image`` process)."""
+"""Omni sleep/wakeup JSON validation and level-2 wake contract (live ``Qwen/Qwen-Image``)."""
 
 from __future__ import annotations
 
@@ -15,9 +15,17 @@ pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni]
 
 _SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
 
+# One module-scoped server for this file. ``--enable-sleep-mode`` is required by the
+# level-2 wake contract test; invalid-JSON cases only hit request validation and are
+# unaffected. Do not mix ``omni_server`` (module) with ``omni_server_function`` here:
+# ``iter_omni_server`` holds ``omni_fixture_lock`` across the module yield, so a
+# function-scoped second server would deadlock / contend for the GPU.
 _QWEN_IMAGE = [
     pytest.param(
-        OmniServerParams(model="Qwen/Qwen-Image"),
+        OmniServerParams(
+            model="Qwen/Qwen-Image",
+            server_args=["--enable-sleep-mode"],
+        ),
         id="qwen_image",
         marks=hardware_marks(res={"cuda": "H100"}),
     ),
@@ -93,3 +101,37 @@ def test_omni_sleep_wakeup_invalid_requests(
         openai_client.send_omni_sleep_http_request(cfg)
     else:
         openai_client.send_omni_wakeup_http_request(cfg)
+
+
+# Keep this after the invalid-JSON cases: sleep(level=2) discards weights and the
+# shared module server is not usable afterward; module teardown then stops it.
+@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
+def test_wakeup_after_level2_sleep_fails(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+) -> None:
+    """Regression for #4473 Repro A via HTTP: wake after sleep(level=2) → 501.
+
+    Moved from ``tests/entrypoints/test_omni_sleep_mode.py`` (BAGEL cold start).
+    Live serve surfaces ``NotImplementedError`` as OpenAI-style JSON with code 501.
+    """
+    assert omni_server is not None
+    sleep_resps = openai_client.send_omni_sleep_http_request(
+        {
+            "json": {"stage_ids": [0], "level": 2},
+            "timeout": 180,
+        }
+    )
+    sleep_resp = sleep_resps[0]
+    assert sleep_resp.status_code == 200, sleep_resp
+    assert sleep_resp.json_body is not None
+    assert sleep_resp.json_body.get("status") == "SUCCESS", sleep_resp.json_body
+
+    openai_client.send_omni_wakeup_http_request(
+        {
+            "json": {"stage_ids": [0]},
+            "timeout": 120,
+            "err_code": 501,
+            "err_message": ("sleep(level=2)", "not yet implemented"),
+        }
+    )

tests/entrypoints/openai_api/test_omni_sleep_wakeup.py

@@ -175,6 +175,26 @@ def test_wakeup_engine_not_support(sleep_incapable_engine):
     assert "wake_up" in response.json()["detail"]
 
 
+def test_wakeup_route_propagates_not_implemented_error(sleep_capable_engine, mocker):
+    """``/v1/omni/wakeup`` must not swallow ``NotImplementedError`` from the engine.
+
+    Route exception-propagation only: the mock authors the exception on purpose.
+    The real level-2 discarded-weight path and live HTTP 501 mapping are covered by
+    ``test_wakeup_after_level2_sleep_fails``.
+    """
+    sleep_capable_engine.wake_up = mocker.AsyncMock(side_effect=NotImplementedError("engine wake_up refused"))
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0}
+    # Re-raise in-process; default TestClient would wrap this as a 500 response.
+    client = TestClient(app, raise_server_exceptions=True)
+
+    with pytest.raises(NotImplementedError, match="engine wake_up refused"):
+        client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+
+    # Failed wake must not clear the sleeping set.
+    assert app.state.sleeping_stages == {0}
+
+
 def test_wakeup_removes_only_requested_stages(sleep_capable_engine):
     app = _make_app(sleep_capable_engine)
     app.state.sleeping_stages = {0, 1, 2}

tests/entrypoints/test_omni_sleep_mode.py

@@ -1,13 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Omni sleep mode: entrypoint-level VRAM/ACK tests (L4) plus H100 multi-TP e2e."""
+"""Omni sleep mode: entrypoint-level VRAM/ACK tests plus multi-stage e2e.
+
+Class-scoped engines amortize BAGEL load across sleep/wake assertions within
+each suite (LLM vs diffusion). Suites are separate classes so only one heavy
+engine is alive at a time. Device cleanup runs once per module.
+"""
 
 import asyncio
 import logging
 import os
 
 import pytest
+import pytest_asyncio
 import torch
 from vllm import SamplingParams
 
@@ -20,12 +26,11 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("OmniTest")
 pytestmark = [
     pytest.mark.advanced_model,
-    pytest.mark.usefixtures("clean_gpu_memory_between_tests"),
 ]
 
 
-def clean_gpu_envs():
-    """clean up GPU environment variables to ensure tests run on all available devices."""
+def clean_device_envs():
+    """Clear device-visibility env vars so tests see all available devices."""
     device_visibility_vars = [
         "CUDA_VISIBLE_DEVICES",  # NVIDIA
         "HIP_VISIBLE_DEVICES",  # AMD ROCm
@@ -103,11 +108,15 @@ def get_dynamic_devices(stage_idx: int, num_stages: int, tp_size: int) -> str:
     return ",".join(device_ids)
 
 
-@pytest.fixture(scope="function")
-async def llm_engine():
-    if current_omni_platform.is_rocm():
-        clean_gpu_envs()
-    model_name = "ByteDance-Seed/BAGEL-7B-MoT"
+async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None:
+    """Best-effort full wake so the next shared-engine test starts active."""
+    try:
+        await engine.wake_up(stage_ids=stage_ids)
+    except Exception as e:
+        logger.warning("ensure_awake failed (stage_ids=%s): %s", stage_ids, e)
+
+
+def _build_llm_stages() -> tuple[list[dict], list[dict]]:
     common_args = {
         "worker_type": "ar",
         "enable_sleep_mode": True,
@@ -133,18 +142,38 @@ async def llm_engine():
         },
     ]
     connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}]
-    engine = AsyncOmni(model=model_name, stages=stages, connectors=connectors, init_timeout=600, enable_sleep_mode=True)
+    return stages, connectors
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _module_device_cleanup():
+    """One device cleanup pass around the module (engines are class-scoped)."""
+    from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+
+    print("\n=== PRE-MODULE DEVICE CLEANUP (sleep_mode) ===")
+    run_pre_test_cleanup()
+    yield
+    print("\n=== POST-MODULE DEVICE CLEANUP (sleep_mode) ===")
+    run_post_test_cleanup()
+
+
+@pytest_asyncio.fixture(scope="class", loop_scope="class")
+async def llm_engine():
+    """Shared 2-stage BAGEL LLM engine for sleep/wake protocol tests."""
+    if current_omni_platform.is_rocm():
+        clean_device_envs()
+    stages, connectors = _build_llm_stages()
+    engine = AsyncOmni(model=MODEL, stages=stages, connectors=connectors, init_timeout=600, enable_sleep_mode=True)
     yield engine
     engine.shutdown()
-    # Subprocess / driver can lag releasing VRAM; brief pause before the next test spins up new workers.
     await asyncio.sleep(1.5)
 
 
-@pytest.fixture(scope="function")
+@pytest_asyncio.fixture(scope="class", loop_scope="class")
 async def diffusion_engine():
+    """Shared BAGEL diffusion TP=2 engine for sleep/wake + generate checks."""
     if current_omni_platform.is_rocm():
-        clean_gpu_envs()
-    model_name = "ByteDance-Seed/BAGEL-7B-MoT"
+        clean_device_envs()
     stages = [
         {
             "stage_id": 0,
@@ -165,28 +194,27 @@ async def diffusion_engine():
             "final_output_type": "image",
         }
     ]
-    engine = AsyncOmni(model=model_name, stages=stages, init_timeout=600, enable_sleep_mode=True)
+    engine = AsyncOmni(model=MODEL, stages=stages, init_timeout=600, enable_sleep_mode=True)
     yield engine
     engine.shutdown()
     await asyncio.sleep(1.5)
 
 
-class TestOmniSleepMode:
-    @pytest.mark.asyncio
+class TestOmniLlmSleepMode:
+    """LLM sleep/wake protocol tests sharing one class-scoped engine."""
+
+    @pytest.mark.asyncio(loop_scope="class")
     @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
     async def test_llm_sleep_ack(self, llm_engine: AsyncOmni):
         """LLM Thinker (GPU0) Signal and Physical Recycling Audit"""
+        device_id = 0
         try:
-            device_id = 0
             used_before = get_device_global_memory_used_gib(device_id)
             acks = await llm_engine.sleep(stage_ids=[0], level=1)
             await asyncio.sleep(1.5)
             used_after = get_device_global_memory_used_gib(device_id)
             drop_gib = used_before - used_after
-            # Verification signal successful
             assert all(get_ack_info(ack, "status") == "SUCCESS" for ack in acks)
-            # Worker-reported delta (can be 0 if get_current_memory_usage does not move) or
-            # GPU-global drop from mem_get_info (sees child worker processes).
             total_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks)
             freed_gib = total_freed_bytes / 1024**3
             logger.info(
@@ -201,9 +229,42 @@ class TestOmniSleepMode:
                 f"ACK={freed_gib:.2f} GiB, global_drop={drop_gib:.2f} GiB"
             )
         finally:
-            llm_engine.shutdown()
+            await _ensure_awake(llm_engine, [0])
 
-    @pytest.mark.asyncio
+    @pytest.mark.asyncio(loop_scope="class")
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_partial_wake_blocks_generate(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro B: generate() must be rejected if kv_cache
+        is still asleep after wake_up(tags=["weights"]), instead of crashing with
+        CUDA illegal memory access."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=1)
+            await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
+            with pytest.raises(RuntimeError, match="partially or fully asleep"):
+                async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
+                    pass
+        finally:
+            await _ensure_awake(llm_engine, [0])
+
+    @pytest.mark.asyncio(loop_scope="class")
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_duplicate_wake_is_idempotent(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro C: duplicate wake_up(tags=None) must be a
+        safe no-op instead of raising a cumem CUDA invalid argument error."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=1)
+            first_acks = await llm_engine.wake_up(stage_ids=[0])
+            assert len(first_acks) > 0, "First wake_up() should return ACKs"
+            second_acks = await llm_engine.wake_up(stage_ids=[0])
+            assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}"
+        finally:
+            await _ensure_awake(llm_engine, [0])
+
+
+class TestOmniDiffusionSleepMode:
+    """Diffusion sleep/wake tests sharing one class-scoped TP=2 engine."""
+
+    @pytest.mark.asyncio(loop_scope="class")
     @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
     async def test_diffusion_sleep_handshake(self, diffusion_engine: AsyncOmni):
         """Diffusion Worker stage signal loop"""
@@ -217,19 +278,15 @@ class TestOmniSleepMode:
             assert len(acks) >= 1, "Expected at least 1 ACK from Diffusion Workers"
             assert all(_get_status(ack) == "SUCCESS" for ack in acks)
             logger.info(f"Success: Received {len(acks)} Diffusion Worker ACKs")
-            logger.info("Testing auto-wakeup before test end...")
             await diffusion_engine.wake_up(stage_ids=[0])
-            logger.info("Test logic finished, triggering manual shutdown...")
         finally:
-            diffusion_engine.shutdown()
-            logger.info("Manual shutdown executed. Test should exit now.")
+            await _ensure_awake(diffusion_engine, [0])
 
-    @pytest.mark.asyncio
+    @pytest.mark.asyncio(loop_scope="class")
     @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
     async def test_cross_device_cleanup(self, diffusion_engine: AsyncOmni):
         """Physical recycling audit: leveraging deterministic data returned by Workers"""
         try:
-            # TP2 uses GPUs 0 and 1; measure whole-GPU usage (includes worker subprocesses).
             used_before = get_device_global_memory_used_gib(0) + get_device_global_memory_used_gib(1)
             acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
             await asyncio.sleep(1.5)
@@ -246,52 +303,90 @@ class TestOmniSleepMode:
             )
             logger.info("SUCCESS: 100% weights offloaded.")
         finally:
-            diffusion_engine.shutdown()
+            await _ensure_awake(diffusion_engine, [0])
 
-    @pytest.mark.asyncio
+    @pytest.mark.asyncio(loop_scope="class")
     @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_diffusion_integrity_bit_level(self, diffusion_engine: AsyncOmni):
-        """Bit-level consistency after Diffusion wake-up (prevent image corruption)"""
+    async def test_diffusion_sleep_wake_generate(self, diffusion_engine: AsyncOmni):
+        """Merged diffusion lifecycle with test points from the former trio:
+
+        - ``test_diffusion_integrity_bit_level``: 4-step/512 baseline generate →
+          sleep → wake → post generate (images present, count unchanged).
+        - ``test_diffusion_vram_lifecycle_audit``: ACK/process VRAM reclaim on
+          sleep; after wake, process VRAM restores to within 3 GiB of initial;
+          post-wake generate smoke.
+        - ``test_diffusion_model_sleep_tp``: sleep ACK SUCCESS → wake → verify
+          generate (same Active→Sleep→Active→Generate path; runs on the shared
+          class TP=2 BagelPipeline engine rather than a second default-config load).
+        """
+        import gc
+
+        device_id = 1
         try:
+            # --- integrity: baseline generate at original resolution/steps ---
             prompt = "A huge swimming pool, with many people swimming."
             sp = OmniDiffusionSamplingParams(num_inference_steps=4, height=512, width=512, seed=42)
             llm_sp = SamplingParams()
 
-            # Baseline Generation
             logger.info("Running Baseline Generation...")
             base_output = None
             async for output in diffusion_engine.generate(prompt, request_id="base", sampling_params_list=[llm_sp, sp]):
                 base_output = output
             assert base_output is not None and len(base_output.images) > 0
             logger.info("Baseline Generation successful.")
-            # Sleep Level 1
+
+            # Snapshot Active VRAM after generate (lifecycle compared Sleep/Wake
+            # against the pre-sleep active watermark).
+            get_vram_info(device_id)
+            torch.accelerator.empty_cache()
+            vram_initial = get_vram_info(device_id)["reserved"]
+            logger.info(f"Diffusion Active VRAM before sleep: {vram_initial:.2f} GiB")
+
+            # --- lifecycle + model_sleep_tp: level-1 sleep reclaim ---
             logger.info("Entering Deep Sleep (VRAM Scavenging)...")
-            await diffusion_engine.sleep(stage_ids=[0], level=1)
-            # Wake-up
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
+            statuses = [get_ack_info(ack, "status") for ack in acks]
+            assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
+
+            reported_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks)
+            reported_freed_gib = reported_freed_bytes / 1024**3
+            logger.info(f"Worker internally reported freed: {reported_freed_gib:.2f} GiB")
+
+            await asyncio.sleep(2)
+            get_vram_info(device_id)
+            torch.accelerator.empty_cache()
+            vram_sleeping = get_vram_info(device_id)["reserved"]
+            logger.info(f"External VRAM measurement during Sleep: {vram_sleeping:.2f} GiB")
+            assert reported_freed_gib > 14.0 or vram_sleeping < 5.0, (
+                f"Reclamation failed. Reported: {reported_freed_gib:.2f}G, Measured: {vram_sleeping:.2f}G"
+            )
+
+            # --- wake + VRAM restore to initial (±3 GiB) ---
             logger.info("Waking up (Reloading Weights)...")
             await diffusion_engine.wake_up(stage_ids=[0])
-
             await asyncio.sleep(2.0)
-            import gc
-
             gc.collect()
+            get_vram_info(device_id)
+            torch.accelerator.empty_cache()
+            vram_restored = get_vram_info(device_id)["reserved"]
+            logger.info(f"VRAM after Wake-up: {vram_restored:.2f} GiB")
+            assert abs(vram_restored - vram_initial) < 3.0, "VRAM failed to restore to initial levels"
 
+            # --- integrity + lifecycle: post-wake generate ---
             logger.info("Running Post-Wakeup Generation...")
             post_output = None
             async for output in diffusion_engine.generate(prompt, request_id="post", sampling_params_list=[llm_sp, sp]):
                 post_output = output
-            # Assert result consistency
             assert post_output is not None
             assert len(base_output.images) == len(post_output.images)
             assert post_output.images[0] is not None
-            logger.info("SUCCESS: Diffusion integrity verified after Sleep/Wake cycle.")
-        except Exception as e:
-            logger.error(f"Integrity test failed: {e}")
-            raise e
+            logger.info("SUCCESS: Diffusion integrity + VRAM lifecycle + sleep/wake generate verified.")
         finally:
-            logger.info("Triggering mandatory cleanup...")
-            diffusion_engine.shutdown()
-            logger.info("Cleanup complete, test exiting.")
+            await _ensure_awake(diffusion_engine, [0])
+
+
+class TestOmniCoordinatedSleepMode:
+    """Dual-engine coordination (kept skipped; do not delete)."""
 
     @pytest.mark.skip(
         reason=(
@@ -301,7 +396,7 @@ class TestOmniSleepMode:
             "Note: this test is fixture-free so skip does not init llm/diffusion engines."
         )
     )
-    @pytest.mark.asyncio
+    @pytest.mark.asyncio(loop_scope="class")
     @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
     async def test_coordinated_cross_device(self, llm_engine: AsyncOmni, diffusion_engine: AsyncOmni):
         """Heterogeneous Coordinated Cleanup Test (Talker and Diffusion on GPU 1)"""
@@ -341,172 +436,22 @@ class TestOmniSleepMode:
             diffusion_engine.shutdown()
             logger.info("All engines scavenged. Ready for next test.")
 
-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_diffusion_vram_lifecycle_audit(self, diffusion_engine: AsyncOmni):
-        """Diffusion memory loop: Active -> Deep Sleep -> Active -> inference sanity check"""
-        device_id = 1
-        try:
-            get_vram_info(device_id)
-            torch.accelerator.empty_cache()
-            vram_initial = get_vram_info(device_id)["reserved"]
-            logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB")
-
-            # Sleep
-            logger.info("Triggering Level 1 Deep Sleep (Partial Weight Offloading)...")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
-
-            reported_freed_bytes = sum(getattr(ack, "freed_bytes", 0) for ack in acks)
-            reported_freed_gib = reported_freed_bytes / 1024**3
-            logger.info(f"Worker internally reported freed: {reported_freed_gib:.2f} GiB")
-
-            await asyncio.sleep(2)
-            get_vram_info(device_id)
-            torch.accelerator.empty_cache()
-
-            vram_sleeping = get_vram_info(device_id)["reserved"]
-            logger.info(f"External VRAM measurement during Sleep: {vram_sleeping:.2f} GiB")
-
-            assert reported_freed_gib > 14.0 or vram_sleeping < 5.0, (
-                f"Reclamation failed. Reported: {reported_freed_gib:.2f}G, Measured: {vram_sleeping:.2f}G"
-            )
-
-            # wake-up
-            logger.info("Triggering Wake-up (Reloading weights to GPU)...")
-            await diffusion_engine.wake_up(stage_ids=[0])
-
-            await asyncio.sleep(2)
-            get_vram_info(device_id)
-            torch.accelerator.empty_cache()
-            vram_restored = get_vram_info(device_id)["reserved"]
-            logger.info(f"VRAM after Wake-up: {vram_restored:.2f} GiB")
-
-            assert abs(vram_restored - vram_initial) < 3.0, "VRAM failed to restore to initial levels"
-
-            # inference sanity check
-            logger.info("Running post-lifecycle inference smoke test...")
-            prompt = "A futuristic lab with glowing lights, high quality."
-            sp = OmniDiffusionSamplingParams(num_inference_steps=2, height=512, width=512, seed=42)
-            llm_sp = SamplingParams()
-
-            base_img_found = False
-            async for output in diffusion_engine.generate(
-                prompt, request_id="lifecycle-check", sampling_params_list=[llm_sp, sp]
-            ):
-                if output.images and output.images[0] is not None:
-                    base_img_found = True
-
-            assert base_img_found, "Inference failed after Wake-up cycle!"
-            logger.info("SUCCESS: Full Diffusion Lifecycle (Active -> Sleep -> Active -> Generate) audited.")
-
-        except Exception as e:
-            logger.error(f"Lifecycle audit failed: {e}")
-            raise e
-        finally:
-            logger.info("Cleaning up engine and scavenging processes...")
-            diffusion_engine.shutdown()
-            await asyncio.sleep(1)
-
-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_level2_sleep_wake_raises(self, llm_engine: AsyncOmni):
-        """Regression for #4473 Repro A: wake_up() after sleep(level=2) must raise
-        NotImplementedError instead of silently producing corrupted output."""
-        try:
-            await llm_engine.sleep(stage_ids=[0], level=2)
-            with pytest.raises(NotImplementedError, match="sleep\\(level=2\\)"):
-                await llm_engine.wake_up(stage_ids=[0])
-        finally:
-            llm_engine.shutdown()
-
-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_partial_wake_blocks_generate(self, llm_engine: AsyncOmni):
-        """Regression for #4473 Repro B: generate() must be rejected if kv_cache
-        is still asleep after wake_up(tags=["weights"]), instead of crashing with
-        CUDA illegal memory access."""
-        try:
-            await llm_engine.sleep(stage_ids=[0], level=1)
-            await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
-            with pytest.raises(RuntimeError, match="partially or fully asleep"):
-                async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
-                    pass
-        finally:
-            llm_engine.shutdown()
-
-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_duplicate_wake_is_idempotent(self, llm_engine: AsyncOmni):
-        """Regression for #4473 Repro C: duplicate wake_up(tags=None) must be a
-        safe no-op instead of raising a cumem CUDA invalid argument error."""
-        try:
-            await llm_engine.sleep(stage_ids=[0], level=1)
-            first_acks = await llm_engine.wake_up(stage_ids=[0])
-            assert len(first_acks) > 0, "First wake_up() should return ACKs"
-            second_acks = await llm_engine.wake_up(stage_ids=[0])
-            assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}"
-        finally:
-            llm_engine.shutdown()
-
 
 # ---------------------------------------------------------------------------
-# H100 (or MI325) e2e: BAGEL / pure diffusion, multi-TP
+# Multi-stage (H100) / pure diffusion TP=1 (2×L4): sleep→wake→generate
+# (BAGEL TP=2 diffusion path covered by TestOmniDiffusionSleepMode)
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.omni
-@pytest.mark.advanced_model
-@pytest.mark.parametrize("tp_size", [1])
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
-@pytest.mark.asyncio
-async def test_diffusion_model_sleep_tp(tp_size: int):
-    """Two-stage BAGEL default config: warmup, sleep all, wake, verify generate."""
-    if current_omni_platform.is_rocm():
-        clean_gpu_envs()
-    num_gpus = torch.accelerator.device_count()
-    if num_gpus < tp_size:
-        pytest.skip(f"Skipping TP={tp_size}")
-
-    engine_args = {
-        "model": MODEL,
-        "enable_sleep_mode": True,
-        "tensor_parallel_size": tp_size,
-        "enforce_eager": True,
-        "trust_remote_code": True,
-        "dtype": "bfloat16",
-        "gpu_memory_utilization": 0.5,
-    }
-
-    engine = AsyncOmni(**engine_args, stage_init_timeout=1200)
-    try:
-        diff_sp = OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256)
-        llm_sp = SamplingParams()
-
-        async for _ in engine.generate("test", sampling_params_list=[llm_sp, diff_sp]):
-            pass
-
-        acks = await engine.sleep(level=1)
-        statuses = [get_ack_info(ack, "status") for ack in acks]
-        assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
-
-        await engine.wake_up()
-        async for _ in engine.generate("verify", sampling_params_list=[llm_sp, diff_sp]):
-            pass
-
-        logger.info("Diffusion TP=%s lifecycle OK", tp_size)
-    finally:
-        engine.shutdown()
-
-
 @pytest.mark.omni
 @pytest.mark.advanced_model
 @pytest.mark.parametrize("tp_size", [1, 2])
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.asyncio
-async def test_multistage_sleep_h100(tp_size: int):
+async def test_multistage_llm_diffusion_sleep_wake(tp_size: int):
     """Explicit 2-stage (llm + diffusion) + connectors; sleep/wake both stages."""
     if current_omni_platform.is_rocm():
-        clean_gpu_envs()
+        clean_device_envs()
     num_gpus = torch.accelerator.device_count()
     if num_gpus < tp_size * 2:
         pytest.skip("Not enough GPUs")
@@ -552,18 +497,22 @@ async def test_multistage_sleep_h100(tp_size: int):
 
 
 @pytest.mark.omni
-@pytest.mark.advanced_model
-@pytest.mark.parametrize("tp_size", [1, 2])
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.asyncio
-async def test_pure_diffusion_scenario(tp_size: int):
-    """Single-stage random diffusion: sleep, wake, generate."""
+async def test_pure_diffusion_sleep_wake():
+    """Single-stage random diffusion (TP=1): sleep → wake → generate.
+
+    Diffusion TP=2 sleep/wake is covered by TestOmniDiffusionSleepMode
+    (shared BAGEL TP=2 engine); this case keeps the default AsyncOmni +
+    qwen_image_random entrypoint path on L4.
+    """
     if current_omni_platform.is_rocm():
-        clean_gpu_envs()
+        clean_device_envs()
     engine_args = {
         "model": MODEL_DIFF,
         "enable_sleep_mode": True,
-        "tensor_parallel_size": tp_size,
+        "tensor_parallel_size": 1,
         "enforce_eager": True,
         "dtype": "bfloat16",
         "gpu_memory_utilization": 0.5,
@@ -578,6 +527,6 @@ async def test_pure_diffusion_scenario(tp_size: int):
             sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256),
         ):
             pass
-        logger.info("Pure diffusion OK (TP=%s)", tp_size)
+        logger.info("Pure diffusion sleep/wake OK")
     finally:
         engine.shutdown()