diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 600a5d4db..ae2e7f4c4 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -12,6 +12,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -29,6 +30,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -46,6 +48,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -64,6 +67,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -81,6 +85,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -98,6 +103,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -116,6 +122,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -134,6 +141,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -152,6 +160,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -178,6 +187,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -197,6 +207,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -226,6 +237,11 @@ steps: env: - name: HF_HOME value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token nodeSelector: node.kubernetes.io/instance-type: gpu-h100-sxm volumes: @@ -261,6 +277,11 @@ steps: env: - name: HF_HOME value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token nodeSelector: node.kubernetes.io/instance-type: gpu-h100-sxm volumes: @@ -272,37 +293,50 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - # - label: "Bagel Text2Img Model Test with H100" - # timeout_in_minutes: 30 - # depends_on: upload-merge-pipeline - # 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: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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 + - label: "Bagel Model Test with H100 (Real Weights)" + timeout_in_minutes: 60 + depends_on: upload-merge-pipeline + commands: + - | + timeout 55m bash -c ' + set -e + export VLLM_WORKER_MULTIPROC_METHOD=spawn + export VLLM_TEST_CLEAN_GPU_MEMORY=1 + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" + pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" + pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" + ' + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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 + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + 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/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 7c8c4a89e..53a4f180c 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -30,6 +30,11 @@ steps: env: - name: HF_HOME value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token nodeSelector: node.kubernetes.io/instance-type: gpu-h100-sxm volumes: @@ -58,6 +63,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -86,6 +92,11 @@ steps: env: - name: HF_HOME value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token nodeSelector: node.kubernetes.io/instance-type: gpu-h100-sxm volumes: @@ -116,6 +127,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -149,6 +161,11 @@ steps: env: - name: HF_HOME value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token nodeSelector: node.kubernetes.io/instance-type: gpu-h100-sxm volumes: diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 854c2eead..b922a868b 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -12,6 +12,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -28,6 +29,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -44,6 +46,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -60,6 +63,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -76,6 +80,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -93,6 +98,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -110,6 +116,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -137,6 +144,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -159,6 +167,7 @@ steps: propagate-environment: true environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -221,6 +230,7 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" volumes: - "/fsx/hf_cache:/fsx/hf_cache" @@ -268,7 +278,7 @@ steps: timeout 30m bash -c ' export VLLM_WORKER_MULTIPROC_METHOD=spawn export VLLM_TEST_CLEAN_GPU_MEMORY=1 - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py + pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" ' agents: queue: "mithril-h100-pool" @@ -288,6 +298,11 @@ steps: env: - name: HF_HOME value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token nodeSelector: node.kubernetes.io/instance-type: gpu-h100-sxm volumes: @@ -306,7 +321,7 @@ steps: timeout 30m bash -c ' export VLLM_WORKER_MULTIPROC_METHOD=spawn export VLLM_TEST_CLEAN_GPU_MEMORY=1 - pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py + pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" ' agents: queue: "mithril-h100-pool" @@ -326,6 +341,11 @@ steps: env: - name: HF_HOME value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token nodeSelector: node.kubernetes.io/instance-type: gpu-h100-sxm volumes: @@ -345,7 +365,7 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn export VLLM_TEST_CLEAN_GPU_MEMORY=1 export VLLM_IMAGE_FETCH_TIMEOUT=60 - pytest -s -v tests/e2e/online_serving/test_bagel_online.py + pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" ' agents: queue: "mithril-h100-pool" @@ -365,6 +385,11 @@ steps: env: - name: HF_HOME value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token nodeSelector: node.kubernetes.io/instance-type: gpu-h100-sxm volumes: diff --git a/tests/conftest.py b/tests/conftest.py index 51a5e9924..75cfaff98 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1236,15 +1236,14 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st port = params.port stage_config_path = params.stage_config_path if run_level == "advanced_model" and stage_config_path is not None: + # Dynamically detect stages from config to avoid KeyError + # for models with fewer stages (e.g., BAGEL has 2, Qwen3-Omni has 3) + with open(stage_config_path, encoding="utf-8") as f: + _cfg = yaml.safe_load(f) or {} + _stage_ids = [s["stage_id"] for s in _cfg.get("stage_args", []) if "stage_id" in s] stage_config_path = modify_stage_config( stage_config_path, - deletes={ - "stage_args": { - 0: ["engine_args.load_format"], - 1: ["engine_args.load_format"], - 2: ["engine_args.load_format"], - } - }, + deletes={"stage_args": {sid: ["engine_args.load_format"] for sid in _stage_ids}}, ) server_args = params.server_args or [] diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml index dbb93344b..7f3a3a6f4 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml @@ -21,6 +21,7 @@ stage_args: enable_prefix_caching: false max_num_batched_tokens: 32768 tensor_parallel_size: 1 + load_format: dummy omni_kv_config: need_send_cache: true kv_transfer_criteria: @@ -54,6 +55,7 @@ stage_args: enable_prefix_caching: false max_num_batched_tokens: 32768 tensor_parallel_size: 1 + load_format: dummy omni_kv_config: need_recv_cache: true engine_input_source: [0] diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml index 721c50248..aa3cc7718 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml @@ -21,6 +21,7 @@ stage_args: enable_prefix_caching: false max_num_batched_tokens: 32768 tensor_parallel_size: 1 + load_format: dummy omni_kv_config: need_send_cache: true kv_transfer_criteria: @@ -53,6 +54,7 @@ stage_args: enable_prefix_caching: false max_num_batched_tokens: 32768 tensor_parallel_size: 1 + load_format: dummy omni_kv_config: need_recv_cache: true engine_input_source: [0] diff --git a/tests/e2e/offline_inference/test_bagel_img2img.py b/tests/e2e/offline_inference/test_bagel_img2img.py index da9df2778..8c734c6a2 100644 --- a/tests/e2e/offline_inference/test_bagel_img2img.py +++ b/tests/e2e/offline_inference/test_bagel_img2img.py @@ -22,6 +22,7 @@ import pytest from PIL import Image from vllm.assets.image import ImageAsset +from tests.conftest import modify_stage_config from tests.utils import hardware_test from vllm_omni.entrypoints.omni import Omni @@ -168,17 +169,39 @@ def _generate_bagel_img2img( 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"}) -def test_bagel_img2img_shared_memory_connector(): +def test_bagel_img2img_shared_memory_connector(run_level): """Test Bagel img2img with shared memory connector.""" input_image = _load_input_image() config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") + config_path = _resolve_stage_config(config_path, run_level) omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300) try: generated_image = _generate_bagel_img2img(omni, input_image) - _validate_pixels(generated_image) + if run_level == "advanced_model": + _validate_pixels(generated_image) finally: omni.close() diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py index 7990ac980..ed369aedd 100644 --- a/tests/e2e/offline_inference/test_bagel_text2img.py +++ b/tests/e2e/offline_inference/test_bagel_text2img.py @@ -28,6 +28,7 @@ from typing import Any import pytest from PIL import Image +from tests.conftest import modify_stage_config from tests.utils import hardware_test from vllm_omni.entrypoints.omni import Omni @@ -158,17 +159,39 @@ def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Ima 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"}) -def test_bagel_text2img_shared_memory_connector(): +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) omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300) try: generated_image = _generate_bagel_image(omni) - _validate_pixels(generated_image) + if run_level == "advanced_model": + _validate_pixels(generated_image) finally: omni.close() @@ -251,9 +274,10 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str: @pytest.mark.core_model +@pytest.mark.advanced_model @pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}) -def test_bagel_text2img_mooncake_connector(): +def test_bagel_text2img_mooncake_connector(run_level): """Test Bagel text2img with Mooncake connector for inter-stage communication.""" MOONCAKE_HOST = "127.0.0.1" MOONCAKE_RPC_PORT = _find_free_port() @@ -291,10 +315,12 @@ def test_bagel_text2img_mooncake_connector(): http_port=MOONCAKE_HTTP_PORT, ) + temp_config_file = _resolve_stage_config(temp_config_file, run_level) omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300) generated_image = _generate_bagel_image(omni) - _validate_pixels(generated_image) + if run_level == "advanced_model": + _validate_pixels(generated_image) finally: if omni: diff --git a/tests/e2e/online_serving/test_bagel_online.py b/tests/e2e/online_serving/test_bagel_online.py index 4056cfdef..a5e26db1e 100644 --- a/tests/e2e/online_serving/test_bagel_online.py +++ b/tests/e2e/online_serving/test_bagel_online.py @@ -22,22 +22,18 @@ Equivalent to running: 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.conftest import OmniServerParams from tests.utils import hardware_test +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + MODEL = "ByteDance-Seed/BAGEL-7B-MoT" STAGE_CONFIGS_PATH = str( Path(__file__).parent.parent / "offline_inference" / "stage_configs" / "bagel_sharedmemory_ci.yaml" @@ -46,202 +42,74 @@ STAGE_CONFIGS_PATH = str( 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() +# Create parameter combinations for model and stage config +test_params = [ + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIGS_PATH, + server_args=["--stage-init-timeout", "300"], + ), +] -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 _build_text2img_messages(prompt: str) -> list[dict]: + """Build OpenAI-format messages for text2img generation.""" + return [ + { + "role": "user", + "content": [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}], + } + ] -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 _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]: + """Build OpenAI-format messages for img2img generation.""" + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}, + }, + ], + } + ] -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}], +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_bagel_text2img_online(omni_server, openai_client) -> None: + """Test Bagel text2img via OpenAI-compatible chat completions API.""" + request_config = { + "model": omni_server.model, + "messages": _build_text2img_messages(TEXT2IMG_PROMPT), + "modalities": ["image"], } - 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 + openai_client.send_diffusion_request(request_config) @pytest.mark.core_model +@pytest.mark.advanced_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(): +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_bagel_img2img_online(omni_server, openai_client) -> None: """Test Bagel img2img via OpenAI-compatible chat completions API.""" input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") + buffer = BytesIO() + input_image.save(buffer, format="JPEG") + image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8") - with BagelOmniServer() as server: - response_data = _send_chat_request( - server.base_url, - IMG2IMG_PROMPT, - modality="img2img", - image=input_image, - ) + request_config = { + "model": omni_server.model, + "messages": _build_img2img_messages(IMG2IMG_PROMPT, image_b64), + "modalities": ["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}" + openai_client.send_diffusion_request(request_config) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index 6b041aece..ee2f471f5 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -95,6 +95,9 @@ def get_qwen_image_layered_pre_process_func( else: image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image) + if isinstance(image, PIL.Image.Image) and image.mode != "RGBA": + image = image.convert("RGBA") + # 1. calculate dimensions image_size = image.size assert request.sampling_params.resolution in [640, 1024], ( @@ -652,6 +655,8 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as width = req.sampling_params.width else: # fallback to run pre-processing in pipeline (debug only) + if isinstance(image, PIL.Image.Image) and image.mode != "RGBA": + image = image.convert("RGBA") image_size = image[0].size if isinstance(image, list) else image.size assert resolution in [640, 1024], f"resolution must be either 640 or 1024, but got {resolution}" calculated_width, calculated_height = calculate_dimensions( diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index a6005290c..56a770461 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -28,17 +28,23 @@ class OmniDiffusionRequest: def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" + # Detect whether user explicitly provided guidance_scale. + # The sentinel default is 0.0 (false-like); any truthy value means + # the caller set it intentionally. We must resolve this BEFORE + # auto-filling guidance_scale_2, otherwise the sentinel leaks into + # guidance_scale_2. + if self.sampling_params.guidance_scale: + self.sampling_params.guidance_scale_provided = True + else: + self.sampling_params.guidance_scale = 1.0 + # Set do_classifier_free_guidance based on guidance scale and negative prompt if self.sampling_params.guidance_scale > 1.0 and any( (not isinstance(p, str) and p.get("negative_prompt")) for p in self.prompts ): self.sampling_params.do_classifier_free_guidance = True + + # Auto-fill guidance_scale_2 from the (now-resolved) guidance_scale + # so downstream code always has a valid value. if self.sampling_params.guidance_scale_2 is None: self.sampling_params.guidance_scale_2 = self.sampling_params.guidance_scale - - # The dataclass default value is 0 (false-like), used to detect whether user explicitly provides this value - # After this check is done, reset this value to old default 1 - if self.sampling_params.guidance_scale: - self.sampling_params.guidance_scale_provided = True - else: - self.sampling_params.guidance_scale = 1.0 diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index accb173e1..972c95c29 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -221,6 +221,7 @@ class DiffusionModelRunner: not getattr(req, "skip_cache_refresh", False) and self.cache_backend is not None and self.cache_backend.is_enabled() + and req.sampling_params.num_inference_steps is not None ): self.cache_backend.refresh(self.pipeline, req.sampling_params.num_inference_steps) diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 08812223d..52bd6031c 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/vllm_omni/entrypoints/async_omni_diffusion.py @@ -172,9 +172,6 @@ class AsyncOmniDiffusion: if request_id is None: request_id = f"diff-{uuid.uuid4().hex[:16]}" - if sampling_params.guidance_scale: - sampling_params.guidance_scale_provided = True - if lora_request is not None: sampling_params.lora_request = lora_request diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 68ec4f404..d58617e17 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2052,20 +2052,26 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): except ValueError: logger.warning("Invalid size format: %s", extra_body.get("size")) - # Get request parameters from extra_body - # Text-to-image parameters (ref: text_to_image.py) - num_inference_steps = extra_body.get("num_inference_steps", 50) + # Get request parameters from extra_body. + # Avoid hardcoded defaults here — let each pipeline's forward() + # method apply its own model-specific default when the user does + # not provide a value. + num_inference_steps = extra_body.get("num_inference_steps") guidance_scale = extra_body.get("guidance_scale") - true_cfg_scale = extra_body.get("true_cfg_scale") # Qwen-Image specific + true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale") seed = extra_body.get("seed") negative_prompt = extra_body.get("negative_prompt") num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1) # Text-to-video parameters (ref: text_to_video.py) num_frames = extra_body.get("num_frames") - guidance_scale_2 = extra_body.get("guidance_scale_2") # For video high-noise CFG + guidance_scale_2 = extra_body.get("guidance_scale_2") lora_body = extra_body.get("lora") + # Qwen-Image-Layered parameters + layers = extra_body.get("layers") + resolution = extra_body.get("resolution") + logger.info( "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s", request_id, @@ -2089,25 +2095,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "negative_prompt": negative_prompt, } gen_params = OmniDiffusionSamplingParams( - num_inference_steps=num_inference_steps, height=height, width=width, num_outputs_per_prompt=num_outputs_per_prompt, seed=seed, ) + # Only override defaults when the user explicitly provides values + if num_inference_steps is not None: + gen_params.num_inference_steps = num_inference_steps if guidance_scale is not None: gen_params.guidance_scale = guidance_scale - - # Add Qwen-Image specific parameter if true_cfg_scale is not None: gen_params.true_cfg_scale = true_cfg_scale - - # Add video generation parameters if set if num_frames is not None: gen_params.num_frames = num_frames if guidance_scale_2 is not None: gen_params.guidance_scale_2 = guidance_scale_2 + if layers is not None: + gen_params.layers = layers + if resolution is not None: + gen_params.resolution = resolution # Parse per-request LoRA (works for both AsyncOmniDiffusion and AsyncOmni). if lora_body and isinstance(lora_body, dict): diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index c0f10af2b..5768c3b6d 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -234,8 +234,9 @@ class OmniDiffusionSamplingParams: step_index: int | None = None boundary_ratio: float | None = None - # Scheduler parameters - num_inference_steps: int = 50 + # Scheduler parameters – ``None`` means "not explicitly set by the caller"; + # each pipeline's ``forward()`` decides its own model-specific default. + num_inference_steps: int | None = None guidance_scale: float = 0.0 guidance_scale_provided: bool = False guidance_scale_2: float | None = None diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index bdcad535a..de248f0f3 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -396,6 +396,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # Keys that should stay on GPU in model_intermediate_buffer to avoid # CPU-to-GPU round-trips on every decode step. self.gpu_resident_buffer_keys: set[str] = { + "audio_codes", "last_talker_hidden", "tts_pad_embed", "tailing_text_hidden", @@ -650,7 +651,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # Stays on GPU - gpu_resident_buffer_keys avoids the CPU round-trip. if hidden_states.numel() == 0: return {} - return {"last_talker_hidden": hidden_states[-1, :].detach()} + last = hidden_states[-1, :].detach() + return {"last_talker_hidden": last} # -------------------- prompt construction helpers -------------------- diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 166cc6231..ecab7af05 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -560,6 +560,29 @@ class GPUARModelRunner(OmniGPUModelRunner): hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output ) + # Pre-copy multimodal tensors to CPU once (not per-request) to avoid + # redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU. + mm_cpu: dict[str, object] = {} + if isinstance(multimodal_outputs, dict) and multimodal_outputs: + for k, v in multimodal_outputs.items(): + try: + if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: + mm_cpu[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, dict): + sub_dict: dict[str, torch.Tensor] = {} + for sk, sv in v.items(): + if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]: + sub_dict[str(sk)] = sv.detach().to("cpu").contiguous() + if sub_dict: + mm_cpu[k] = sub_dict + elif isinstance(v, list): + element = v[0] + if isinstance(element, torch.Tensor): + element = element.detach().to("cpu").contiguous() + mm_cpu[k] = element + except Exception as e: + logger.error(f"Error in merge multimodal outputs: {e}") + pooler_output: list[dict[str, object]] = [] for rid in req_ids_output_copy: idx = req_id_to_index_output_copy[rid] @@ -568,28 +591,20 @@ class GPUARModelRunner(OmniGPUModelRunner): end = start + sched hidden_slice = hidden_states_cpu[start:end] payload: dict[str, object] = {"hidden": hidden_slice} - if isinstance(multimodal_outputs, dict) and multimodal_outputs: + if mm_cpu: mm_payload: dict[str, object] = {} - for k, v in multimodal_outputs.items(): - try: - if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: - mm_payload[k] = v.detach().to("cpu")[start:end].contiguous() - elif isinstance(v, dict): - sub_dict: dict[str, torch.Tensor] = {} - for sk, sv in v.items(): - if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]: - sub_dict[str(sk)] = sv.detach().to("cpu")[start:end].contiguous() - if sub_dict: - mm_payload[k] = sub_dict - elif isinstance(v, list): - element = v[0] - if isinstance(element, torch.Tensor): - element = element.detach().to("cpu").contiguous() - mm_payload[k] = element - except Exception as e: - logger.error(f"Error in merge multimodal outputs: {e}") - if mm_payload: - payload.update(mm_payload) + for k, v in mm_cpu.items(): + if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: + mm_payload[k] = v[start:end].contiguous() + elif isinstance(v, dict): + mm_payload[k] = {sk: sv[start:end].contiguous() for sk, sv in v.items()} + elif isinstance(v, torch.Tensor): + # List-derived tensor payloads are request-invariant; clone to + # avoid accidental cross-request aliasing on downstream mutation. + mm_payload[k] = v.clone() + else: + mm_payload[k] = v + payload.update(mm_payload) pooler_output.append(payload) with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): if self.routed_experts_initialized: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 1b1d3d9cf..c77fc7d69 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1313,14 +1313,16 @@ class OmniGPUModelRunner(GPUModelRunner): None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc ): req_embeds, code_predictor_codes = self.talker_mtp(req_input_ids, req_embeds, last_talker_hidden, text_step) - # update the inputs_embeds and code_predictor_codes - code_predictor_codes_cpu = code_predictor_codes.detach().to("cpu").contiguous() + # code_predictor_codes stays on GPU here; _update_intermediate_buffer + # keeps it device-resident when the key is in gpu_resident_buffer_keys. + # D2H is deferred to sample_tokens where hidden_states.to("cpu") already + # syncs the stream, avoiding a per-step cudaStreamSynchronize. out_key = getattr(self.model, "talker_mtp_output_key", "code_predictor_codes") for idx, req_id in enumerate(decode_req_ids): req_index = self.input_batch.req_ids.index(req_id) start_offset = int(self.query_start_loc.cpu[req_index]) inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1] - update_dict = {out_key: code_predictor_codes_cpu[idx : idx + 1]} + update_dict = {out_key: code_predictor_codes[idx : idx + 1]} self._merge_additional_information_update(req_id, update_dict) def _model_forward(