vllm-omni/tests/entrypoints/test_omni_sleep_mode.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""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

from tests.helpers.mark import hardware_test
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("OmniTest")
pytestmark = [
    pytest.mark.advanced_model,
]


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
        "ZE_AFFINITY_MASK",  # Intel XPU
        "ONEAPI_DEVICE_SELECTOR",  # Intel OneAPI
        "ASCEND_RT_VISIBLE_DEVICES",  # Huawei NPU (CAN)
    ]
    for key in device_visibility_vars:
        os.environ.pop(key, None)


def get_vram_info(device_id: int) -> dict:
    """Per-**process** CUDA allocator stats (GiB). Does not see other processes' usage."""
    try:
        if current_omni_platform.is_rocm():
            num_gpus = torch.accelerator.device_count()
            safe_id = device_id if device_id < num_gpus else 0
            torch.accelerator.synchronize(safe_id)
            return {
                "reserved": torch.cuda.memory_reserved(safe_id) / 1024**3,
                "allocated": torch.cuda.memory_allocated(safe_id) / 1024**3,
            }
        else:
            with torch.cuda.device(device_id):
                torch.accelerator.synchronize()
                return {
                    "reserved": torch.cuda.memory_reserved() / 1024**3,
                    "allocated": torch.cuda.memory_allocated() / 1024**3,
                }
    except Exception as e:
        logger.warning(f"memory skip ({device_id}): {e}")
        return {"reserved": 0.0, "allocated": 0.0}


def get_device_global_memory_used_gib(device_id: int) -> float:
    """GPU-wide memory in use (GiB), includes all processes (driver view).

    Uses ``torch.cuda.mem_get_info`` (CUDA / ROCm when available). Use this when stages run in
    **separate worker processes**; ``memory_reserved`` in the test process does not reflect
    child workers' allocations.
    """
    try:
        with torch.cuda.device(device_id):
            torch.accelerator.synchronize()
            free_b, total_b = torch.cuda.mem_get_info()
        return (total_b - free_b) / 1024**3
    except Exception as e:
        logger.warning("get_device_global_memory_used_gib(%s): %s", device_id, e)
        return 0.0


def get_ack_info(ack, key, default=None):
    """
    Since ACKs in a distributed environment can be either objects or dictionaries,
    this tool ensures compatibility.
    """
    if hasattr(ack, key):
        return getattr(ack, key)
    if isinstance(ack, dict):
        return ack.get(key, default)
    return default


MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
MODEL_DIFF = "riverclouds/qwen_image_random"


def get_dynamic_devices(stage_idx: int, num_stages: int, tp_size: int) -> str:
    total_gpus = torch.accelerator.device_count()
    gpus_per_stage = tp_size
    start_idx = stage_idx * gpus_per_stage
    if start_idx + gpus_per_stage > total_gpus:
        start_idx = start_idx % total_gpus
    device_ids = [str(start_idx + i) for i in range(gpus_per_stage)]
    return ",".join(device_ids)


async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None:
    """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:
        logger.warning("ensure_awake failed (stage_ids=%s): %s", stage_ids, e)
    try:
        await engine.resume_generation(stage_ids=stage_ids)
    except Exception as e:
        logger.warning("ensure_resume 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,
        "dtype": "bfloat16",
        "trust_remote_code": True,
        "max_model_len": 2048,
        "max_num_batched_tokens": 8192,
        "enforce_eager": True,
    }
    stages = [
        {
            "stage_id": 0,
            "stage_type": "llm",
            "runtime": {"process": True, "devices": "0", "max_batch_size": 1},
            "engine_args": {**common_args, "model_stage": "thinker", "gpu_memory_utilization": 0.1},
        },
        {
            "stage_id": 1,
            "stage_type": "llm",
            "engine_input_source": [0],
            "runtime": {"process": True, "devices": "1", "max_batch_size": 1, "connector_type": "queue"},
            "engine_args": {**common_args, "model_stage": "talker", "gpu_memory_utilization": 0.1},
        },
    ]
    connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}]
    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.clean import cleanup_test_environment

    print("\n=== PRE-MODULE DEVICE CLEANUP (sleep_mode) ===")
    cleanup_test_environment()
    yield
    print("\n=== POST-MODULE DEVICE CLEANUP (sleep_mode) ===")
    cleanup_test_environment()


@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()
    await asyncio.sleep(1.5)


@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_device_envs()
    stages = [
        {
            "stage_id": 0,
            "stage_type": "diffusion",
            "runtime": {"process": True, "devices": "0,1", "max_batch_size": 1},
            "engine_args": {
                "model_stage": "base",
                "gpu_memory_utilization": 0.1,
                "model_class_name": "BagelPipeline",
                "enable_sleep_mode": True,
                "enforce_eager": True,
                "max_num_batched_tokens": 8192,
                "parallel_config": {
                    "tensor_parallel_size": 2,
                },
            },
            "final_output": True,
            "final_output_type": "image",
        }
    ]
    engine = AsyncOmni(model=MODEL, stages=stages, init_timeout=600, enable_sleep_mode=True)
    yield engine
    engine.shutdown()
    await asyncio.sleep(1.5)


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:
            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
            assert all(get_ack_info(ack, "status") == "SUCCESS" for ack in acks)
            total_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks)
            freed_gib = total_freed_bytes / 1024**3
            logger.info(
                "Thinker: ACK freed=%.2f GiB, global GPU used drop=%.2f GiB (before=%.2f, after=%.2f)",
                freed_gib,
                drop_gib,
                used_before,
                used_after,
            )
            assert freed_gib > 5.0 or drop_gib > 3.0, (
                "Expected either ACK freed_bytes or global VRAM drop after sleep. "
                f"ACK={freed_gib:.2f} GiB, global_drop={drop_gib:.2f} GiB"
            )
        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_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"])
            # wake_up restores weights but keeps the AR admission hold.
            # Resume so generate() can reach the sleeping-tags rejection.
            await llm_engine.resume_generation(stage_ids=[0])
            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"""
        try:
            logger.info("Starting Diffusion Worker Handshake Test")
            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)

            def _get_status(ack):
                return ack.status if hasattr(ack, "status") else ack.get("status")

            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")
            await diffusion_engine.wake_up(stage_ids=[0])
        finally:
            await _ensure_awake(diffusion_engine, [0])

    @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:
            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)
            used_after = get_device_global_memory_used_gib(0) + get_device_global_memory_used_gib(1)
            drop_gib = used_before - used_after
            total_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks)
            freed_gb = total_freed_bytes / 1024**3
            logger.info("Physical reclamation summary from workers:")
            logger.info(f"- Total Workers: {len(acks)}")
            logger.info(f"- Total Freed (ACK): {freed_gb:.2f} GiB, global used drop: {drop_gib:.2f} GiB")
            assert freed_gb > 14.0 or drop_gib > 8.0, (
                "Expected either ACK freed_bytes or global VRAM drop on GPUs 0+1. "
                f"ACK={freed_gb:.2f} GiB, global_drop={drop_gib:.2f} GiB"
            )
            logger.info("SUCCESS: 100% weights offloaded.")
        finally:
            await _ensure_awake(diffusion_engine, [0])

    @pytest.mark.asyncio(loop_scope="class")
    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
    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()

            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.")

            # 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)...")
            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])
            # AR stages keep the frontend admission gate held after wake.
            await diffusion_engine.resume_generation(stage_ids=[0])
            await asyncio.sleep(2.0)
            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 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 + VRAM lifecycle + sleep/wake generate verified.")
        finally:
            await _ensure_awake(diffusion_engine, [0])


class TestOmniCoordinatedSleepMode:
    """Dual-engine coordination (kept skipped; do not delete)."""

    @pytest.mark.skip(
        reason=(
            "Flaky/CI: dual AsyncOmni can fail with "
            "RuntimeError: Orchestrator init failed, StageDiffusionProc died during handshake. "
            "Re-enable when stable (no OOM on coordinated talker+diffusion). "
            "Note: this test is fixture-free so skip does not init llm/diffusion engines."
        )
    )
    @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)"""
        device_id = 1
        try:
            logger.info(f"Waking up both engines on GPU {device_id}...")
            await llm_engine.wake_up(stage_ids=[1])
            await diffusion_engine.wake_up(stage_ids=[0])

            get_vram_info(device_id)
            torch.accelerator.empty_cache()
            await asyncio.sleep(2)

            initial_vram = get_vram_info(device_id)["reserved"]
            logger.info(f"GPU {device_id} Peak Pressure: {initial_vram:.2f} GiB")

            # coordinated sleep
            logger.info("Issuing concurrent SLEEP commands...")
            await llm_engine.sleep(stage_ids=[1], level=2)
            await asyncio.sleep(1.0)
            await diffusion_engine.sleep(stage_ids=[0], level=2)

            await asyncio.sleep(3.0)
            torch.accelerator.empty_cache()

            final_vram = get_vram_info(device_id)["reserved"]
            logger.info(f"GPU {device_id} Final VRAM after coordinated sleep: {final_vram:.2f} GiB")

            assert initial_vram - final_vram > 15.0 or final_vram < 8.0
            logger.info(f"SUCCESS: Heterogeneous VRAM drop verified on GPU {device_id}.")
        except Exception as e:
            logger.error(f"Coordinated test failed: {e}")
            raise e
        finally:
            logger.info("Triggering mandatory cleanup for both engines...")
            llm_engine.shutdown()
            diffusion_engine.shutdown()
            logger.info("All engines scavenged. Ready for next test.")


# ---------------------------------------------------------------------------
# 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, 2])
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.asyncio
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_device_envs()
    num_gpus = torch.accelerator.device_count()
    if num_gpus < tp_size * 2:
        pytest.skip("Not enough GPUs")

    stages = []
    for i in range(2):
        devs = get_dynamic_devices(i, 2, tp_size)
        stages.append(
            {
                "stage_id": i,
                "stage_type": "llm" if i == 0 else "diffusion",
                "runtime": {"process": True, "devices": devs},
                "engine_args": {
                    "model": MODEL,
                    "model_stage": "thinker" if i == 0 else "base",
                    "tensor_parallel_size": tp_size,
                    "gpu_memory_utilization": 0.4,
                    "dtype": "bfloat16",
                    "enable_sleep_mode": True,
                    "trust_remote_code": True,
                },
            }
        )

    connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}]

    engine = AsyncOmni(
        model=MODEL, stages=stages, connectors=connectors, enable_sleep_mode=True, stage_init_timeout=1200
    )
    try:
        sp = OmniDiffusionSamplingParams(num_inference_steps=2)
        async for _ in engine.generate("warmup", sampling_params_list=[SamplingParams(), sp]):
            pass

        acks = await engine.sleep(stage_ids=[0, 1], level=1)
        assert len(acks) == 2

        await engine.wake_up(stage_ids=[0, 1])
        await engine.resume_generation(stage_ids=[0, 1])
        async for _ in engine.generate("verify", sampling_params_list=[SamplingParams(), sp]):
            pass
    finally:
        engine.shutdown()


@pytest.mark.omni
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.asyncio
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_device_envs()
    engine_args = {
        "model": MODEL_DIFF,
        "enable_sleep_mode": True,
        "tensor_parallel_size": 1,
        "enforce_eager": True,
        "dtype": "bfloat16",
        "gpu_memory_utilization": 0.5,
    }

    engine = AsyncOmni(**engine_args, stage_init_timeout=1200)
    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),
        ):
            pass
        logger.info("Pure diffusion sleep/wake OK")
    finally:
        engine.shutdown()