584 lines
24 KiB
Python
584 lines
24 KiB
Python
# 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."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
|
|
import pytest
|
|
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,
|
|
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."""
|
|
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)
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
async def llm_engine():
|
|
if current_omni_platform.is_rocm():
|
|
clean_gpu_envs()
|
|
model_name = "ByteDance-Seed/BAGEL-7B-MoT"
|
|
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"}]
|
|
engine = AsyncOmni(model=model_name, 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")
|
|
async def diffusion_engine():
|
|
if current_omni_platform.is_rocm():
|
|
clean_gpu_envs()
|
|
model_name = "ByteDance-Seed/BAGEL-7B-MoT"
|
|
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_name, stages=stages, init_timeout=600, enable_sleep_mode=True)
|
|
yield engine
|
|
engine.shutdown()
|
|
await asyncio.sleep(1.5)
|
|
|
|
|
|
class TestOmniSleepMode:
|
|
@pytest.mark.asyncio
|
|
@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"""
|
|
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(
|
|
"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:
|
|
llm_engine.shutdown()
|
|
|
|
@pytest.mark.asyncio
|
|
@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")
|
|
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.")
|
|
|
|
@pytest.mark.asyncio
|
|
@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)
|
|
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:
|
|
diffusion_engine.shutdown()
|
|
|
|
@pytest.mark.asyncio
|
|
@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)"""
|
|
try:
|
|
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
|
|
logger.info("Entering Deep Sleep (VRAM Scavenging)...")
|
|
await diffusion_engine.sleep(stage_ids=[0], level=1)
|
|
# Wake-up
|
|
logger.info("Waking up (Reloading Weights)...")
|
|
await diffusion_engine.wake_up(stage_ids=[0])
|
|
|
|
await asyncio.sleep(2.0)
|
|
import gc
|
|
|
|
gc.collect()
|
|
|
|
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
|
|
finally:
|
|
logger.info("Triggering mandatory cleanup...")
|
|
diffusion_engine.shutdown()
|
|
logger.info("Cleanup complete, test exiting.")
|
|
|
|
@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
|
|
@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.")
|
|
|
|
@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
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@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):
|
|
"""Explicit 2-stage (llm + diffusion) + connectors; sleep/wake both stages."""
|
|
if current_omni_platform.is_rocm():
|
|
clean_gpu_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])
|
|
async for _ in engine.generate("verify", sampling_params_list=[SamplingParams(), sp]):
|
|
pass
|
|
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_pure_diffusion_scenario(tp_size: int):
|
|
"""Single-stage random diffusion: sleep, wake, generate."""
|
|
if current_omni_platform.is_rocm():
|
|
clean_gpu_envs()
|
|
engine_args = {
|
|
"model": MODEL_DIFF,
|
|
"enable_sleep_mode": True,
|
|
"tensor_parallel_size": tp_size,
|
|
"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()
|
|
async for _ in engine.generate(
|
|
"test",
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256),
|
|
):
|
|
pass
|
|
logger.info("Pure diffusion OK (TP=%s)", tp_size)
|
|
finally:
|
|
engine.shutdown()
|