diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py b/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py index 0a8240c4b..ad3f659f1 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py +++ b/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"), + } + ) diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py index cf7d6e234..c47eb8da6 100644 --- a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py +++ b/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} diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py index efbb19fff..0377e1a9a 100644 --- a/tests/entrypoints/test_omni_sleep_mode.py +++ b/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()