diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index a2d0cdfe0..f7f3e4b05 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -305,3 +305,38 @@ steps: hostPath: path: /mnt/hf-cache type: DirectoryOrCreate + + - label: "Bagel Text2Img Model Test with H100" + timeout_in_minutes: 30 + depends_on: image-build + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate diff --git a/pyproject.toml b/pyproject.toml index 3e1149099..e1fafa606 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,8 @@ dev = [ "psutil>=7.2.0", "soundfile>=0.13.1", "imageio[ffmpeg]>=0.6.0", - "opencv-python>=4.12.0.88" + "opencv-python>=4.12.0.88", + "mooncake-transfer-engine==0.3.8.post1" ] docs = [ diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml new file mode 100644 index 000000000..4c6a67e8b --- /dev/null +++ b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml @@ -0,0 +1,85 @@ +# stage config for running BAGEL with Mooncake connector for CI e2e tests. +# This config is optimized for single GPU tests with Mooncake inter-stage communication. + +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: BagelForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.60 + enforce_eager: true + trust_remote_code: true + engine_output_type: text + distributed_executor_backend: mp + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 1 + omni_kv_config: + need_send_cache: true + kv_transfer_criteria: + type: prefill_finished + final_output: true + final_output_type: text + is_comprehension: true + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 2048 + seed: 52 + detokenize: true + repetition_penalty: 1.05 + output_connectors: + to_stage_1: mooncake_connector + - stage_id: 1 + stage_type: diffusion + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: dit + gpu_memory_utilization: 0.4 + enforce_eager: true + trust_remote_code: true + engine_output_type: image + distributed_executor_backend: mp + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 1 + omni_kv_config: + need_recv_cache: true + engine_input_source: [0] + final_output: true + final_output_type: image + is_comprehension: false + default_sampling_params: + seed: 52 + input_connectors: + from_stage_0: mooncake_connector + +# Top-level runtime config with Mooncake connector +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + connectors: + mooncake_connector: + name: MooncakeConnector + extra: + host: "${MOONCAKE_HOST}" + metadata_server: "http://${MOONCAKE_HOST}:${MOONCAKE_HTTP_PORT}/metadata" + master: "${MOONCAKE_HOST}:${MOONCAKE_RPC_PORT}" + segment: 64000000 + localbuf: 64000000 + proto: tcp + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml new file mode 100644 index 000000000..67ff7309f --- /dev/null +++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml @@ -0,0 +1,83 @@ +# stage config for running BAGEL with SharedMemory connector for CI e2e tests. +# This config is optimized for single GPU tests with SharedMemory inter-stage communication. + +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: BagelForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.65 + enforce_eager: true + trust_remote_code: true + engine_output_type: text + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 1 + omni_kv_config: + need_send_cache: true + kv_transfer_criteria: + type: prefill_finished #or special token generated + final_output: true + final_output_type: text + is_comprehension: true + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 2048 + seed: 52 + detokenize: True + repetition_penalty: 1.05 + + - stage_id: 1 + stage_type: diffusion + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: dit + gpu_memory_utilization: 0.4 + enforce_eager: true + trust_remote_code: true + engine_output_type: image + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 1 + omni_kv_config: + need_recv_cache: true + engine_input_source: [0] + + final_output: true + final_output_type: image + is_comprehension: false + default_sampling_params: + seed: 52 + +# Runtime edges +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + + # Distributed connectors configuration (optional) + # More connectors will be supported in the future. + connectors: + shared_memory_connector: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 # 64KB threshold + + + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py new file mode 100644 index 000000000..777af52a7 --- /dev/null +++ b/tests/e2e/offline_inference/test_bagel_text2img.py @@ -0,0 +1,308 @@ +# 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" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +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.utils import hardware_test +from vllm_omni.entrypoints.omni import Omni + +# 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": (68, 107, 134)}, + {"position": (400, 50), "rgb": (95, 139, 166)}, + {"position": (700, 100), "rgb": (99, 122, 151)}, + {"position": (150, 400), "rgb": (111, 125, 153)}, + {"position": (512, 512), "rgb": (97, 107, 131)}, + {"position": (700, 400), "rgb": (48, 64, 98)}, + {"position": (100, 700), "rgb": (79, 63, 84)}, + {"position": (400, 700), "rgb": (40, 58, 79)}, + {"position": (700, 700), "rgb": (60, 75, 103)}, + {"position": (256, 256), "rgb": (97, 128, 156)}, +] + +# Maximum allowed difference per color channel +PIXEL_TOLERANCE = 5 + +# Default test prompt +DEFAULT_PROMPT = "<|im_start|>A futuristic city skyline at twilight, cyberpunk style<|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, max_tokens: int = 1, num_inference_steps: int = 15) -> list: + """Configure sampling parameters for Bagel text2img generation. + + Args: + omni: The Omni instance to get default params from. + max_tokens: Maximum tokens for the first stage. + num_inference_steps: Number of inference steps for the diffusion stage. + + Returns: + Configured sampling params list. + """ + params_list = omni.default_sampling_params_list + params_list[0].max_tokens = max_tokens # type: ignore + if len(params_list) > 1: + params_list[1].num_inference_steps = num_inference_steps # type: ignore + 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: + for stage_out in 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 + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +def test_bagel_text2img_shared_memory_connector(): + """Test Bagel text2img with shared memory connector.""" + config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") + omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path) + + try: + generated_image = _generate_bagel_image(omni) + _validate_pixels(generated_image) + finally: + omni.close() + + +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 _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.diffusion +@hardware_test(res={"cuda": "H100"}) +def test_bagel_text2img_mooncake_connector(): + """Test Bagel text2img with Mooncake connector for inter-stage communication.""" + 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 + omni = 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, + ) + + omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file) + + generated_image = _generate_bagel_image(omni) + _validate_pixels(generated_image) + + finally: + if omni: + omni.close() + 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 diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml index ee2b17729..c60f3b700 100644 --- a/vllm_omni/model_executor/stage_configs/bagel.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel.yaml @@ -11,7 +11,7 @@ stage_args: model_arch: BagelForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.4 + gpu_memory_utilization: 0.35 enforce_eager: true trust_remote_code: true engine_output_type: text @@ -42,7 +42,7 @@ stage_args: max_batch_size: 1 engine_args: model_stage: dit - gpu_memory_utilization: 0.4 + gpu_memory_utilization: 0.55 enforce_eager: true trust_remote_code: true engine_output_type: image diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml index eddc844d5..c19302c76 100644 --- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml @@ -11,7 +11,7 @@ stage_args: model_arch: BagelForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.4 + gpu_memory_utilization: 0.35 enforce_eager: true trust_remote_code: true engine_output_type: text @@ -41,11 +41,11 @@ stage_args: - stage_id: 1 stage_type: diffusion runtime: - devices: "1" + devices: "0" max_batch_size: 1 engine_args: model_stage: dit - gpu_memory_utilization: 0.4 + gpu_memory_utilization: 0.55 enforce_eager: true trust_remote_code: true engine_output_type: image