369 lines
12 KiB
Python
369 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
End-to-end test for Bagel text2img generation.
|
|
|
|
This test validates that the Bagel model generates images that match
|
|
expected reference pixel values within a ±5 tolerance.
|
|
|
|
Equivalent to running:
|
|
python3 examples/offline_inference/bagel/end2end.py \
|
|
--prompts "A futuristic city skyline at twilight, cyberpunk style" \
|
|
--modality text2img --step 15
|
|
"""
|
|
|
|
import os
|
|
|
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
|
import signal
|
|
import socket
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from tests.conftest import OmniRunner, modify_stage_config
|
|
from tests.utils import hardware_test
|
|
from vllm_omni import Omni
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
# Reference pixel data extracted from the known-good output image
|
|
# Each entry contains (x, y) position and expected (R, G, B) values
|
|
# "Generated with seed=52, num_inference_steps=15,
|
|
# prompt='A futuristic city skyline at twilight, cyberpunk style'"
|
|
REFERENCE_PIXELS = [
|
|
{"position": (100, 100), "rgb": (121, 118, 100)},
|
|
{"position": (400, 50), "rgb": (163, 162, 143)},
|
|
{"position": (700, 100), "rgb": (170, 156, 127)},
|
|
{"position": (150, 400), "rgb": (129, 127, 112)},
|
|
{"position": (512, 512), "rgb": (135, 61, 59)},
|
|
{"position": (700, 400), "rgb": (205, 107, 43)},
|
|
{"position": (100, 700), "rgb": (197, 177, 157)},
|
|
{"position": (400, 700), "rgb": (139, 107, 86)},
|
|
{"position": (700, 700), "rgb": (247, 205, 146)},
|
|
{"position": (256, 256), "rgb": (171, 160, 153)},
|
|
]
|
|
|
|
if current_omni_platform.is_rocm():
|
|
REFERENCE_PIXELS = [
|
|
{"position": (100, 100), "rgb": (123, 119, 100)},
|
|
{"position": (400, 50), "rgb": (162, 161, 142)},
|
|
{"position": (700, 100), "rgb": (171, 156, 127)},
|
|
{"position": (150, 400), "rgb": (131, 128, 112)},
|
|
{"position": (512, 512), "rgb": (134, 61, 59)},
|
|
{"position": (700, 400), "rgb": (204, 107, 43)},
|
|
{"position": (100, 700), "rgb": (201, 180, 165)},
|
|
{"position": (400, 700), "rgb": (140, 108, 87)},
|
|
{"position": (700, 700), "rgb": (247, 205, 145)},
|
|
{"position": (256, 256), "rgb": (171, 160, 153)},
|
|
]
|
|
|
|
# Maximum allowed difference per color channel
|
|
PIXEL_TOLERANCE = 5
|
|
|
|
# Default test prompt
|
|
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
|
|
|
|
|
|
def _find_free_port() -> int:
|
|
"""Find and return a free ephemeral port by binding to port 0."""
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.bind(("127.0.0.1", 0))
|
|
s.listen(1)
|
|
port = s.getsockname()[1]
|
|
return port
|
|
|
|
|
|
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
|
|
"""Configure sampling parameters for Bagel text2img generation.
|
|
|
|
Args:
|
|
omni: The Omni instance to get default params from.
|
|
num_inference_steps: Number of inference steps for the diffusion stage.
|
|
|
|
Returns:
|
|
Configured sampling params list.
|
|
"""
|
|
params_list = omni.default_sampling_params_list
|
|
if len(params_list) > 1:
|
|
params_list[1].num_inference_steps = num_inference_steps # type: ignore
|
|
params_list[1].extra_args = { # type: ignore
|
|
"cfg_text_scale": 4.0,
|
|
"cfg_img_scale": 1.5,
|
|
}
|
|
return params_list
|
|
|
|
|
|
def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
|
|
"""Extract the generated image from Omni outputs.
|
|
|
|
Args:
|
|
omni_outputs: List of outputs from omni.generate().
|
|
|
|
Returns:
|
|
The first generated PIL Image, or None if no image found.
|
|
"""
|
|
for req_output in omni_outputs:
|
|
if images := getattr(req_output, "images", None):
|
|
return images[0]
|
|
if hasattr(req_output, "request_output") and req_output.request_output:
|
|
stage_out = req_output.request_output
|
|
if hasattr(stage_out, "images") and stage_out.images:
|
|
return stage_out.images[0]
|
|
return None
|
|
|
|
|
|
def _validate_pixels(
|
|
image: Image.Image,
|
|
reference_pixels: list[dict[str, Any]] = REFERENCE_PIXELS,
|
|
tolerance: int = PIXEL_TOLERANCE,
|
|
) -> None:
|
|
"""Validate that image pixels match expected reference values.
|
|
|
|
Args:
|
|
image: The PIL Image to validate.
|
|
reference_pixels: List of dicts with 'position' (x, y) and 'rgb' (R, G, B).
|
|
tolerance: Maximum allowed difference per color channel.
|
|
|
|
Raises:
|
|
AssertionError: If any pixel differs beyond tolerance.
|
|
"""
|
|
for ref in reference_pixels:
|
|
x, y = ref["position"]
|
|
expected = ref["rgb"]
|
|
actual = image.getpixel((x, y))[:3]
|
|
assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), (
|
|
f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}"
|
|
)
|
|
|
|
|
|
def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Image:
|
|
"""Generate an image using Bagel model with configured parameters.
|
|
|
|
Args:
|
|
omni: The Omni instance to use for generation.
|
|
prompt: The text prompt for image generation.
|
|
|
|
Returns:
|
|
The generated PIL Image.
|
|
|
|
Raises:
|
|
AssertionError: If no image is generated or size is incorrect.
|
|
"""
|
|
params_list = _configure_sampling_params(omni)
|
|
|
|
omni_outputs = list(
|
|
omni.generate(
|
|
prompts=[{"prompt": prompt, "modalities": ["image"]}],
|
|
sampling_params_list=params_list,
|
|
)
|
|
)
|
|
|
|
generated_image = _extract_generated_image(omni_outputs)
|
|
assert generated_image is not None, "No images generated"
|
|
assert generated_image.size == (1024, 1024), f"Expected 1024x1024, got {generated_image.size}"
|
|
|
|
return generated_image
|
|
|
|
|
|
def _resolve_stage_config(config_path: str, run_level: str) -> str:
|
|
"""Resolve stage config based on run level.
|
|
|
|
For advanced_model (real weights), strip load_format: dummy so the model
|
|
falls back to loading real weights from HuggingFace.
|
|
"""
|
|
if run_level == "advanced_model":
|
|
return modify_stage_config(
|
|
config_path,
|
|
deletes={
|
|
"stage_args": {
|
|
0: ["engine_args.load_format"],
|
|
1: ["engine_args.load_format"],
|
|
}
|
|
},
|
|
)
|
|
return config_path
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
|
def test_bagel_text2img_shared_memory_connector(run_level):
|
|
"""Test Bagel text2img with shared memory connector."""
|
|
config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
|
|
config_path = _resolve_stage_config(config_path, run_level)
|
|
with OmniRunner(
|
|
"ByteDance-Seed/BAGEL-7B-MoT",
|
|
stage_configs_path=config_path,
|
|
) as runner:
|
|
generated_image = _generate_bagel_image(runner.omni)
|
|
if run_level == "advanced_model":
|
|
_validate_pixels(generated_image)
|
|
|
|
|
|
def _wait_for_port(host: str, port: int, timeout: int = 30) -> bool:
|
|
"""Wait for a port to become available.
|
|
|
|
Args:
|
|
host: The host address.
|
|
port: The port number.
|
|
timeout: Maximum seconds to wait.
|
|
|
|
Returns:
|
|
True if port becomes available, False otherwise.
|
|
"""
|
|
for _ in range(timeout):
|
|
try:
|
|
with socket.create_connection((host, port), timeout=1):
|
|
return True
|
|
except (TimeoutError, ConnectionRefusedError):
|
|
time.sleep(1)
|
|
return False
|
|
|
|
|
|
def _is_mooncake_master_available() -> bool:
|
|
"""Check if mooncake_master binary is present and can actually execute."""
|
|
import shutil
|
|
|
|
binary = shutil.which("mooncake_master")
|
|
if binary is None:
|
|
return False
|
|
try:
|
|
result = subprocess.run(
|
|
[binary, "--help"],
|
|
capture_output=True,
|
|
timeout=5,
|
|
)
|
|
return result.returncode != 127
|
|
except (subprocess.TimeoutExpired, OSError):
|
|
return True
|
|
|
|
|
|
def _cleanup_mooncake_processes(timeout_secs: int = 5) -> None:
|
|
"""Clean up any existing mooncake_master processes.
|
|
|
|
Args:
|
|
timeout_secs: Maximum seconds to wait for graceful termination.
|
|
"""
|
|
subprocess.run(
|
|
["pkill", "-f", "mooncake_master"],
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
)
|
|
start_time = time.time()
|
|
while time.time() - start_time < timeout_secs:
|
|
result = subprocess.run(
|
|
["pgrep", "-f", "mooncake_master"],
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
)
|
|
if result.returncode != 0:
|
|
break
|
|
time.sleep(0.5)
|
|
else:
|
|
subprocess.run(
|
|
["pkill", "-9", "-f", "mooncake_master"],
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
)
|
|
time.sleep(1)
|
|
|
|
|
|
def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
|
|
"""Load Mooncake config from YAML and substitute placeholders.
|
|
|
|
Args:
|
|
host: Mooncake host address.
|
|
rpc_port: RPC port for Mooncake master.
|
|
http_port: HTTP metadata server port.
|
|
|
|
Returns:
|
|
Path to the temporary config file with substituted values.
|
|
"""
|
|
config_path = str(Path(__file__).parent / "stage_configs" / "bagel_mooncake_ci.yaml")
|
|
with open(config_path) as f:
|
|
config_content = f.read()
|
|
|
|
# Substitute placeholders
|
|
config_content = config_content.replace("${MOONCAKE_HOST}", host)
|
|
config_content = config_content.replace("${MOONCAKE_RPC_PORT}", str(rpc_port))
|
|
config_content = config_content.replace("${MOONCAKE_HTTP_PORT}", str(http_port))
|
|
|
|
# Write to temp file
|
|
temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
|
|
temp_file.write(config_content)
|
|
temp_file.close()
|
|
return temp_file.name
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@hardware_test(res={"cuda": "H100"})
|
|
def test_bagel_text2img_mooncake_connector(run_level):
|
|
"""Test Bagel text2img with Mooncake connector for inter-stage communication."""
|
|
if not _is_mooncake_master_available():
|
|
pytest.skip("mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)")
|
|
MOONCAKE_HOST = "127.0.0.1"
|
|
MOONCAKE_RPC_PORT = _find_free_port()
|
|
MOONCAKE_HTTP_PORT = _find_free_port()
|
|
MOONCAKE_METRICS_PORT = _find_free_port()
|
|
|
|
mooncake_master_proc = None
|
|
temp_config_file = None
|
|
|
|
try:
|
|
_cleanup_mooncake_processes()
|
|
|
|
# Start mooncake_master
|
|
mooncake_master_proc = subprocess.Popen(
|
|
[
|
|
"mooncake_master",
|
|
f"--rpc_port={MOONCAKE_RPC_PORT}",
|
|
"--enable_http_metadata_server=true",
|
|
"--http_metadata_server_host=0.0.0.0",
|
|
f"--http_metadata_server_port={MOONCAKE_HTTP_PORT}",
|
|
f"--metrics_port={MOONCAKE_METRICS_PORT}",
|
|
],
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
preexec_fn=os.setsid,
|
|
)
|
|
|
|
assert _wait_for_port(MOONCAKE_HOST, MOONCAKE_RPC_PORT), "mooncake_master failed to start"
|
|
|
|
# Create temp config and initialize Omni
|
|
temp_config_file = _load_mooncake_config(
|
|
host=MOONCAKE_HOST,
|
|
rpc_port=MOONCAKE_RPC_PORT,
|
|
http_port=MOONCAKE_HTTP_PORT,
|
|
)
|
|
|
|
temp_config_file = _resolve_stage_config(temp_config_file, run_level)
|
|
with OmniRunner(
|
|
"ByteDance-Seed/BAGEL-7B-MoT",
|
|
stage_configs_path=temp_config_file,
|
|
stage_init_timeout=300,
|
|
) as runner:
|
|
generated_image = _generate_bagel_image(runner.omni)
|
|
if run_level == "advanced_model":
|
|
_validate_pixels(generated_image)
|
|
|
|
finally:
|
|
if temp_config_file:
|
|
try:
|
|
os.unlink(temp_config_file)
|
|
except OSError:
|
|
pass
|
|
if mooncake_master_proc:
|
|
try:
|
|
os.killpg(os.getpgid(mooncake_master_proc.pid), signal.SIGKILL)
|
|
except OSError:
|
|
pass
|