vllm-omni/tests/e2e/online_serving/test_bagel_online.py

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

"""
End-to-end online serving test for Bagel text2img and img2img generation.

This test validates that the Bagel model can serve image generation requests
via the OpenAI-compatible chat completions API.

Equivalent to running:
    vllm-omni serve "ByteDance-Seed/BAGEL-7B-MoT" --omni --port 8091

    # text2img
    python3 examples/online_serving/bagel/openai_chat_client.py \\
        --prompt "A cute cat" --modality text2img

    # img2img
    python3 examples/online_serving/bagel/openai_chat_client.py \\
        --prompt "Let the woman wear a blue dress" --modality img2img \\
        --image-url women.jpg
"""

import base64
import os
import signal
import socket
import subprocess
import sys
import time
from io import BytesIO
from pathlib import Path
from typing import Any

import pytest
import requests
from PIL import Image
from vllm.assets.image import ImageAsset

from tests.utils import hardware_test

MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
STAGE_CONFIGS_PATH = str(
    Path(__file__).parent.parent / "offline_inference" / "stage_configs" / "bagel_sharedmemory_ci.yaml"
)

TEXT2IMG_PROMPT = "A cute cat"
IMG2IMG_PROMPT = "Change the grass color to red"


class BagelOmniServer:
    """Context manager to start/stop a vLLM-Omni server for Bagel model tests."""

    def __init__(
        self,
        model: str = MODEL,
        stage_configs_path: str = STAGE_CONFIGS_PATH,
        env_dict: dict[str, str] | None = None,
    ) -> None:
        self.model = model
        self.stage_configs_path = stage_configs_path
        self.env_dict = env_dict
        self.proc: subprocess.Popen | None = None
        self.host = "127.0.0.1"
        self.port = _find_free_port()

    @property
    def base_url(self) -> str:
        return f"http://{self.host}:{self.port}"

    def _start_server(self) -> None:
        env = os.environ.copy()
        if self.env_dict is not None:
            env.update(self.env_dict)

        cmd = [
            sys.executable,
            "-m",
            "vllm_omni.entrypoints.cli.main",
            "serve",
            self.model,
            "--omni",
            "--host",
            self.host,
            "--port",
            str(self.port),
            "--stage-configs-path",
            self.stage_configs_path,
            "--stage-init-timeout",
            "300",
        ]

        self.proc = subprocess.Popen(
            cmd,
            env=env,
            start_new_session=True,
        )

        try:
            if not _wait_for_port(self.host, self.port, timeout=600, proc=self.proc):
                self.terminate()
                raise RuntimeError(f"Server failed to start within 600 seconds on {self.host}:{self.port}")
        except Exception:
            self.terminate()
            raise

    def __enter__(self):
        self._start_server()
        return self

    def terminate(self) -> None:
        if self.proc:
            try:
                os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
            except ProcessLookupError:
                pass
            try:
                self.proc.wait(timeout=30)
            except subprocess.TimeoutExpired:
                try:
                    os.killpg(os.getpgid(self.proc.pid), signal.SIGKILL)
                except ProcessLookupError:
                    pass
                self.proc.wait()
            self.proc = None

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.terminate()


def _find_free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        s.listen(1)
        return s.getsockname()[1]


def _wait_for_port(host: str, port: int, timeout: int = 600, proc: subprocess.Popen | None = None) -> bool:
    start = time.time()
    while time.time() - start < timeout:
        if proc is not None and proc.poll() is not None:
            # Server process exited early
            return False
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
                sock.settimeout(1)
                if sock.connect_ex((host, port)) == 0:
                    return True
        except Exception:
            pass
        time.sleep(2)
    return False


def _send_chat_request(
    server_url: str,
    prompt: str,
    *,
    modality: str = "text2img",
    image: Image.Image | None = None,
    timeout: int = 300,
) -> dict[str, Any]:
    """Send a chat completion request matching the openai_chat_client.py format."""
    content: list[dict[str, Any]] = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]

    if image is not None:
        buffer = BytesIO()
        image.save(buffer, format="JPEG")
        b64_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
        content.append(
            {
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{b64_data}"},
            }
        )

    payload: dict[str, Any] = {
        "messages": [{"role": "user", "content": content}],
    }

    if modality in ("text2img", "img2img"):
        payload["modalities"] = ["image"]

    resp = requests.post(
        f"{server_url}/v1/chat/completions",
        headers={"Content-Type": "application/json"},
        json=payload,
        timeout=timeout,
    )
    resp.raise_for_status()
    return resp.json()


def _extract_image_from_response(data: dict[str, Any]) -> Image.Image | None:
    """Extract the generated PIL Image from a chat completion response."""
    for choice in data.get("choices", []):
        content = choice.get("message", {}).get("content")
        if isinstance(content, list) and content:
            first_item = content[0]
            if isinstance(first_item, dict) and "image_url" in first_item:
                url = first_item["image_url"].get("url", "")
                if url.startswith("data:image"):
                    _, b64 = url.split(",", 1)
                    return Image.open(BytesIO(base64.b64decode(b64)))
    return None


@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
def test_bagel_text2img_online():
    """Test Bagel text2img via OpenAI-compatible chat completions API."""
    with BagelOmniServer() as server:
        response_data = _send_chat_request(
            server.base_url,
            TEXT2IMG_PROMPT,
            modality="text2img",
        )

        image = _extract_image_from_response(response_data)
        assert image is not None, f"No image in response: {response_data}"
        image.load()

        w, h = image.size
        assert w > 0 and h > 0, f"Invalid image size: {image.size}"


@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
def test_bagel_img2img_online():
    """Test Bagel img2img via OpenAI-compatible chat completions API."""
    input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")

    with BagelOmniServer() as server:
        response_data = _send_chat_request(
            server.base_url,
            IMG2IMG_PROMPT,
            modality="img2img",
            image=input_image,
        )

        image = _extract_image_from_response(response_data)
        assert image is not None, f"No image in response: {response_data}"
        image.load()

        w, h = image.size
        assert w > 0 and h > 0, f"Invalid image size: {image.size}"