From 8a9add1c81d87b34800df4def9b763b6a0117c92 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Mon, 20 Apr 2026 14:23:42 +0800 Subject: [PATCH] [CI] Restructure vLLM-Omni Test Layout, Fixture Scope, and Support Modules (#2620) Signed-off-by: wangyu <410167048@qq.com> --- .buildkite/test-nightly.yml | 2 +- .../l4_functionality_tests.inc.md | 2 +- docs/contributing/ci/tests_markers.md | 12 +- docs/contributing/ci/tests_style.md | 14 +- tests/conftest.py | 3580 +---------------- tests/core/test_prefix_cache.py | 44 +- tests/dfx/conftest.py | 2 +- tests/dfx/perf/scripts/run_benchmark.py | 2 +- tests/dfx/stability/conftest.py | 130 +- tests/dfx/stability/helpers.py | 117 + .../scripts/test_benchmark_stability.py | 2 +- .../cache/test_teacache_extractors.py | 2 +- .../distributed/test_ulysses_uaa_perf.py | 2 +- .../lora/{conftest.py => helpers.py} | 0 tests/diffusion/lora/test_lora_manager.py | 2 +- .../diffusion/models/bagel/test_bagel_lora.py | 2 +- .../models/flux2/test_flux2_transformer_tp.py | 2 +- .../models/glm_image/test_glm_image_sp.py | 25 +- .../quantization/test_quantization_quality.py | 2 +- .../diffusion/test_diffusion_model_runner.py | 2 +- .../diffusion/test_diffusion_step_pipeline.py | 2 +- tests/e2e/accuracy/conftest.py | 15 +- tests/e2e/accuracy/{utils.py => helpers.py} | 16 +- tests/e2e/accuracy/test_gebench_h100_smoke.py | 4 +- .../accuracy/test_gedit_bench_h100_smoke.py | 4 +- tests/e2e/accuracy/test_qwen_image.py | 15 +- tests/e2e/accuracy/test_qwen_image_edit.py | 15 +- tests/e2e/accuracy/test_qwen_image_layered.py | 15 +- .../test_wan22_i2v_video_similarity.py | 4 +- .../test_async_omni_collective_rpc.py | 2 +- .../test_async_omni_qwen_image_generate.py | 2 +- .../custom_pipeline/test_worker_extension.py | 2 +- .../offline_inference/test_bagel_img2img.py | 7 +- .../e2e/offline_inference/test_bagel_lora.py | 7 +- .../offline_inference/test_bagel_text2img.py | 7 +- .../test_bagel_understanding.py | 5 +- tests/e2e/offline_inference/test_cache_dit.py | 4 +- .../e2e/offline_inference/test_cosyvoice3.py | 4 +- .../test_diffusion_cpu_offload.py | 5 +- .../test_diffusion_layerwise_offload.py | 4 +- .../offline_inference/test_diffusion_lora.py | 2 +- .../e2e/offline_inference/test_dynin_omni.py | 2 +- .../offline_inference/test_expert_parallel.py | 4 +- .../test_flux_autoround_w4a16.py | 5 +- .../offline_inference/test_flux_kontext.py | 2 +- .../test_hunyuanimage3_text2img.py | 2 +- .../test_ltx2_cfg_parallel_parity.py | 2 +- .../e2e/offline_inference/test_magi_human.py | 4 +- .../offline_inference/test_mammoth_moda2.py | 6 +- .../offline_inference/test_ming_flash_omni.py | 6 +- tests/e2e/offline_inference/test_omnivoice.py | 4 +- .../e2e/offline_inference/test_ovis_image.py | 2 +- .../test_quantization_fp8.py | 4 +- .../offline_inference/test_qwen2_5_omni.py | 6 +- .../e2e/offline_inference/test_qwen3_omni.py | 8 +- .../offline_inference/test_qwen3_tts_base.py | 4 +- .../test_qwen3_tts_customvoice.py | 4 +- .../test_qwen_image_diffusion_batching.py | 4 +- .../test_sequence_parallel.py | 4 +- .../test_stable_audio_expansion.py | 4 +- tests/e2e/offline_inference/test_t2i_model.py | 2 +- tests/e2e/offline_inference/test_t2v_model.py | 2 +- tests/e2e/offline_inference/test_teacache.py | 4 +- .../test_vae_decode_parallelism.py | 2 +- tests/e2e/offline_inference/test_voxcpm.py | 10 +- tests/e2e/offline_inference/test_voxcpm2.py | 4 +- .../e2e/offline_inference/test_voxtral_tts.py | 5 +- .../test_zimage_parallelism.py | 10 +- .../online_serving/test_bagel_expansion.py | 11 +- tests/e2e/online_serving/test_bagel_online.py | 4 +- .../e2e/online_serving/test_cosyvoice3_tts.py | 4 +- .../test_dynin_omni_expansion.py | 8 +- .../online_serving/test_flux2_expansion.py | 8 +- .../test_flux2_klein_inpaint_expansion.py | 2 +- .../test_flux_2_dev_expansion.py | 9 +- .../test_flux_kontext_expansion.py | 9 +- .../test_hunyuan_video_15_expansion.py | 8 +- .../e2e/online_serving/test_image_gen_edit.py | 2 +- .../test_images_generations_lora.py | 4 +- .../test_longcat_image_edit_expansion.py | 11 +- .../test_longcat_image_expansion.py | 9 +- tests/e2e/online_serving/test_mimo_audio.py | 11 +- .../online_serving/test_ming_flash_omni.py | 9 +- .../online_serving/test_nextstep_expansion.py | 4 +- tests/e2e/online_serving/test_omnivoice.py | 5 +- tests/e2e/online_serving/test_qwen2_5_omni.py | 13 +- tests/e2e/online_serving/test_qwen3_omni.py | 13 +- .../test_qwen3_omni_expansion.py | 17 +- .../test_qwen3_omni_realtime_websocket.py | 8 +- .../e2e/online_serving/test_qwen3_tts_base.py | 5 +- .../test_qwen3_tts_base_expansion.py | 5 +- .../online_serving/test_qwen3_tts_batch.py | 10 +- .../test_qwen3_tts_customvoice.py | 5 +- .../test_qwen3_tts_customvoice_expansion.py | 5 +- .../test_qwen3_tts_speaker_embedding.py | 5 +- .../test_qwen3_tts_websocket.py | 5 +- .../test_qwen_image_edit_expansion.py | 11 +- .../test_qwen_image_expansion.py | 9 +- .../test_qwen_image_layered_expansion.py | 12 +- .../e2e/online_serving/test_sd3_expansion.py | 8 +- .../test_video_generation_api.py | 4 +- tests/e2e/online_serving/test_voxtral_tts.py | 4 +- .../online_serving/test_wan22_expansion.py | 10 +- .../test_wan_2_1_vace_expansion.py | 8 +- .../online_serving/test_zimage_expansion.py | 8 +- tests/engine/test_async_omni_engine_abort.py | 10 +- .../openai_api/test_text_splitter.py | 2 +- tests/examples/conftest.py | 354 +- tests/examples/helpers.py | 353 ++ .../offline_inference/test_text_to_image.py | 6 +- .../online_serving/test_qwen2_5_omni.py | 8 +- .../online_serving/test_qwen3_omni.py | 8 +- .../online_serving/test_text_to_image.py | 7 +- tests/helpers/__init__.py | 8 + tests/helpers/assertions.py | 506 +++ tests/helpers/env.py | 280 ++ tests/helpers/fixtures/__init__.py | 1 + tests/helpers/fixtures/env.py | 59 + tests/helpers/fixtures/log.py | 7 + tests/helpers/fixtures/run_args.py | 17 + tests/helpers/fixtures/runtime.py | 137 + tests/helpers/mark.py | 135 + tests/helpers/media.py | 643 +++ .../utils.py => helpers/process.py} | 92 +- tests/helpers/runtime.py | 1413 +++++++ tests/helpers/stage_config.py | 475 +++ .../cosyvoice3/test_cosyvoice3_components.py | 2 +- tests/test_config_factory.py | 6 +- tests/utils.py | 837 ---- tools/pre_commit/check_pickle_imports.py | 3 +- 130 files changed, 4549 insertions(+), 5306 deletions(-) create mode 100644 tests/dfx/stability/helpers.py rename tests/diffusion/lora/{conftest.py => helpers.py} (100%) rename tests/e2e/accuracy/{utils.py => helpers.py} (91%) create mode 100644 tests/examples/helpers.py create mode 100644 tests/helpers/__init__.py create mode 100644 tests/helpers/assertions.py create mode 100644 tests/helpers/env.py create mode 100644 tests/helpers/fixtures/__init__.py create mode 100644 tests/helpers/fixtures/env.py create mode 100644 tests/helpers/fixtures/log.py create mode 100644 tests/helpers/fixtures/run_args.py create mode 100644 tests/helpers/fixtures/runtime.py create mode 100644 tests/helpers/mark.py create mode 100644 tests/helpers/media.py rename tests/{e2e/offline_inference/utils.py => helpers/process.py} (58%) create mode 100644 tests/helpers/runtime.py create mode 100644 tests/helpers/stage_config.py delete mode 100644 tests/utils.py diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 02d8cced4..c33d7b4d1 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -552,7 +552,7 @@ steps: - label: ":full_moon: Diffusion X2V · Accuracy Test" timeout_in_minutes: 180 commands: - - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model + - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m advanced_model --run-level advanced_model agents: queue: "mithril-h100-pool" plugins: diff --git a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md index 69d6ad828..ab4deecd6 100644 --- a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md @@ -40,7 +40,7 @@ Currently all the features are available in online serving mode. Hence, only nee - Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/). - To maximize code reuse, you may refer to - `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper. - - `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`. + - `tests/helpers/mark.py` for `@hardware_test(...)` and `hardware_marks`. - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases. - Doc: add a concise docstring for each test function. - Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py). diff --git a/docs/contributing/ci/tests_markers.md b/docs/contributing/ci/tests_markers.md index 7c1ba1c73..7628db284 100644 --- a/docs/contributing/ci/tests_markers.md +++ b/docs/contributing/ci/tests_markers.md @@ -38,7 +38,7 @@ Defined in `pyproject.toml`: ### Example usage for markers ```python -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test @pytest.mark.core_model @pytest.mark.omni @@ -53,7 +53,7 @@ def test_video_to_audio() ### Decorator: `@hardware_test` -This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions: +This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/helpers/mark.py` performs the following actions: 1. **Applies platform and resource markers** Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`). @@ -105,7 +105,7 @@ This decorator is intended to make hardware-aware, cross-platform test authoring `hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function. ```python -from tests.utils import hardware_marks +from tests.helpers.mark import hardware_marks MULTI_CARD_MARKS = hardware_marks( res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2 @@ -133,9 +133,9 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator "distributed_tpu: Tests that require multiple TPU devices", ] ``` -2. **Implement a marker construction function for your platform** in `vllm-omni/tests/utils.py`: +2. **Implement a marker construction function for your platform** in `vllm-omni/tests/helpers/mark.py`: ```python - # In vllm-omni/tests/utils.py + # In vllm-omni/tests/helpers/mark.py def tpu_marks(*, res: str, num_cards: int): test_platform = pytest.mark.tpu @@ -175,4 +175,4 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator - Plug into `hardware_marks` - You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation! -See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`). +See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`). diff --git a/docs/contributing/ci/tests_style.md b/docs/contributing/ci/tests_style.md index 392f00472..3a8cb0f12 100644 --- a/docs/contributing/ci/tests_style.md +++ b/docs/contributing/ci/tests_style.md @@ -221,15 +221,13 @@ from pathlib import Path import openai import pytest -from tests.conftest import ( - OmniServer, - convert_audio_to_text, +from tests.helpers.media import ( + convert_audio_bytes_to_text, cosine_similarity_text, - dummy_messages_from_mix_data, generate_synthetic_video, - merge_base64_and_convert_to_text, ) -from tests.utils import get_deploy_config_path +from tests.helpers.runtime import OmniServer, dummy_messages_from_mix_data +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.platforms import current_omni_platform # Edit: model name and stage config path @@ -406,7 +404,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N # PURPOSE: Verify text and audio outputs convey the same information # CUSTOMIZATION: Adjust similarity threshold (0.9) based on accuracy requirements assert audio_data is not None, "No audio output is generated" - audio_content = merge_base64_and_convert_to_text(audio_data) + audio_content = convert_audio_bytes_to_text(audio_data) print(f"text content is: {text_content}") print(f"audio content is: {audio_content}") similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) @@ -429,7 +427,7 @@ from pathlib import Path import pytest from vllm.assets.video import VideoAsset -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test from ..multi_stages.conftest import OmniRunner # Optional: set process start method for workers diff --git a/tests/conftest.py b/tests/conftest.py index 83752521f..77075f952 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3520 +1,62 @@ -import atexit -import base64 -import datetime -import io -import json -import math -import os -import random -import re -import tempfile - -import requests - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -# Set CPU device for CI environments without GPU -if "VLLM_TARGET_DEVICE" not in os.environ: - os.environ["VLLM_TARGET_DEVICE"] = "cpu" - -import concurrent.futures -import contextlib -import gc -import multiprocessing -import socket -import subprocess -import sys -import threading -import time -import uuid -from collections.abc import Generator -from dataclasses import dataclass -from io import BytesIO -from pathlib import Path -from typing import Any, NamedTuple - -import cv2 -import numpy as np -import psutil -import pytest -import soundfile as sf -import torch -import yaml -from openai import OpenAI, omit -from PIL import Image -from transformers import pipeline -from vllm import TextPrompt -from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from vllm.logger import init_logger -from vllm.utils.network_utils import get_open_port - -from vllm_omni.config.stage_config import resolve_deploy_yaml -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniSamplingParams -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform - -logger = init_logger(__name__) - - -PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None -PromptImageInput = list[Any] | Any | None -PromptVideoInput = list[Any] | Any | None - -_GENDER_PIPELINE = None -# transformers.Pipeline is not thread-safe; concurrent e2e requests must serialize inference. -_GENDER_PIPELINE_LOCK = threading.Lock() - -# int16 mono PCM from /v1/audio/speech when response_format=pcm (Qwen3-TTS code2wav output rate). -_PCM_SPEECH_SAMPLE_RATE_HZ = 24_000 - - -class OmniServerParams(NamedTuple): - model: str - port: int | None = None - stage_config_path: str | None = None - server_args: list[str] | None = None - env_dict: dict[str, str] | None = None - use_omni: bool = True - use_stage_cli: bool = False - init_timeout: int | None = None - stage_init_timeout: int | None = None # None defers to the server's own default (300 s) - - -def assert_image_diffusion_response( - response, - request_config: dict[str, Any], - run_level: str = None, -) -> None: - """ - Validate image diffusion response. - - Expected request_config schema: - { - "request_type": "image", - "extra_body": { - "num_outputs_per_prompt": 1, - "width": ..., - "height": ..., - ... - } - } - """ - assert response.images is not None, "Image response is None" - assert len(response.images) > 0, "No images in response" - - extra_body = request_config.get("extra_body") or {} - - num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt") - if num_outputs_per_prompt is not None: - assert len(response.images) == num_outputs_per_prompt, ( - f"Expected {num_outputs_per_prompt} images, got {len(response.images)}" - ) - - if run_level == "advanced_model": - width = extra_body.get("width") - height = extra_body.get("height") - - if width is not None or height is not None: - for img in response.images: - assert_image_valid(img, width=width, height=height) - - -def assert_video_diffusion_response( - response, - request_config: dict[str, Any], - run_level: str = None, -) -> None: - """ - Validate video diffusion response. - - Expected request_config schema: - { - "request_type": "video", - "form_data": { - "prompt": "...", - "num_frames": ..., - "width": ..., - "height": ..., - "fps": ..., - ... - } - } - """ - form_data = request_config.get("form_data", {}) - - assert response.videos is not None, "Video response is None" - assert len(response.videos) > 0, "No videos in response" - - expected_frames = _maybe_int(form_data.get("num_frames")) - expected_width = _maybe_int(form_data.get("width")) - expected_height = _maybe_int(form_data.get("height")) - expected_fps = _maybe_int(form_data.get("fps")) - - for vid_bytes in response.videos: - assert_video_valid( - vid_bytes, - num_frames=expected_frames, - width=expected_width, - height=expected_height, - fps=expected_fps, - ) - - -def assert_audio_diffusion_response( - response, - request_config: dict[str, Any], - run_level: str = None, -) -> None: - """ - Validate audio diffusion response. - """ - raise NotImplementedError("Audio validation is not implemented yet") - - -def _maybe_int(value: Any) -> int | None: - if value is None: - return None - return int(value) - - -def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, height: int | None = None): - """Assert the file is a loadable image with optional exact dimensions.""" - if isinstance(image, Path): - assert image.exists(), f"Image not found: {image}" - image = Image.open(image) - image.load() - assert image.width > 0 and image.height > 0 - if width is not None: - assert image.width == width, f"Expected width={width}, got {image.width}" - if height is not None: - assert image.height == height, f"Expected height={height}, got {image.height}" - return image - - -def assert_video_valid( - video: Path | bytes | BytesIO, - *, - num_frames: int | None = None, - width: int | None = None, - height: int | None = None, - fps: float | None = None, -) -> dict[str, int | float]: - """Assert the MP4 has the expected resolution and exact frame count.""" - temp_path = None - cap = None - try: - # Normalize input to file path - if isinstance(video, Path): - if not video.exists(): - raise AssertionError(f"Video file not found: {video}") - video_path = str(video) - else: - # Create temp file for bytes/BytesIO - suffix = ".mp4" - with tempfile.NamedTemporaryFile(delete=False, suffix=suffix, mode="wb") as tmp: - if isinstance(video, bytes): - tmp.write(video) - elif isinstance(video, BytesIO): - tmp.write(video.getvalue()) - else: - raise TypeError(f"Unsupported video type: {type(video)}") - temp_path = Path(tmp.name) - video_path = str(temp_path) - - # Open video capture - cap = cv2.VideoCapture(video_path) - if not cap.isOpened(): - raise AssertionError(f"Failed to open video: {video_path}") - - # Extract properties - actual_num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - actual_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - actual_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - actual_fps = cap.get(cv2.CAP_PROP_FPS) - - actual_num_frames = 0 - while True: - ok, _frame = cap.read() - if not ok: - break - actual_num_frames += 1 - - # Basic validity checks - if actual_num_frames <= 0: - raise AssertionError(f"Invalid frame count: {actual_num_frames} (must be > 0)") - if actual_width <= 0 or actual_height <= 0: - raise AssertionError(f"Invalid dimensions: {actual_width}x{actual_height} (must be > 0)") - if actual_fps <= 0: - raise AssertionError(f"Invalid FPS: {actual_fps} (must be > 0)") - - # Validate against expectations - if num_frames is not None: - expected_num_frames = (num_frames // 4) * 4 + 1 - assert actual_num_frames == expected_num_frames, ( - f"Frame count mismatch: expected {num_frames}, got {actual_num_frames}" - ) - if width is not None: - assert actual_width == width, f"Width mismatch: expected {width}px, got {actual_width}px" - if height is not None: - assert actual_height == height, f"Height mismatch: expected {height}px, got {actual_height}px" - if fps is not None: - # Use tolerance for float comparison (codec rounding) - assert abs(actual_fps - fps) < 0.5, f"FPS mismatch: expected {fps}, got {actual_fps:.2f}" - - return {"num_frames": actual_num_frames, "width": actual_width, "height": actual_height, "fps": actual_fps} - - except Exception as e: - print(f"ERROR: {type(e).__name__}: {e}", flush=True) - raise - - finally: - # Cleanup resources - if cap is not None: - cap.release() - if temp_path and temp_path.exists(): - try: - temp_path.unlink() - except OSError: - pass - - -def assert_audio_valid( - audio_or_path: Path | np.ndarray, - *, - sample_rate: int, - channels: int, - duration_s: float, -) -> None: - """Assert WAV file or (batch, channels, samples) ndarray matches expected audio format.""" - expected_samples = int(duration_s * sample_rate) - if isinstance(audio_or_path, np.ndarray): - audio = audio_or_path - assert audio.ndim == 3, f"Expected audio ndim=3 (batch, channels, samples), got shape {audio.shape}" - assert audio.shape[0] == 1, f"Expected batch size 1, got {audio.shape[0]}" - assert audio.shape[1] == channels, f"Expected {channels} channels, got {audio.shape[1]}" - assert audio.shape[2] == expected_samples, ( - f"Expected {expected_samples} samples ({duration_s}s @ {sample_rate} Hz), got {audio.shape[2]}" - ) - return - - path = audio_or_path - assert path.exists(), f"Audio not found: {path}" - info = sf.info(str(path)) - assert info.samplerate == sample_rate, f"Expected sample_rate={sample_rate}, got {info.samplerate}" - assert info.channels == channels, f"Expected {channels} channel(s), got {info.channels}" - assert info.frames == expected_samples, ( - f"Expected {expected_samples} frames ({duration_s}s @ {sample_rate} Hz), got {info.frames}" - ) - - -def decode_b64_image(b64: str): - img = Image.open(BytesIO(base64.b64decode(b64))) - img.load() - return img - - -@pytest.fixture(scope="session") -def model_prefix() -> str: - """Optional model-path prefix from MODEL_PREFIX env var. - Useful if models are downloaded to non-default local directories. - """ - prefix = os.environ.get("MODEL_PREFIX", "") - return f"{prefix.rstrip('/')}/" if prefix else "" - - -@pytest.fixture(autouse=True) -def default_vllm_config(): - """Set a default VllmConfig for all tests. - - This fixture is auto-used for all tests to ensure that any test - that directly instantiates vLLM CustomOps (e.g., RMSNorm, LayerNorm) - or model components has the required VllmConfig context. - - This fixture is required for vLLM 0.14.0+ where CustomOp initialization - requires a VllmConfig context set via set_current_vllm_config(). - """ - from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config - - # Use CPU device if no GPU is available (e.g., in CI environments) - has_gpu = torch.cuda.is_available() and torch.cuda.device_count() > 0 - device = "cuda" if has_gpu else "cpu" - device_config = DeviceConfig(device=device) - - with set_current_vllm_config(VllmConfig(device_config=device_config)): - yield - - -@pytest.fixture(autouse=True) -def clean_gpu_memory_between_tests(): - print("\n=== PRE-TEST GPU CLEANUP ===") - _run_pre_test_cleanup() - yield - _run_post_test_cleanup() - - -@pytest.fixture(autouse=True) -def log_test_name_before_test(request): - print(f"--- Running test: {request.node.name}") - yield - - -def _run_pre_test_cleanup(enable_force=False): - if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force: - print("\nPre-test GPU cleanup skipped(Default off is typical when one worker/instance runs many tests.)\n") - return - - print("\nPre-test GPU status:") - - num_gpus = torch.cuda.device_count() - if num_gpus > 0: - try: - from tests.utils import wait_for_gpu_memory_to_clear - - wait_for_gpu_memory_to_clear( - devices=list(range(num_gpus)), - threshold_ratio=0.05, - ) - except Exception as e: - print(f"Pre-test cleanup note: {e}") - - -def _run_post_test_cleanup(enable_force=False): - if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force: - print("GPU cleanup disabled") - return - - if torch.cuda.is_available(): - gc.collect() - torch.cuda.empty_cache() - - print("Post-test GPU status:") - _print_gpu_processes() - - -def _print_gpu_processes(): - """Print GPU information including nvidia-smi and system processes""" - - print("\n" + "=" * 80) - print("NVIDIA GPU Information (nvidia-smi)") - print("=" * 80) - - try: - nvidia_result = subprocess.run( - ["nvidia-smi"], - capture_output=True, - text=True, - timeout=5, - ) - - if nvidia_result.returncode == 0: - lines = nvidia_result.stdout.strip().split("\n") - for line in lines[:20]: - print(line) - - if len(lines) > 20: - print(f"... (showing first 20 of {len(lines)} lines)") - else: - print("nvidia-smi command failed") - - except (subprocess.TimeoutExpired, FileNotFoundError): - print("nvidia-smi not available or timed out") - except Exception as e: - print(f"Error running nvidia-smi: {e}") - - print("\n" + "=" * 80) - print("Detailed GPU Processes (nvidia-smi pmon)") - print("=" * 80) - - try: - pmon_result = subprocess.run( - ["nvidia-smi", "pmon", "-c", "1"], - capture_output=True, - text=True, - timeout=3, - ) - - if pmon_result.returncode == 0 and pmon_result.stdout.strip(): - print(pmon_result.stdout) - else: - print("No active GPU processes found via nvidia-smi pmon") - - except Exception: - print("nvidia-smi pmon not available") - - print("\n" + "=" * 80) - print("System Processes with GPU keywords") - print("=" * 80) - - -def dummy_messages_from_mix_data( - system_prompt: dict[str, Any] = None, - video_data_url: Any = None, - audio_data_url: Any = None, - image_data_url: Any = None, - content_text: str = None, -): - """Create messages with video、image、audio data URL for OpenAI API.""" - - if content_text is not None: - content = [{"type": "text", "text": content_text}] - else: - content = [] - - media_items = [] - if isinstance(video_data_url, list): - for video_url in video_data_url: - media_items.append((video_url, "video")) - else: - media_items.append((video_data_url, "video")) - - if isinstance(image_data_url, list): - for url in image_data_url: - media_items.append((url, "image")) - else: - media_items.append((image_data_url, "image")) - - if isinstance(audio_data_url, list): - for url in audio_data_url: - media_items.append((url, "audio")) - else: - media_items.append((audio_data_url, "audio")) - - content.extend( - {"type": f"{media_type}_url", f"{media_type}_url": {"url": url}} - for url, media_type in media_items - if url is not None - ) - messages = [{"role": "user", "content": content}] - if system_prompt is not None: - messages = [system_prompt] + messages - return messages - - -def generate_synthetic_audio( - duration: int, # seconds - num_channels: int, # 1：Mono，2：Stereo 5：5.1 surround sound - sample_rate: int = 48000, # Default use 48000Hz. - save_to_file: bool = False, -) -> dict[str, Any]: - """ - Generate TTS speech with pyttsx3 and return base64 string. - """ - - import pyttsx3 - import soundfile as sf - - def _pick_voice(engine: pyttsx3.Engine) -> str | None: - voices = engine.getProperty("voices") - if not voices: - return None - - preferred_tokens = ( - "natural", - "jenny", - "sonia", - "susan", - "zira", - "aria", - "hazel", - "samantha", - "ava", - "allison", - "female", - "woman", - "english-us", - "en-us", - "english", - ) - discouraged_tokens = ( - "espeak", - "robot", - "mbrola", - "microsoft david", - "male", - "man", - ) - - best_voice = voices[0] - best_score = float("-inf") - for voice in voices: - voice_text = f"{getattr(voice, 'id', '')} {getattr(voice, 'name', '')}".lower() - voice_languages = " ".join( - lang.decode(errors="ignore") if isinstance(lang, bytes) else str(lang) - for lang in getattr(voice, "languages", []) - ).lower() - combined_text = f"{voice_text} {voice_languages}" - score = 0 - for idx, token in enumerate(preferred_tokens): - if token in combined_text: - score += 20 - idx - for token in discouraged_tokens: - if token in combined_text: - score -= 10 - if "english" in combined_text or "en_" in combined_text or "en-" in combined_text: - score += 4 - if "en-us" in combined_text or "english-us" in combined_text: - score += 4 - if score > best_score: - best_score = score - best_voice = voice - - return best_voice.id - - def _resample_audio(audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray: - if src_sr == dst_sr or len(audio) == 0: - return audio.astype(np.float32) - - src_len = audio.shape[0] - dst_len = max(1, int(round(src_len * float(dst_sr) / float(src_sr)))) - src_idx = np.arange(src_len, dtype=np.float32) - dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32) - - resampled_channels: list[np.ndarray] = [] - for ch in range(audio.shape[1]): - resampled_channels.append(np.interp(dst_idx, src_idx, audio[:, ch]).astype(np.float32)) - return np.stack(resampled_channels, axis=1) - - def _match_channels(audio: np.ndarray, target_channels: int) -> np.ndarray: - current_channels = audio.shape[1] - if current_channels == target_channels: - return audio.astype(np.float32) - if target_channels == 1: - return np.mean(audio, axis=1, keepdims=True, dtype=np.float32) - if current_channels == 1: - return np.repeat(audio, target_channels, axis=1).astype(np.float32) - - collapsed = np.mean(audio, axis=1, keepdims=True, dtype=np.float32) - return np.repeat(collapsed, target_channels, axis=1).astype(np.float32) - - def _trim_silence(audio: np.ndarray, threshold: float = 0.01) -> np.ndarray: - if len(audio) == 0: - return audio - energy = np.max(np.abs(audio), axis=1) - voiced = np.where(energy > threshold)[0] - if len(voiced) == 0: - return audio - start = max(0, int(voiced[0]) - int(sample_rate * 0.02)) - end = min(len(audio), int(voiced[-1]) + int(sample_rate * 0.04) + 1) - return audio[start:end] - - def _enhance_speech(audio: np.ndarray) -> np.ndarray: - if len(audio) == 0: - return audio.astype(np.float32) - enhanced = audio.astype(np.float32).copy() - enhanced -= np.mean(enhanced, axis=0, keepdims=True, dtype=np.float32) - if len(enhanced) > 1: - preemphasis = enhanced.copy() - preemphasis[1:] = enhanced[1:] - 0.94 * enhanced[:-1] - enhanced = 0.7 * enhanced + 0.3 * preemphasis - # Mild dynamic-range compression for ASR/TTS robustness. - enhanced = np.sign(enhanced) * np.sqrt(np.abs(enhanced)) - # Light fade to avoid clicks after trimming/repeating. - fade = min(len(enhanced) // 4, max(1, int(sample_rate * 0.01))) - if fade > 1: - ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32) - ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32) - enhanced[:fade] *= ramp_in[:, None] - enhanced[-fade:] *= ramp_out[:, None] - peak = float(np.max(np.abs(enhanced))) - if peak > 1e-8: - enhanced = enhanced / peak * 0.95 - return enhanced.astype(np.float32) - - phrase_text = "test" - num_samples = int(sample_rate * max(1, duration)) - audio_data = np.zeros((num_samples, num_channels), dtype=np.float32) - - engine = pyttsx3.init() - engine.setProperty("rate", 112) - engine.setProperty("volume", 1.0) - selected_voice = _pick_voice(engine) - if selected_voice is not None: - engine.setProperty("voice", selected_voice) - - temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) - temp_wav.close() - - try: - engine.save_to_file(phrase_text, temp_wav.name) - engine.runAndWait() - engine.stop() - - ready = False - for _ in range(50): - if os.path.exists(temp_wav.name) and os.path.getsize(temp_wav.name) > 44: - ready = True - break - time.sleep(0.1) - - if not ready: - raise RuntimeError("pyttsx3 did not produce a WAV file in time.") - - tts_audio, tts_sr = sf.read(temp_wav.name, dtype="float32", always_2d=True) - finally: - if os.path.exists(temp_wav.name): - os.unlink(temp_wav.name) - - if len(tts_audio) == 0: - raise RuntimeError("pyttsx3 produced an empty WAV file.") - - tts_audio = _resample_audio(tts_audio, tts_sr, sample_rate) - tts_audio = _match_channels(tts_audio, num_channels) - tts_audio = _trim_silence(tts_audio, threshold=0.012) - tts_audio = _enhance_speech(tts_audio) - - lead_silence = min(int(sample_rate * 0.02), num_samples // 8) - pause_samples = int(sample_rate * 0.18) - start = lead_silence - phrase_len = tts_audio.shape[0] - - while start < num_samples: - take = min(phrase_len, num_samples - start) - audio_data[start : start + take] = tts_audio[:take] - start += phrase_len + pause_samples - - max_amp = float(np.max(np.abs(audio_data))) - if max_amp > 0: - audio_data = audio_data / max_amp * 0.95 - - audio_bytes: bytes | None = None - output_path: str | None = None - result: dict[str, Any] = { - "np_array": audio_data.copy(), - } - - if save_to_file: - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - output_path = f"audio_{num_channels}ch_{timestamp}.wav" - - try: - sf.write(output_path, audio_data, sample_rate, format="WAV", subtype="PCM_16") - print(f"Audio saved: {output_path}") - - with open(output_path, "rb") as f: - audio_bytes = f.read() - except Exception as e: - print(f"Save failed: {e}") - save_to_file = False - - # If not saving or save failed, create in memory - if not save_to_file or audio_bytes is None: - buffer = io.BytesIO() - sf.write(buffer, audio_data, sample_rate, format="WAV", subtype="PCM_16") - buffer.seek(0) - audio_bytes = buffer.read() - - # Return result - base64_audio = base64.b64encode(audio_bytes).decode("utf-8") - result["base64"] = base64_audio - # Always include file_path to avoid KeyError in callers. - result["file_path"] = output_path if save_to_file and output_path else None - - return result - - -def _mux_mp4_bytes_with_synthetic_audio( - video_mp4_bytes: bytes, - *, - num_frames: int, - fps: float = 30.0, - sample_rate: int = 48000, -) -> bytes: - """ - Mux a video-only MP4 with mono TTS audio from :func:`generate_synthetic_audio` (AAC). - - Audio length is at least the video duration in whole seconds (rounded up); ffmpeg - ``-shortest`` trims to the video when the WAV is longer. - - Uses ffmpeg from ``imageio_ffmpeg`` when available, else ``ffmpeg`` on PATH. - If TTS or mux fails, returns ``video_mp4_bytes`` unchanged. - - Mux subprocess does **not** use ``capture_output=True``: ffmpeg can block writing - to a full stderr pipe while :func:`subprocess.run` waits for exit (classic deadlock). - """ - duration_sec = num_frames / fps if fps > 0 else 0.0 - # generate_synthetic_audio(duration=int) uses at least 1s of buffer internally - duration_int = max(1, int(math.ceil(duration_sec))) - - try: - audio_result = generate_synthetic_audio( - duration=duration_int, - num_channels=1, - sample_rate=sample_rate, - save_to_file=False, - ) - audio_pcm = audio_result["np_array"] - except Exception as e: - logger.warning("Synthetic video: generate_synthetic_audio failed (%s); using video-only MP4.", e) - return video_mp4_bytes - - try: - import imageio_ffmpeg - - ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe() - except Exception: - ffmpeg_exe = "ffmpeg" - - import tempfile - - try: - with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp: - vid_path = os.path.join(tmp, "video.mp4") - wav_path = os.path.join(tmp, "audio.wav") - out_path = os.path.join(tmp, "out.mp4") - with open(vid_path, "wb") as f: - f.write(video_mp4_bytes) - sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16") - cmd = [ - ffmpeg_exe, - "-y", - "-nostdin", - "-hide_banner", - "-loglevel", - "error", - "-i", - vid_path, - "-i", - wav_path, - "-c:v", - "copy", - "-c:a", - "aac", - "-b:a", - "128k", - "-shortest", - "-movflags", - "+faststart", - out_path, - ] - subprocess.run( - cmd, - check=True, - stdin=subprocess.DEVNULL, - timeout=300, - ) - with open(out_path, "rb") as f: - return f.read() - except ( - FileNotFoundError, - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - OSError, - ) as e: - logger.warning("Synthetic video: audio mux failed (%s); using video-only MP4.", e) - return video_mp4_bytes - - -def generate_synthetic_video( - width: int, - height: int, - num_frames: int, - save_to_file: bool = False, - *, - embed_audio: bool = False, -) -> dict[str, Any]: - """Generate synthetic video with bouncing balls and base64 MP4. - - When ``embed_audio`` is True, muxes mono AAC from :func:`generate_synthetic_audio` - (TTS + ffmpeg) into the MP4; otherwise returns video-only MP4 (faster when tests do - not need an audio track). - """ - - import cv2 - import imageio - - # Create random balls - num_balls = random.randint(3, 8) - balls = [] - - for _ in range(num_balls): - radius = min(width, height) // 8 - if radius < 1: - raise ValueError(f"Video dimensions ({width}x{height}) are too small for synthetic video generation") - x = random.randint(radius, width - radius) - y = random.randint(radius, height - radius) - - speed = random.uniform(3.0, 8.0) - angle = random.uniform(0, 2 * math.pi) - vx = speed * math.cos(angle) - vy = speed * math.sin(angle) - - # OpenCV uses BGR format, but imageio expects RGB - # We'll create in BGR first, then convert to RGB later - color_bgr = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255)) - - balls.append({"x": x, "y": y, "vx": vx, "vy": vy, "radius": radius, "color_bgr": color_bgr}) - - # Generate video frames - video_frames = [] - - for frame_idx in range(num_frames): - # Create black background (BGR format) - frame_bgr = np.zeros((height, width, 3), dtype=np.uint8) - - for ball in balls: - # Update position - ball["x"] += ball["vx"] - ball["y"] += ball["vy"] - - # Boundary collision detection - if ball["x"] - ball["radius"] <= 0 or ball["x"] + ball["radius"] >= width: - ball["vx"] = -ball["vx"] - ball["x"] = max(ball["radius"], min(width - ball["radius"], ball["x"])) - - if ball["y"] - ball["radius"] <= 0 or ball["y"] + ball["radius"] >= height: - ball["vy"] = -ball["vy"] - ball["y"] = max(ball["radius"], min(height - ball["radius"], ball["y"])) - - # Use cv2 to draw circle - x, y = int(ball["x"]), int(ball["y"]) - radius = ball["radius"] - - # Draw solid circle (main circle) - cv2.circle(frame_bgr, (x, y), radius, ball["color_bgr"], -1) - - # Add simple 3D effect: draw a brighter center - if radius > 3: # Only add highlight when radius is large enough - highlight_radius = max(1, radius // 2) - highlight_x = max(highlight_radius, min(x - radius // 4, width - highlight_radius)) - highlight_y = max(highlight_radius, min(y - radius // 4, height - highlight_radius)) - - # Create highlight color (brighter) - highlight_color = tuple(min(c + 40, 255) for c in ball["color_bgr"]) - cv2.circle(frame_bgr, (highlight_x, highlight_y), highlight_radius, highlight_color, -1) - - # Convert BGR to RGB for imageio - frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) - video_frames.append(frame_rgb) - - video_array = np.array(video_frames) - result = { - "np_array": video_array, - } - saved_file_path = None - - fps = 30 - buffer = io.BytesIO() - writer_kwargs = { - "format": "mp4", - "fps": fps, - "codec": "libx264", - "quality": 7, - "pixelformat": "yuv420p", - "macro_block_size": 16, - "ffmpeg_params": [ - "-preset", - "medium", - "-crf", - "23", - "-movflags", - "+faststart", - "-pix_fmt", - "yuv420p", - "-vf", - f"scale={width}:{height}", - ], - } - - try: - with imageio.get_writer(buffer, **writer_kwargs) as writer: - for frame in video_frames: - writer.append_data(frame) - buffer.seek(0) - video_only_bytes = buffer.read() - except Exception as e: - print(f"Warning: Failed to encode synthetic video: {e}") - raise - - if embed_audio: - video_bytes = _mux_mp4_bytes_with_synthetic_audio(video_only_bytes, num_frames=num_frames, fps=float(fps)) - else: - video_bytes = video_only_bytes - - if save_to_file: - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - output_path = f"video_{width}x{height}_{timestamp}.mp4" - try: - with open(output_path, "wb") as f: - f.write(video_bytes) - saved_file_path = output_path - print(f"Video saved to: {saved_file_path}") - except Exception as e: - print(f"Warning: Failed to save video to file {output_path}: {e}") - - base64_video = base64.b64encode(video_bytes).decode("utf-8") - - result["base64"] = base64_video - if save_to_file and saved_file_path: - result["file_path"] = saved_file_path - - return result - - -def generate_synthetic_image(width: int, height: int, save_to_file: bool = False) -> dict[str, Any]: - """Generate synthetic image with randomly colored squares and return base64 string.""" - from PIL import Image, ImageDraw - - # Create white background - image = Image.new("RGB", (width, height), (255, 255, 255)) - draw = ImageDraw.Draw(image) - - # Generate random number of squares - num_squares = random.randint(3, 8) - - for _ in range(num_squares): - # Random square size - square_size = random.randint(min(width, height) // 8, min(width, height) // 4) - - # Random position - x = random.randint(0, width - square_size - 1) - y = random.randint(0, height - square_size - 1) - - # Random color - color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) - - # Random border width - border_width = random.randint(1, 5) - - # Draw square - draw.rectangle([x, y, x + square_size, y + square_size], fill=color, outline=(0, 0, 0), width=border_width) - - image_array = np.array(image) - result = {"np_array": image_array.copy()} - - # Handle file saving - image_bytes = None - saved_file_path = None - - if save_to_file: - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - output_path = f"image_{width}x{height}_{timestamp}.jpg" - - try: - # Save image to file - image.save(output_path, format="JPEG", quality=85, optimize=True) - saved_file_path = output_path - print(f"Image saved to: {saved_file_path}") - - # Read file for base64 encoding - with open(output_path, "rb") as f: - image_bytes = f.read() - - except Exception as e: - print(f"Warning: Failed to save image to file {output_path}: {e}") - save_to_file = False - - # If not saving or save failed, create in memory - if not save_to_file or image_bytes is None: - buffer = io.BytesIO() - image.save(buffer, format="JPEG", quality=85, optimize=True) - buffer.seek(0) - image_bytes = buffer.read() - - # Generate base64 - base64_image = base64.b64encode(image_bytes).decode("utf-8") - - # Return result - result["base64"] = base64_image - if save_to_file and saved_file_path: - result["file_path"] = saved_file_path - - return result - - -def preprocess_text(text): - import opencc - - word_to_num = { - "zero": "0", - "one": "1", - "two": "2", - "three": "3", - "four": "4", - "five": "5", - "six": "6", - "seven": "7", - "eight": "8", - "nine": "9", - "ten": "10", - } - - for word, num in word_to_num.items(): - pattern = r"\b" + re.escape(word) + r"\b" - text = re.sub(pattern, num, text, flags=re.IGNORECASE) - - text = re.sub(r"[^\w\s]", "", text) - text = re.sub(r"\s+", " ", text) - cc = opencc.OpenCC("t2s") - text = cc.convert(text) - - # Special handling for spaces between Chinese characters: - # - Keep single spaces between English words/numbers - # - Remove spaces only when surrounded by Chinese characters on both sides to prevent incorrect word segmentation - text = re.sub(r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])", "", text) - - return text.lower().strip() - - -def cosine_similarity_text(text1, text2, n: int = 3): - from collections import Counter - - if not text1 or not text2: - return 0.0 - - text1 = preprocess_text(text1) - text2 = preprocess_text(text2) - print(f"cosine similarity text1 is: {text1}, text2 is: {text2}") - - ngrams1 = [text1[i : i + n] for i in range(len(text1) - n + 1)] - ngrams2 = [text2[i : i + n] for i in range(len(text2) - n + 1)] - - counter1 = Counter(ngrams1) - counter2 = Counter(ngrams2) - - all_ngrams = set(counter1.keys()) | set(counter2.keys()) - vec1 = [counter1.get(ng, 0) for ng in all_ngrams] - vec2 = [counter2.get(ng, 0) for ng in all_ngrams] - - dot_product = sum(a * b for a, b in zip(vec1, vec2)) - norm1 = sum(a * a for a in vec1) ** 0.5 - norm2 = sum(b * b for b in vec2) ** 0.5 - - if norm1 == 0 or norm2 == 0: - return 0.0 - return dot_product / (norm1 * norm2) - - -def convert_audio_to_text(audio_data): - """ - Convert base64 encoded audio data to text using speech recognition. - """ - audio_data = base64.b64decode(audio_data) - output_path = f"./test_{uuid.uuid4().hex}.wav" - with open(output_path, "wb") as audio_file: - audio_file.write(audio_data) - - print(f"audio data is saved: {output_path}") - text = convert_audio_file_to_text(output_path=output_path) - return text - - -def _merge_base64_audio_to_segment(base64_list: list[str]): - """Merge a list of base64-encoded audio chunks into one pydub AudioSegment.""" - from pydub import AudioSegment - - merged = None - for b64 in base64_list: - raw = base64.b64decode(b64.split(",", 1)[-1]) - seg = AudioSegment.from_file(io.BytesIO(raw)) - merged = seg if merged is None else merged + seg - return merged - - -@contextlib.contextmanager -def _serialize_whisper_small_model_download(): - """Serialize Whisper ``small`` cache writes across processes (Linux; ``fcntl``).""" - import fcntl - - lock_path = Path.home() / ".cache" / "whisper" / ".small_model_download.lock" - lock_path.parent.mkdir(parents=True, exist_ok=True) - f = open(lock_path, "a+b") - try: - fcntl.flock(f.fileno(), fcntl.LOCK_EX) - yield - finally: - fcntl.flock(f.fileno(), fcntl.LOCK_UN) - f.close() - - -def _whisper_transcribe_in_current_process(output_path: str) -> str: - import whisper - - # Multi-GPU: use last visible device to avoid colliding with default device 0; single device uses 0. - device_index = None - if current_omni_platform.is_available(): - n = current_omni_platform.get_device_count() - if n == 1: - device_index = 0 - elif n > 1: - device_index = n - 1 - - if device_index is not None: - torch_device = current_omni_platform.get_torch_device(device_index) - current_omni_platform.set_device(torch_device) - device = str(torch_device) - use_accelerator = True - else: - use_accelerator = False - device = "cpu" - with _serialize_whisper_small_model_download(): - model = whisper.load_model("small", device=device) - try: - text = model.transcribe( - output_path, - temperature=0.0, - word_timestamps=True, - condition_on_previous_text=False, - )["text"] - finally: - del model - gc.collect() - if use_accelerator: - current_omni_platform.synchronize() - current_omni_platform.empty_cache() - - return text or "" - - -def convert_audio_file_to_text(output_path: str) -> str: - """Convert an audio file to text in an isolated subprocess (spawn).""" - ctx = multiprocessing.get_context("spawn") - with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor: - future = executor.submit(_whisper_transcribe_in_current_process, output_path) - return future.result() - - -def convert_audio_bytes_to_text(raw_bytes: bytes) -> str: - """ - Write container audio bytes (WAV, etc.) to a temp WAV file suitable for Whisper/ffmpeg. - Normalizes with soundfile to PCM_16 WAV when possible to avoid codec issues. - """ - output_path = f"./test_{uuid.uuid4().hex}.wav" - data, samplerate = sf.read(io.BytesIO(raw_bytes)) - sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16") - text = convert_audio_file_to_text(output_path) - return text - - -def modify_stage_config( - yaml_path: str, - updates: dict[str, Any] = None, - deletes: dict[str, Any] = None, -) -> str: - """ - Modify configurations in a YAML file, supporting both top-level and stage-specific modifications, - including addition, modification, and deletion of configurations. - - Args: - yaml_path: Path to the YAML configuration file. - updates: Dictionary containing both top-level and stage-specific modifications to add or update. - Format: { - 'async_chunk': True, - 'stage_args': { - 0: {'engine_args.max_model_len': 5800}, - 1: {'engine_args.max_num_seqs': 2} - } - } - deletes: Dictionary containing configurations to delete. - Format: { - 'old_config': None, # Delete entire key - 'stage_args': { - 0: ['engine_args.old_param'], - 1: ['runtime.unused_setting'] - } - } - - Returns: - str: Path to the newly created modified YAML file with timestamp suffix. - """ - path = Path(yaml_path) - if not path.exists(): - raise FileNotFoundError(f"yaml does not exist: {path}") - - try: - with open(yaml_path, encoding="utf-8") as f: - config = yaml.safe_load(f) or {} - except Exception as e: - raise ValueError(f"Cannot parse YAML file: {e}") - - # Helper function to apply update - def apply_update(config_dict: dict, key_path: str, value: Any) -> None: - """Apply update to dictionary using dot-separated path.""" - # Handle direct list assignment (e.g., engine_input_source: [1, 2]) - if "." not in key_path: - # Simple key, set directly - config_dict[key_path] = value - return - - current = config_dict - keys = key_path.split(".") - - for i in range(len(keys) - 1): - key = keys[i] - - # Handle list indices - if key.isdigit() and isinstance(current, list): - index = int(key) - if index < 0: - raise ValueError(f"Negative list index not allowed: {index}") - if index >= len(current): - # Expand list if needed - while len(current) <= index: - # If we need to go deeper (more keys after this), create a dict - # Otherwise, create None placeholder - current.append({} if i < len(keys) - 2 else None) - current = current[index] - elif isinstance(current, dict): - # Handle dictionary keys - if key not in current: - # If there are more keys after this, create appropriate structure - if i < len(keys) - 1: - # Check if next key is a digit (list index) or string (dict key) - if keys[i + 1].isdigit(): - current[key] = [] - else: - current[key] = {} - else: - # This is the last key, create based on value type - current[key] = [] if isinstance(value, list) else {} - elif not isinstance(current[key], (dict, list)) and i < len(keys) - 1: - # If current value is not dict/list but we need to go deeper, replace it - if keys[i + 1].isdigit(): - current[key] = [] - else: - current[key] = {} - current = current[key] - else: - # Current is not a dict or list, cannot traverse further - raise TypeError( - f"Cannot access {'.'.join(keys[: i + 1])} as a dict/list. It's a {type(current).__name__}" - ) - - # Set the final value - last_key = keys[-1] - if isinstance(current, list) and last_key.isdigit(): - # Setting a value in a list by index - index = int(last_key) - if index < 0: - raise ValueError(f"Negative list index not allowed: {index}") - if index >= len(current): - # Expand list if needed - while len(current) <= index: - current.append(None) - current[index] = value - elif isinstance(current, dict): - # Special case: if the value is a list and we're setting a top-level key - # Example: updating engine_input_source with [1, 2] - current[last_key] = value - else: - # Current is not a dict, cannot set key - raise TypeError(f"Cannot set value at {key_path}. Current type is {type(current).__name__}, expected dict.") - - # Helper function to delete by path - def delete_by_path(config_dict: dict, path: str) -> None: - """Delete configuration by dot-separated path.""" - if not path: - return - - current = config_dict - keys = path.split(".") - - # Traverse to the parent - for i in range(len(keys) - 1): - key = keys[i] - - # Handle list indices - if key.isdigit() and isinstance(current, list): - index = int(key) - if index < 0 or index >= len(current): - raise KeyError(f"List index {index} out of bounds") - current = current[index] - elif isinstance(current, dict): - if key not in current: - raise KeyError(f"Path {'.'.join(keys[: i + 1])} does not exist") - current = current[key] - else: - raise TypeError( - f"Cannot access {'.'.join(keys[: i + 1])} as a dict/list. It's a {type(current).__name__}" - ) - - # Delete the item - last_key = keys[-1] - - if isinstance(current, list) and last_key.isdigit(): - index = int(last_key) - if index < 0 or index >= len(current): - raise KeyError(f"List index {index} out of bounds") - del current[index] - elif isinstance(current, dict) and last_key in current: - del current[last_key] - else: - print(f"Path {path} does not exist") - - _stage_key = "stages" if "stages" in config else "stage_args" - - # Apply deletions first - if deletes: - for key, value in deletes.items(): - if key in ("stage_args", "stages"): - if value and isinstance(value, dict): - stage_args = config.get(_stage_key, []) - if not stage_args: - raise ValueError("stage_args does not exist in config") - - for stage_id, delete_paths in value.items(): - if not delete_paths: - continue - - # Find stage by ID - target_stage = None - for stage in stage_args: - if stage.get("stage_id") == int(stage_id): - target_stage = stage - break - - if target_stage is None: - continue - - # Delete specified paths in this stage - # Avoid shadowing the original YAML Path used for the output filename below. - for delete_path in delete_paths: - if delete_path: # Skip empty paths - delete_by_path(target_stage, delete_path) - elif "." in key: - # Delete using dot-separated path - delete_by_path(config, key) - elif value is None and key in config: - # Delete entire key - del config[key] - - # Apply updates - if updates: - for key, value in updates.items(): - if key in ("stage_args", "stages"): - if value and isinstance(value, dict): - stage_args = config.get(_stage_key, []) - if not stage_args: - raise ValueError("stage_args does not exist in config") - - for stage_id, stage_updates in value.items(): - # Find stage by ID - target_stage = None - for stage in stage_args: - if stage.get("stage_id") == int(stage_id): - target_stage = stage - break - - if target_stage is None: - available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s] - raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}") - - # Apply updates to this stage - for update_path, val in stage_updates.items(): - # Check if this is a simple key (not dot-separated) - # Example: 'engine_input_source' vs 'engine_args.max_model_len' - if "." not in update_path: - # Direct key assignment (e.g., updating a list value) - target_stage[update_path] = val - else: - # Dot-separated path (e.g., nested dict access) - apply_update(target_stage, update_path, val) - elif "." in key: - # Apply using dot-separated path - apply_update(config, key, value) - else: - # Direct top-level key - config[key] = value - - # Unique suffix: multiple modify_stage_config calls in one process often run - # within the same second (e.g. test_qwen3_omni_expansion imports both - # get_chunk_config and get_batch_token_config). int(time.time()) would collide - # and the later write would overwrite the earlier YAML on disk. - # Keep generated configs outside the repo and delete them when pytest exits. - output_fd, output_path = tempfile.mkstemp(prefix=f"{path.stem}_", suffix=".yaml") - atexit.register(Path(output_path).unlink, missing_ok=True) - - with os.fdopen(output_fd, "w", encoding="utf-8") as f: - yaml.dump(config, f, default_flow_style=None, sort_keys=False, allow_unicode=True, indent=2) - - return str(output_path) - - -class OmniServer: - """Omniserver for vLLM-Omni tests.""" - - def __init__( - self, - model: str, - serve_args: list[str], - *, - port: int | None = None, - env_dict: dict[str, str] | None = None, - use_omni: bool = True, - ) -> None: - _run_pre_test_cleanup(enable_force=True) - _run_post_test_cleanup(enable_force=True) - cleanup_dist_env_and_memory() - self.model = model - self.serve_args = serve_args - self.env_dict = env_dict - self.use_omni = use_omni - self.proc: subprocess.Popen | None = None - self.host = "127.0.0.1" - if port is None: - self.port = get_open_port() - else: - self.port = port - - def _start_server(self) -> None: - """Start the vLLM-Omni server subprocess.""" - env = os.environ.copy() - env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - if self.env_dict is not None: - env.update(self.env_dict) - - cmd = [ - sys.executable, - "-m", - "vllm_omni.entrypoints.cli.main", - "serve", - self.model, - "--host", - self.host, - "--port", - str(self.port), - ] - if self.use_omni: - cmd.append("--omni") - cmd += self.serve_args - - print(f"Launching OmniServer with: {' '.join(cmd)}") - self.proc = subprocess.Popen( - cmd, - env=env, - cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), # Set working directory to vllm-omni root - ) - - # Wait for server to be ready - max_wait = 1200 # 20 minutes - start_time = time.time() - while time.time() - start_time < max_wait: - # Check for process status - ret = self.proc.poll() - if ret is not None: - raise RuntimeError(f"Server processes exited with code {ret} before becoming ready.") - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.settimeout(1) - result = sock.connect_ex((self.host, self.port)) - if result == 0: - print(f"Server ready on {self.host}:{self.port}") - return - time.sleep(2) - - raise RuntimeError(f"Server failed to start within {max_wait} seconds") - - def _kill_process_tree(self, pid): - """kill process and its children with verification""" - try: - parent = psutil.Process(pid) - children = parent.children(recursive=True) - - # Get all PIDs first - all_pids = [pid] + [child.pid for child in children] - - # Terminate children - for child in children: - try: - child.terminate() - except psutil.NoSuchProcess: - pass - - # Wait for children - gone, still_alive = psutil.wait_procs(children, timeout=10) - - # Kill remaining children - for child in still_alive: - try: - child.kill() - except psutil.NoSuchProcess: - pass - - # Terminate parent - try: - parent.terminate() - parent.wait(timeout=10) - except (psutil.NoSuchProcess, psutil.TimeoutExpired): - try: - parent.kill() - except psutil.NoSuchProcess: - pass - - # VERIFICATION: Check if all processes are gone - time.sleep(1) # Give system time - alive_processes = [] - for check_pid in all_pids: - if psutil.pid_exists(check_pid): - alive_processes.append(check_pid) - - if alive_processes: - print(f"Warning: Processes still alive: {alive_processes}") - # Optional: Try system kill - import subprocess - - for alive_pid in alive_processes: - try: - subprocess.run(["kill", "-9", str(alive_pid)], timeout=2) - except Exception as e: - print(f"Cleanup failed: {e}") - - except psutil.NoSuchProcess: - pass - - def __enter__(self): - self._start_server() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - if self.proc: - self._kill_process_tree(self.proc.pid) - _run_pre_test_cleanup(enable_force=True) - _run_post_test_cleanup(enable_force=True) - cleanup_dist_env_and_memory() - - -class OmniServerStageCli(OmniServer): - """Omni server harness that exercises the stage CLI flow.""" - - def __init__( - self, - model: str, - stage_config_path: str, - serve_args: list[str] | None = None, - *, - stage_ids: list[int] | None = None, - port: int | None = None, - env_dict: dict[str, str] | None = None, - ) -> None: - super().__init__(model, serve_args or [], port=port, env_dict=env_dict, use_omni=True) - self.stage_config_path = stage_config_path - self.master_port = get_open_port() - self.visible_device_list = self._load_visible_device_list(env_dict) - resolved_cfg = resolve_deploy_yaml(stage_config_path) - # Dump the resolved deploy config so CI logs show each stage's - # gpu_memory_utilization / max_model_len / max_num_seqs after - # base_config inheritance and overlay merge — essential when - # diagnosing OOMs that depend on the merged values. - print( - f"[OmniServerStageCli] Resolved deploy config from {stage_config_path}:\n" - f"{yaml.safe_dump(resolved_cfg, sort_keys=False, default_flow_style=False)}", - flush=True, - ) - self.stage_runtime_devices = self._load_stage_runtime_devices(resolved_cfg) - self.stage_ids = stage_ids or self._load_stage_ids(resolved_cfg) - if 0 not in self.stage_ids: - raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}") - self.stage_procs: dict[int, subprocess.Popen] = {} - self.proc = None - - @staticmethod - def _stage_entries(cfg: dict) -> list[dict]: - """Return the list of stage entries from either legacy (``stage_args``) - or new-schema (``stages``) deploy YAMLs.""" - return cfg.get("stage_args") or cfg.get("stages") or [] - - @staticmethod - def _load_stage_ids(resolved_config: dict) -> list[int]: - stage_ids = [ - stage["stage_id"] for stage in OmniServerStageCli._stage_entries(resolved_config) if "stage_id" in stage - ] - if not stage_ids: - raise ValueError("No stage IDs found in resolved config") - return stage_ids - - @staticmethod - def _load_stage_runtime_devices(resolved_config: dict) -> dict[int, str]: - runtime_devices: dict[int, str] = {} - for stage in OmniServerStageCli._stage_entries(resolved_config): - stage_id = stage.get("stage_id") - # New schema: stage.devices is flat at stage level. - # Legacy schema: stage.runtime.devices is nested. - devices = stage.get("devices") or stage.get("runtime", {}).get("devices") - if stage_id is not None and devices: - runtime_devices[int(stage_id)] = str(devices) - return runtime_devices - - @classmethod - def _parse_device_list(cls, devices: str | int) -> list[str]: - if isinstance(devices, int): - if devices < 0: - raise ValueError("Device IDs must be non-negative integers") - return [str(devices)] - return [token.strip() for token in str(devices).split(",") if token.strip()] - - @classmethod - def _load_visible_device_list(cls, env_dict: dict[str, str] | None) -> list[str] | None: - env = os.environ.copy() - if env_dict is not None: - env.update(env_dict) - - env_var = getattr(current_omni_platform, "device_control_env_var", None) - if env_var and env_var in env: - return [token.strip() for token in env[env_var].split(",") if token.strip()] - return None - - @classmethod - def _map_stage_devices(cls, stage_id: int, visible_device_list: list[str] | None, devices: str) -> str: - device_list = cls._parse_device_list(devices) - - if visible_device_list is None: - return ",".join(device_list) - - if not all(device.isdigit() for device in device_list): - raise ValueError("Logical devices must be non-negative integers") - - logical_ids = [int(device) for device in device_list] - if logical_ids and max(logical_ids) >= len(visible_device_list): - raise ValueError( - f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices" - ) - - return ",".join(visible_device_list[idx] for idx in logical_ids) - - def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str) -> None: - mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, devices) - env_var = getattr(current_omni_platform, "device_control_env_var", None) - if env_var: - env[env_var] = mapped_devices - - def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]: - cmd = [ - sys.executable, - "-m", - "vllm_omni.entrypoints.cli.main", - "serve", - self.model, - "--omni", - "--stage-configs-path", - self.stage_config_path, - "--stage-id", - str(stage_id), - "--omni-master-address", - self.host, - "--omni-master-port", - str(self.master_port), - ] - - if headless: - cmd.append("--headless") - else: - cmd += ["--host", self.host, "--port", str(self.port)] - - cmd += self.serve_args - return cmd - - def _launch_stage(self, stage_id: int, *, headless: bool) -> None: - env = os.environ.copy() - env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - if self.env_dict is not None: - env.update(self.env_dict) - - devices = self.stage_runtime_devices.get(stage_id) - if devices: - self._set_stage_device_env(stage_id, env, devices) - - cmd = self._build_stage_cmd(stage_id, headless=headless) - print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}") - # Capture each subprocess's stdout+stderr to a per-stage log file so - # debugging "Stage N exited before API server ready" doesn't rely on - # guessing; the file is surfaced in the RuntimeError message. - log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_{self.master_port}.log" - self._stage_log_paths = getattr(self, "_stage_log_paths", {}) - self._stage_log_paths[stage_id] = log_path - log_fh = open(log_path, "w", buffering=1) # noqa: SIM115 - closed in __exit__ - self._stage_log_files = getattr(self, "_stage_log_files", {}) - self._stage_log_files[stage_id] = log_fh - proc = subprocess.Popen( - cmd, - env=env, - cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - stdout=log_fh, - stderr=subprocess.STDOUT, - ) - self.stage_procs[stage_id] = proc - if stage_id == 0: - self.proc = proc - - def _ensure_stage_processes_alive(self) -> None: - for stage_id, proc in self.stage_procs.items(): - ret = proc.poll() - if ret is not None: - log_path = getattr(self, "_stage_log_paths", {}).get(stage_id) - tail = "" - if log_path and log_path.exists(): - try: - with open(log_path, encoding="utf-8", errors="replace") as f: - lines = f.readlines() - tail = "\n=== Last 60 lines of stage {} log ({}) ===\n{}".format( - stage_id, log_path, "".join(lines[-60:]) or "" - ) - except Exception as exc: # pragma: no cover - diagnostic only - tail = f"\n" - raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.{tail}") - - def _start_server(self) -> None: - ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]] - - self._launch_stage(0, headless=False) - time.sleep(2) - self._ensure_stage_processes_alive() - - for stage_id in ordered_stage_ids[1:]: - self._launch_stage(stage_id, headless=True) - - max_wait = 1200 - start_time = time.time() - while time.time() - start_time < max_wait: - self._ensure_stage_processes_alive() - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.settimeout(1) - result = sock.connect_ex((self.host, self.port)) - if result == 0: - print(f"OmniServerStageCli ready on {self.host}:{self.port}") - return - time.sleep(2) - - raise RuntimeError(f"OmniServerStageCli failed to start within {max_wait} seconds") - - def _dump_stage_logs_for_debug(self, head_lines: int = 300, tail_lines: int = 500) -> None: - """Tail each stage's subprocess log back to stdout on teardown. - - Stage subprocesses redirect stdout/stderr to ``/tmp/omni_stage_*.log`` - so we don't spam the main CI stream while tests run; but that also - hides engine init (KV cache size, Available KV cache memory, vLLM - engine config) when things go wrong. Dump them here so buildkite - captures them post-run. Head covers engine init; tail covers - whatever state the stage was in when it was torn down. - """ - log_paths = getattr(self, "_stage_log_paths", {}) or {} - for stage_id in sorted(log_paths): - log_path = log_paths[stage_id] - if not log_path or not log_path.exists(): - continue - try: - with open(log_path, encoding="utf-8", errors="replace") as f: - lines = f.readlines() - except Exception as exc: # pragma: no cover - diagnostic only - print(f"[OmniServerStageCli] stage {stage_id} log read failed: {exc}", flush=True) - continue - total = len(lines) - if total <= head_lines + tail_lines: - head_chunk = lines - tail_chunk = [] - elided = 0 - else: - head_chunk = lines[:head_lines] - tail_chunk = lines[-tail_lines:] - elided = total - head_lines - tail_lines - print(f"\n=== stage {stage_id} log HEAD ({log_path}) ===", flush=True) - print("".join(head_chunk).rstrip("\n"), flush=True) - if tail_chunk: - print(f"\n... [{elided} lines elided] ...", flush=True) - print(f"\n=== stage {stage_id} log TAIL ({log_path}) ===", flush=True) - print("".join(tail_chunk).rstrip("\n"), flush=True) - print(f"=== end stage {stage_id} log ===\n", flush=True) - - def __exit__(self, exc_type, exc_val, exc_tb): - self._dump_stage_logs_for_debug() - for stage_id in sorted(self.stage_procs, reverse=True): - proc = self.stage_procs[stage_id] - if proc.poll() is None: - self._kill_process_tree(proc.pid) - _run_pre_test_cleanup(enable_force=True) - _run_post_test_cleanup(enable_force=True) - cleanup_dist_env_and_memory() - - -def pytest_addoption(parser): - parser.addoption( - "--run-level", - action="store", - default="core_model", - choices=["core_model", "advanced_model"], - help="Test level to run: L2, L3", - ) - - -@pytest.fixture(scope="session") -def run_level(request) -> str: - """A command-line argument that specifies the level of tests to run in this session. - See https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/CI_5levels/""" - return request.config.getoption("--run-level") - - -_omni_server_lock = threading.Lock() - - -@pytest.fixture(scope="module") -def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]: - """Start vLLM-Omni through the standard or stage-CLI launcher. - - The fixture stays module-scoped because multi-stage initialization is costly. - The ``use_stage_cli`` flag on ``OmniServerParams`` routes the setup through the - stage-CLI harness while still reusing the same fixture grouping semantics. - """ - with _omni_server_lock: - params: OmniServerParams = request.param - model = model_prefix + params.model - port = params.port - stage_config_path = params.stage_config_path - if run_level == "advanced_model" and stage_config_path is not None: - with open(stage_config_path, encoding="utf-8") as f: - cfg = yaml.safe_load(f) or {} - # Strip ``load_format: dummy`` (CI overlay default) so advanced_model - # tests use real weights. New schema (``stages:``) writes the field - # flat at stage level; legacy schema (``stage_args:``) nests it as - # ``engine_args.load_format``. Handle both. - new_schema_stages = cfg.get("stages") - stage_key = "stages" if new_schema_stages is not None else "stage_args" - delete_path = "load_format" if new_schema_stages is not None else "engine_args.load_format" - stage_entries = cfg.get(stage_key, []) - stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage] - stage_config_path = modify_stage_config( - stage_config_path, - deletes={stage_key: {stage_id: [delete_path] for stage_id in stage_ids}}, - ) - - server_args = params.server_args or [] - if params.use_omni and params.stage_init_timeout is not None: - server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)] - else: - server_args = [*server_args, "--stage-init-timeout", "600"] - if params.init_timeout is not None: - server_args = [*server_args, "--init-timeout", str(params.init_timeout)] - else: - server_args = [*server_args, "--init-timeout", "900"] - if params.use_stage_cli: - if not params.use_omni: - raise ValueError("omni_server with use_stage_cli=True requires use_omni=True") - if stage_config_path is None: - raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path") - server_args += ["--stage-configs-path", stage_config_path] - - with OmniServerStageCli( - model, - stage_config_path, - server_args, - port=port, - env_dict=params.env_dict, - ) as server: - print("OmniServer started successfully") - yield server - print("OmniServer stopping...") - else: - if stage_config_path is not None: - server_args += ["--stage-configs-path", stage_config_path] - - with ( - OmniServer( - model, - server_args, - port=port, - env_dict=params.env_dict, - use_omni=params.use_omni, - ) - if port - else OmniServer( - model, - server_args, - env_dict=params.env_dict, - use_omni=params.use_omni, - ) - ) as server: - print("OmniServer started successfully") - yield server - print("OmniServer stopping...") - - print("OmniServer stopped") - - -@dataclass -class OmniResponse: - text_content: str | None = None - audio_data: list[str] | None = None - audio_content: str | None = None - audio_format: str | None = None - audio_bytes: bytes | None = None - similarity: float | None = None - e2e_latency: float | None = None - success: bool = False - error_message: str | None = None - cached_tokens: int | None = None - - -@dataclass -class DiffusionResponse: - text_content: str | None = None - images: list[Image.Image] | None = None - audios: list[Any] | None = None - videos: list[Any] | None = None - e2e_latency: float | None = None - success: bool = False - error_message: str | None = None - - -def _load_gender_pipeline(): - """ - Lazy-load a cached audio-classification pipeline for gender. - - We prefer the pipeline wrapper because it encapsulates processor/model loading - and avoids direct AutoProcessor.from_pretrained call sites in this file. - """ - global _GENDER_PIPELINE - if _GENDER_PIPELINE is not None: - return _GENDER_PIPELINE - - model_name = "7wolf/wav2vec2-base-gender-classification" - try: - # device=-1 forces CPU for pipeline. - _GENDER_PIPELINE = pipeline( - task="audio-classification", - model=model_name, - device=-1, - ) - return _GENDER_PIPELINE - except Exception as exc: # pragma: no cover - best-effort fallback - print(f"Warning: failed to create gender pipeline '{model_name}': {exc}") - _GENDER_PIPELINE = None - return None - - -def _median_pitch_hz_from_autocorr(mono: np.ndarray, sr: int) -> float | None: - """ - Rough median F0 (Hz) over short-time frames. Used to debias wav2vec2 gender head on TTS, - which often labels lower-pitched synthetic speech as female under load or on clean signals. - Returns None if the clip is too short or mostly unvoiced. - """ - x = np.asarray(mono, dtype=np.float64) - x = x - np.mean(x) - if x.size < int(0.15 * sr): - return None - frame_len = int(0.04 * sr) - hop = max(frame_len // 2, 1) - f0_min_hz, f0_max_hz = 70.0, 400.0 - lag_min = max(1, int(sr / f0_max_hz)) - lag_max = min(frame_len - 2, int(sr / f0_min_hz)) - if lag_max <= lag_min: - return None - win = np.hamming(frame_len) - pitches: list[float] = [] - for start in range(0, int(x.shape[0]) - frame_len, hop): - frame = x[start : start + frame_len] * win - frame = frame - np.mean(frame) - if float(np.sqrt(np.mean(frame**2))) < 1e-4: - continue - ac = np.correlate(frame, frame, mode="full")[frame_len - 1 :] - ac = ac / (float(ac[0]) + 1e-12) - region = ac[lag_min : lag_max + 1] - peak_rel = int(np.argmax(region)) - peak_lag = peak_rel + lag_min - if peak_lag <= 0: - continue - f0 = float(sr) / float(peak_lag) - if f0_min_hz <= f0 <= f0_max_hz: - pitches.append(f0) - if len(pitches) < 4: - return None - return float(np.median(np.asarray(pitches, dtype=np.float64))) - - -def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str: - """ - Estimate voice gender from audio using a small pre-trained classification model. - - Uses a cached `audio-classification` pipeline to classify the clip. - Returns 'male' / 'female' when the model confidence is >= 0.9 and the label - maps to one of these; otherwise returns 'unknown'. If the model is unavailable - or inference fails, returns 'unknown' to keep tests stable. - - Under concurrent tests, a global lock serializes pipeline calls (the HF pipeline is not - thread-safe). A coarse F0 median can correct systematic "male -> female" errors on TTS audio. - """ - data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32", always_2d=True) - if data.size == 0: - raise ValueError("Empty audio") - mono = np.mean(data, axis=1) - - try: - target_sr = 16000 - if int(sr) != target_sr and mono.size > 1: - src_len = int(mono.shape[0]) - dst_len = max(1, int(round(src_len * float(target_sr) / float(sr)))) - src_idx = np.arange(src_len, dtype=np.float32) - dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32) - mono = np.interp(dst_idx, src_idx, mono.astype(np.float32, copy=False)).astype(np.float32) - sr = target_sr - - median_f0 = _median_pitch_hz_from_autocorr(mono, sr) - - clf = _load_gender_pipeline() - if clf is None: - print("gender model not available, returning 'unknown'") - return "unknown" - - # transformers pipeline returns a list of {label, score} (highest score first). - with _GENDER_PIPELINE_LOCK: - outputs = clf(mono, sampling_rate=sr) - if not outputs: - return "unknown" - - top = outputs[0] - label = str(top.get("label", "")).lower() - conf = float(top.get("score", 0.0)) - - if conf < 0.5: - gender = "unknown" - # Some models use non-English labels (e.g., Russian). Normalize to 'male'/'female'. - elif ("female" in label) or ("жен" in label): - gender = "female" - elif ("male" in label) or ("муж" in label): - gender = "male" - else: - gender = "unknown" - - # Debias: wav2vec2 gender heads often call TTS / band-limited male speech "female". - # Low median F0 (~speech male range) + female label -> trust pitch when score is not overwhelming. - if gender == "female" and median_f0 is not None and median_f0 < 165.0 and conf < 0.88: - print(f"gender pitch assist: reclassifying female->male (median_f0={median_f0:.1f} Hz, conf={conf:.3f})") - gender = "male" - elif gender == "male" and median_f0 is not None and median_f0 > 230.0 and conf < 0.88: - print(f"gender pitch assist: reclassifying male->female (median_f0={median_f0:.1f} Hz, conf={conf:.3f})") - gender = "female" - - print( - f"gender classifier: label={label}, conf={conf:.3f}, gender={gender}" - + (f", median_f0={median_f0:.1f}Hz" if median_f0 is not None else "") - ) - return gender - except Exception as exc: # pragma: no cover - best-effort fallback - print(f"Warning: gender classification failed, returning 'unknown': {exc}") - return "unknown" - - -_PRESET_VOICE_GENDER_MAP: dict[str, str] = { - "serena": "female", - "uncle_fu": "male", - "chelsie": "female", - "clone": "female", - "ethan": "male", +""" +Root pytest entrypoint for the vLLM-Omni test suite. + +- `tests/conftest.py` stays thin: plugin registration + compatibility re-exports. +- Importable utilities live under `tests/helpers/`. +- Fixtures live under `tests/helpers/fixtures/` and are loaded via `pytest_plugins`. +""" + +from __future__ import annotations + +pytest_plugins = ( + "tests.helpers.fixtures.env", + "tests.helpers.fixtures.log", + "tests.helpers.fixtures.run_args", + "tests.helpers.fixtures.runtime", +) + + +def pytest_terminal_summary(terminalreporter, exitstatus, config): + # Marker for Buildkite log folding before pytest summary lines. + terminalreporter.write_line("--- Running Summary") + + +# Backward-compatible lazy re-exports. +# (Many tests still import from `tests.conftest`; migrate these imports to `tests.helpers.*` over time.) +# Keep these lazy so conftest import does not trigger heavy helper dependencies. +_ASSERTION_EXPORT_NAMES = ( + "assert_audio_speech_response", + "assert_diffusion_response", + "assert_image_diffusion_response", + "assert_image_valid", + "assert_omni_response", + "assert_video_diffusion_response", + "assert_video_valid", +) +_MEDIA_EXPORT_NAMES = ( + "convert_audio_bytes_to_text", + "convert_audio_file_to_text", + "cosine_similarity_text", + "decode_b64_image", + "generate_synthetic_audio", + "generate_synthetic_image", + "generate_synthetic_video", +) +_STAGE_CONFIG_EXPORT_NAMES = ("modify_stage_config",) +_RUNTIME_EXPORT_NAMES = ( + "DiffusionResponse", + "OmniResponse", + "OmniRunner", + "OmniRunnerHandler", + "OmniServer", + "OmniServerParams", + "OmniServerStageCli", + "OpenAIClientHandler", + "dummy_messages_from_mix_data", +) +_LAZY_EXPORT_MODULES = { + **{name: "tests.helpers.assertions" for name in _ASSERTION_EXPORT_NAMES}, + **{name: "tests.helpers.media" for name in _MEDIA_EXPORT_NAMES}, + **{name: "tests.helpers.stage_config" for name in _STAGE_CONFIG_EXPORT_NAMES}, + **{name: "tests.helpers.runtime" for name in _RUNTIME_EXPORT_NAMES}, } - - -def _assert_preset_voice_gender_from_audio( - audio_bytes: bytes | None, - voice_name: str | None, -) -> None: - """If ``voice_name`` matches a known preset, assert classifier gender matches (skip when unknown).""" - if not voice_name or not audio_bytes: - return - key = str(voice_name).lower() - expected_gender = _PRESET_VOICE_GENDER_MAP.get(key) - if expected_gender is None: - return - estimated_gender = _estimate_voice_gender_from_audio(audio_bytes) - print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}") - if estimated_gender != "unknown": - assert estimated_gender == expected_gender, ( - f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}" - ) - - -# Threshold aligned with _compute_pcm_hnr_db docstring (clean clone vs distorted). -_MIN_PCM_SPEECH_HNR_DB = 1.0 - - -def _compute_pcm_hnr_db(pcm_samples: np.ndarray, sr: int = _PCM_SPEECH_SAMPLE_RATE_HZ) -> float: - """Compute mean Harmonic-to-Noise Ratio (dB) for speech quality. - - Clean cloned speech has HNR > 1.2 dB; distorted speech (e.g. lost - ref_code decoder context) drops below 1.0 dB. - """ - frame_len = int(0.03 * sr) # 30ms frames - hop = frame_len // 2 - hnr_values: list[float] = [] - - for start in range(0, len(pcm_samples) - frame_len, hop): - frame = pcm_samples[start : start + frame_len].astype(np.float32, copy=False) - frame = frame - np.mean(frame) - if np.max(np.abs(frame)) < 0.01: - continue - ac = np.correlate(frame, frame, mode="full")[len(frame) - 1 :] - ac = ac / (ac[0] + 1e-10) - min_lag = int(sr / 400) - max_lag = min(int(sr / 80), len(ac)) - if min_lag >= max_lag: - continue - peak = float(np.max(ac[min_lag:max_lag])) - if 0 < peak < 1: - hnr_values.append(10 * np.log10(peak / (1 - peak + 1e-10))) - - return float(np.mean(hnr_values)) if hnr_values else 0.0 - - -def _assert_pcm_int16_speech_hnr(audio_bytes: bytes) -> None: - """Validate harmonic-to-noise ratio on raw int16 PCM from /v1/audio/speech.""" - assert audio_bytes is not None and len(audio_bytes) >= 2, "missing PCM bytes" - assert len(audio_bytes) % 2 == 0, "PCM byte length must be aligned to int16" - pcm_samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0 - hnr = _compute_pcm_hnr_db(pcm_samples) - print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {_MIN_PCM_SPEECH_HNR_DB} dB)") - assert hnr >= _MIN_PCM_SPEECH_HNR_DB, ( - f"Audio distortion detected: HNR={hnr:.2f} dB < {_MIN_PCM_SPEECH_HNR_DB} dB. " - "Voice clone decoder may be losing ref_code speaker context on later chunks." - ) - - -def assert_omni_response(response: OmniResponse, request_config: dict[str, Any], run_level): - """ - Validate response results. - - Args: - response: OmniResponse object - - Raises: - AssertionError: When the response does not meet validation criteria - """ - assert response.success, "The request failed." - e2e_latency = response.e2e_latency - if e2e_latency is not None: - print(f"the e2e latency is: {e2e_latency}") - - modalities = request_config.get("modalities", ["text", "audio"]) - - if run_level == "advanced_model": - if "audio" in modalities: - assert response.audio_content is not None, "No audio output is generated" - print(f"audio content is: {response.audio_content}") - speaker = request_config.get("speaker") - if speaker: - _assert_preset_voice_gender_from_audio( - response.audio_bytes, - speaker, - ) - - if "text" in modalities: - assert response.text_content is not None, "No text output is generated" - print(f"text content is: {response.text_content}") - - # Verify image description - word_types = ["text", "image", "audio", "video"] - keywords_dict = request_config.get("key_words", {}) - for word_type in word_types: - keywords = keywords_dict.get(word_type) - if "text" in modalities: - if keywords: - text_lower = response.text_content.lower() - assert any(str(kw).lower() in text_lower for kw in keywords), ( - "The output does not contain any of the keywords." - ) - else: - if keywords: - audio_lower = response.audio_content.lower() - assert any(str(kw).lower() in audio_lower for kw in keywords), ( - "The output does not contain any of the keywords." - ) - - # Verify similarity (Whisper transcript vs streamed/detokenized text) - if "text" in modalities and "audio" in modalities: - assert response.similarity is not None and response.similarity > 0.9, ( - "The audio content is not same as the text" - ) - print(f"similarity is: {response.similarity}") - - -def assert_audio_speech_response( - response: OmniResponse, - request_config: dict[str, Any], - run_level: str, -) -> None: - """ - Validate /v1/audio/speech response: success, optional format check, transcription similarity - and gender (non-PCM only for advanced_model), and int16 PCM HNR when response_format is pcm. - """ - assert response.success, "The request failed." - - req_fmt = request_config.get("response_format") - - if req_fmt == "pcm" and response.audio_bytes: - _assert_pcm_int16_speech_hnr(response.audio_bytes) - if response.audio_format: - assert "pcm" in response.audio_format.lower(), ( - f"Expected audio/pcm content-type, got {response.audio_format!r}" - ) - - elif req_fmt == "wav" and response.audio_format: - assert req_fmt in response.audio_format, ( - f"The response audio format {response.audio_format} don't match the request audio format {req_fmt}" - ) - - e2e_latency = response.e2e_latency - if e2e_latency is not None: - print(f"the avg e2e latency is: {e2e_latency}") - - if run_level == "advanced_model" and req_fmt != "pcm": - # Text–audio semantic similarity check (skipped for raw PCM: no Whisper transcript). - expected_text = request_config.get("input") - if expected_text: - transcript = (response.audio_content or "").strip() - print(f"audio content is: {transcript}") - print(f"input text is: {expected_text}") - similarity = cosine_similarity_text(transcript.lower(), expected_text.lower()) - print(f"Cosine similarity: {similarity:.3f}") - assert similarity > 0.9, ( - f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'" - ) - - # Voice gender consistency check (preset names in ``_PRESET_VOICE_GENDER_MAP``). - # When the estimator returns 'unknown', we treat it as inconclusive and do NOT fail the test. - _assert_preset_voice_gender_from_audio( - response.audio_bytes, - request_config.get("voice"), - ) - - -def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None): - """ - Validate diffusion response results. - - Dispatcher that routes validation to modality-specific assert functions. - - Args: - response: DiffusionResponse object. - request_config: Request configuration dictionary. - run_level: Test run level (e.g. "core_model", "advanced_model") - - Raises: - AssertionError: When the response does not meet validation criteria - KeyError: When the request_config does not contain necessary parameters for validation - """ - assert response.success, "The request failed." - - e2e_latency = response.e2e_latency - if e2e_latency is not None: - print(f"the avg e2e is: {e2e_latency}") - - has_any_content = any(content is not None for content in (response.images, response.videos, response.audios)) - assert has_any_content, "Response contains no images, videos, or audios" - - if response.images is not None: - assert_image_diffusion_response( - response=response, - request_config=request_config, - run_level=run_level, - ) - - if response.videos is not None: - assert_video_diffusion_response( - response=response, - request_config=request_config, - run_level=run_level, - ) - - if response.audios is not None: - assert_audio_diffusion_response( - response=response, - request_config=request_config, - run_level=run_level, - ) - - -class OpenAIClientHandler: - """ - OpenAI client handler class, encapsulating both streaming and non-streaming response processing logic. - - This class integrates OpenAI API request sending, response handling, and validation functionality, - supporting both single request and concurrent request modes. - """ - - def __init__( - self, host: str = "127.0.0.1", port: int = get_open_port(), api_key: str = "EMPTY", run_level: str = None - ): - """ - Initialize the OpenAI client. - - Args: - host: vLLM-Omni server host address - port: vLLM-Omni server port - api_key: API key (defaults to "EMPTY") - """ - self.base_url = f"http://{host}:{port}" - self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key) - self.run_level = run_level - - def _process_stream_omni_response(self, chat_completion) -> OmniResponse: - """ - Process streaming responses. - - Args: - chat_completion: OpenAI streaming response object - request_config: Request configuration dictionary - - Returns: - OmniResponse: Processed response object - """ - result = OmniResponse() - start_time = time.perf_counter() - - try: - text_content = "" - audio_data = [] - - for chunk in chat_completion: - for choice in chunk.choices: - # Get content data - if hasattr(choice, "delta"): - content = getattr(choice.delta, "content", None) - else: - content = None - - # Get modality type - modality = getattr(chunk, "modality", None) - - # Process content based on modality type - if modality == "audio" and content: - audio_data.append(content) - elif modality == "text" and content: - text_content += content if content else "" - - # Calculate end-to-end latency - result.e2e_latency = time.perf_counter() - start_time - - # Process audio and text content - audio_content = None - similarity = None - - if audio_data or text_content: - if audio_data: - merged_seg = _merge_base64_audio_to_segment(audio_data) - wav_buf = BytesIO() - merged_seg.export(wav_buf, format="wav") - result.audio_bytes = wav_buf.getvalue() - audio_content = convert_audio_bytes_to_text(result.audio_bytes) - if audio_content and text_content: - similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) - - # Populate result object - result.text_content = text_content - result.audio_data = audio_data - result.audio_content = audio_content - result.similarity = similarity - result.success = True - - except Exception as e: - result.error_message = f"Stream processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def _process_non_stream_omni_response(self, chat_completion) -> OmniResponse: - """ - Process non-streaming responses. - - Args: - chat_completion: OpenAI non-streaming response object - request_config: Request configuration dictionary - - Returns: - OmniResponse: Processed response object - """ - result = OmniResponse() - start_time = time.perf_counter() - - try: - audio_data = None - text_content = None - - # Iterate through all choices - for choice in chat_completion.choices: - # Process audio data - if hasattr(choice.message, "audio") and choice.message.audio is not None: - audio_message = choice.message - audio_data = audio_message.audio.data - - # Process text content - if hasattr(choice.message, "content") and choice.message.content is not None: - text_content = choice.message.content - - # Extract cached_tokens for prefix caching tests - usage = getattr(chat_completion, "usage", None) - if usage and (details := getattr(usage, "prompt_tokens_details", None)): - result.cached_tokens = details.cached_tokens - - # Calculate end-to-end latency - result.e2e_latency = time.perf_counter() - start_time - - # Process audio and text content - audio_content = None - similarity = None - - if audio_data or text_content: - if audio_data: - result.audio_bytes = base64.b64decode(audio_data) - audio_content = convert_audio_bytes_to_text(result.audio_bytes) - if audio_content and text_content: - similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) - - # Populate result object - result.text_content = text_content - result.audio_content = audio_content - result.similarity = similarity - result.success = True - - except Exception as e: - result.error_message = f"Non-stream processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def _process_diffusion_response(self, chat_completion) -> DiffusionResponse: - """ - Process diffusion responses (image generation/editing). - - Args: - chat_completion: OpenAI response object - - Returns: - DiffusionResponse: Processed response object - """ - result = DiffusionResponse() - start_time = time.perf_counter() - - try: - images = [] - # [TODO] reading video and audio output from API response for later validation - - for choice in chat_completion.choices: - if hasattr(choice.message, "content") and choice.message.content is not None: - content = choice.message.content - if isinstance(content, list): - for item in content: - if isinstance(item, dict): - image_url = item.get("image_url", {}).get("url") - else: - image_url_obj = getattr(item, "image_url", None) - image_url = getattr(image_url_obj, "url", None) if image_url_obj else None - if image_url and image_url.startswith("data:image"): - b64_data = image_url.split(",", 1)[1] - img = decode_b64_image(b64_data) - images.append(img) - - result.e2e_latency = time.perf_counter() - start_time - result.images = images if images else None - result.success = True - - except Exception as e: - result.error_message = f"Diffusion response processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def _process_stream_audio_speech_response(self, response, *, response_format: str | None = None) -> OmniResponse: - """ - Process streaming /v1/audio/speech responses into an OmniResponse. - - This mirrors _process_stream_omni_response but operates on low-level - audio bytes and produces an OmniResponse with audio_content filled - from Whisper transcription. - """ - result = OmniResponse() - start_time = time.perf_counter() - - try: - # Aggregate all audio bytes from the streaming response. - data = bytearray() - - # Preferred OpenAI helper. - if hasattr(response, "iter_bytes") and callable(getattr(response, "iter_bytes")): - for chunk in response.iter_bytes(): - if chunk: - data.extend(chunk) - else: - # Generic iterable-of-bytes fallback (e.g., generator or list of chunks). - try: - iterator = iter(response) - except TypeError: - iterator = None - - if iterator is not None: - for chunk in iterator: - if not chunk: - continue - if isinstance(chunk, (bytes, bytearray)): - data.extend(chunk) - elif hasattr(chunk, "data"): - data.extend(chunk.data) # type: ignore[arg-type] - elif hasattr(chunk, "content"): - data.extend(chunk.content) # type: ignore[arg-type] - else: - raise TypeError(f"Unsupported stream chunk type: {type(chunk)}") - else: - raise TypeError(f"Unsupported audio speech streaming response type: {type(response)}") - - raw_bytes = bytes(data) - if response_format == "pcm": - transcript = None - else: - transcript = convert_audio_bytes_to_text(raw_bytes) - - # Populate OmniResponse. - result.audio_bytes = raw_bytes - result.audio_content = transcript - result.e2e_latency = time.perf_counter() - start_time - result.success = True - result.audio_format = getattr(response, "response", None) - if result.audio_format is not None: - result.audio_format = result.audio_format.headers.get("content-type", "") - - except Exception as e: - result.error_message = f"Audio speech stream processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def _process_non_stream_audio_speech_response( - self, response, *, response_format: str | None = None - ) -> OmniResponse: - """ - Process non-streaming /v1/audio/speech responses into an OmniResponse. - - This mirrors _process_non_stream_omni_response but for the binary - audio payload returned by audio.speech.create. - """ - result = OmniResponse() - start_time = time.perf_counter() - - try: - # OpenAI non-streaming audio.speech.create returns HttpxBinaryResponseContent (.read() or .content) - if hasattr(response, "read") and callable(getattr(response, "read")): - raw_bytes = response.read() - elif hasattr(response, "content"): - raw_bytes = response.content # type: ignore[assignment] - else: - raise TypeError(f"Unsupported audio speech response type: {type(response)}") - - if response_format == "pcm": - transcript = None - else: - transcript = convert_audio_bytes_to_text(raw_bytes) - - result.audio_bytes = raw_bytes - result.audio_content = transcript - result.e2e_latency = time.perf_counter() - start_time - result.success = True - result.audio_format = getattr(response, "response", None) - if result.audio_format is not None: - result.audio_format = result.audio_format.headers.get("content-type", "") - - except Exception as e: - result.error_message = f"Audio speech non-stream processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: - """ - Send OpenAI requests. - - Args: - request_config: Request configuration dictionary containing parameters like model, messages, stream. - Optional ``use_audio_in_video`` (bool): when true, sets - ``extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}`` for Qwen-Omni video+audio - extraction. - Optional top-level ``speaker`` (str): Qwen3-Omni preset TTS speaker name; sent as - ``extra_body["speaker"]`` to ``chat.completions.create``. - request_num: Number of requests, defaults to 1 (single request) - - Returns: - List[OmniResponse]: List of response objects - """ - - responses = [] - stream = request_config.get("stream", False) - modalities = request_config.get("modalities", ["text", "audio"]) - - extra_body: dict[str, Any] = {} - if "speaker" in request_config: - extra_body["speaker"] = request_config["speaker"] - if request_config.get("use_audio_in_video"): - mm = dict(extra_body.get("mm_processor_kwargs") or {}) - mm["use_audio_in_video"] = True - extra_body["mm_processor_kwargs"] = mm - extra_body_arg: dict[str, Any] | None = extra_body if extra_body else None - - create_kwargs: dict[str, Any] = { - "model": request_config.get("model"), - "messages": request_config.get("messages"), - "stream": stream, - "modalities": modalities, - } - if extra_body_arg is not None: - create_kwargs["extra_body"] = extra_body_arg - - if request_num == 1: - # Send single request - chat_completion = self.client.chat.completions.create(**create_kwargs) - - if stream: - response = self._process_stream_omni_response(chat_completion) - else: - response = self._process_non_stream_omni_response(chat_completion) - - assert_omni_response(response, request_config, run_level=self.run_level) - responses.append(response) - - else: - # Send concurrent requests: run create + process in worker so e2e_latency includes full round-trip. - def _one_omni_request(): - start = time.perf_counter() - worker_kwargs: dict[str, Any] = { - "model": request_config.get("model"), - "messages": request_config.get("messages"), - "modalities": modalities, - "stream": stream, - } - if extra_body_arg is not None: - worker_kwargs["extra_body"] = extra_body_arg - chat_completion = self.client.chat.completions.create(**worker_kwargs) - if stream: - response = self._process_stream_omni_response(chat_completion) - else: - response = self._process_non_stream_omni_response(chat_completion) - response.e2e_latency = time.perf_counter() - start - return response - - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = [executor.submit(_one_omni_request) for _ in range(request_num)] - for future in concurrent.futures.as_completed(futures): - response = future.result() - assert_omni_response(response, request_config, run_level=self.run_level) - responses.append(response) - - return responses - - def send_audio_speech_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: - """ - Call the /v1/audio/speech endpoint using the same configuration-dict - style as send_omni_request, but via the OpenAI Python client's - audio.speech APIs. - - Expected keys in request_config: - - model: model name/path (required) - - input: text to synthesize (required) - - response_format: audio format such as "wav" or "pcm" (optional) - - task_type, ref_text, ref_audio: TTS-specific extras (optional, passed via extra_body) - - timeout: request timeout in seconds (float, optional, default 120.0) - - stream: whether to use streaming API (bool, optional, default False) - """ - timeout = float(request_config.get("timeout", 120.0)) - - model = request_config["model"] - text_input = request_config["input"] - stream = bool(request_config.get("stream", False)) - voice = request_config.get("voice", None) - - # Standard OpenAI param: use omit when not provided to keep default behavior. - response_format = request_config.get("response_format", omit) - - # Qwen3-TTS custom fields, forwarded via extra_body. - extra_body: dict[str, Any] = {} - # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params. - for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens"): - if key in request_config: - extra_body[key] = request_config[key] - - responses: list[OmniResponse] = [] - - speech_fmt: str | None = None if response_format is omit else str(response_format).lower() - - if request_num == 1: - if stream: - # Use streaming response helper. - with self.client.audio.speech.with_streaming_response.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) as resp: - omni_resp = self._process_stream_audio_speech_response(resp, response_format=speech_fmt) - else: - # Non-streaming response. - resp = self.client.audio.speech.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) - omni_resp = self._process_non_stream_audio_speech_response(resp, response_format=speech_fmt) - - assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) - responses.append(omni_resp) - return responses - else: - # request_num > 1: concurrent requests (use same params as single-request path) - - if stream: - - def _stream_task(): - with self.client.audio.speech.with_streaming_response.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) as resp: - return self._process_stream_audio_speech_response(resp, response_format=speech_fmt) - - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = [executor.submit(_stream_task) for _ in range(request_num)] - for future in concurrent.futures.as_completed(futures): - omni_resp = future.result() - assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) - responses.append(omni_resp) - else: - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = [] - for _ in range(request_num): - future = executor.submit( - self.client.audio.speech.create, - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) - futures.append(future) - - for future in concurrent.futures.as_completed(futures): - resp = future.result() - omni_resp = self._process_non_stream_audio_speech_response(resp, response_format=speech_fmt) - assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) - responses.append(omni_resp) - - return responses - - def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[DiffusionResponse]: - """ - Send OpenAI requests for diffusion models. - - Args: - request_config: Request configuration dictionary containing parameters like model, messages - request_num: Number of requests to send concurrently, defaults to 1 (single request) - Returns: - List[DiffusionResponse]: List of response objects - """ - responses: list[DiffusionResponse] = [] - stream = request_config.get("stream", False) - modalities = request_config.get("modalities", omit) # Most diffusion models don't require modalities param - extra_body = request_config.get("extra_body", None) - - if stream: - raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test") - - if request_num == 1: - # Send single request - chat_completion = self.client.chat.completions.create( - model=request_config.get("model"), - messages=request_config.get("messages"), - extra_body=extra_body, - modalities=modalities, - ) - - response = self._process_diffusion_response(chat_completion) - assert_diffusion_response(response, request_config, run_level=self.run_level) - responses.append(response) - - else: - # Send concurrent requests - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = [] - - # Submit all request tasks - for _ in range(request_num): - future = executor.submit( - self.client.chat.completions.create, - model=request_config.get("model"), - messages=request_config.get("messages"), - modalities=modalities, - extra_body=extra_body, - ) - futures.append(future) - - # Process completed tasks - for future in concurrent.futures.as_completed(futures): - chat_completion = future.result() - response = self._process_diffusion_response(chat_completion) - assert_diffusion_response(response, request_config, run_level=self.run_level) - responses.append(response) - - return responses - - def send_video_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: - """ - Send native /v1/videos requests. - """ - if request_num != 1: - raise NotImplementedError("Concurrent video diffusion requests are not currently implemented") - - if request_config.get("stream", False): - raise NotImplementedError("Streaming is not currently implemented for video diffusion e2e test") - - form_data = request_config.get("form_data") - if not isinstance(form_data, dict): - raise ValueError("Video request_config must contain 'form_data'") - - if not form_data.get("prompt"): - raise ValueError("Video request_config['form_data'] must contain 'prompt'") - - normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None} - - files: dict[str, tuple[str, BytesIO, str]] = {} - image_reference = request_config.get("image_reference") - if image_reference: - if image_reference.startswith("data:image"): - header, encoded = image_reference.split(",", 1) - content_type = header.split(";")[0].removeprefix("data:") - extension = content_type.split("/")[-1] - file_data = base64.b64decode(encoded) - - files["input_reference"] = ( - f"reference.{extension}", - BytesIO(file_data), - content_type, - ) - else: - normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference}) - - result = DiffusionResponse() - start_time = time.perf_counter() - - try: - create_url = self._build_url("/v1/videos") - response = requests.post( - create_url, - data=normalized_form_data, - files=files, - headers={"Accept": "application/json"}, - timeout=60, - ) - response.raise_for_status() - - job_data = response.json() - video_id = job_data["id"] - - self._wait_until_video_completed(video_id) - - video_content = self._download_video_content(video_id) - - result.success = True - result.videos = [video_content] - result.e2e_latency = time.perf_counter() - start_time - - assert_diffusion_response(result, request_config, run_level=self.run_level) - - except Exception as e: - result.success = False - result.error_message = f"Diffusion response processing error: {e}" - assert False, result.error_message - - return [result] - - def _wait_until_video_completed( - self, - video_id: str, - poll_interval_seconds: int = 2, - timeout_seconds: int = 300, - ) -> None: - status_url = self._build_url(f"/v1/videos/{video_id}") - deadline = time.monotonic() + timeout_seconds - - while time.monotonic() < deadline: - status_resp = requests.get( - status_url, - headers={"Accept": "application/json"}, - timeout=30, - ) - status_resp.raise_for_status() - - status_data = status_resp.json() - current_status = status_data["status"] - - if current_status == "completed": - return - - if current_status == "failed": - error_msg = status_data.get("last_error", "Unknown error") - raise RuntimeError(f"Job failed: {error_msg}") - - time.sleep(poll_interval_seconds) - - raise TimeoutError(f"Video job {video_id} did not complete within {timeout_seconds}s") - - def _download_video_content(self, video_id: str) -> bytes: - download_url = self._build_url(f"/v1/videos/{video_id}/content") - video_resp = requests.get(download_url, stream=True, timeout=60) - video_resp.raise_for_status() - - video_bytes = BytesIO() - for chunk in video_resp.iter_content(chunk_size=8192): - if chunk: - video_bytes.write(chunk) - - return video_bytes.getvalue() - - def _build_url(self, path: str) -> str: - return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}" - - -@pytest.fixture -def openai_client(request: pytest.FixtureRequest, run_level: str): - """Create OpenAIClientHandler fixture to facilitate communication with OmniServer - with encapsulated request sending, concurrent requests, response handling, and validation.""" - server = request.getfixturevalue("omni_server") - return OpenAIClientHandler(host=server.host, port=server.port, api_key="EMPTY", run_level=run_level) - - -class OmniRunner: - """ - Offline test runner for Omni models. - """ - - def __init__( - self, - model_name: str, - seed: int = 42, - stage_init_timeout: int = 600, - batch_timeout: int = 10, - init_timeout: int = 900, - shm_threshold_bytes: int = 65536, - log_stats: bool = False, - stage_configs_path: str | None = None, - **kwargs, - ) -> None: - """ - Initialize an OmniRunner for testing. - - Args: - model_name: The model name or path - seed: Random seed for reproducibility - stage_init_timeout: Timeout for initializing a single stage in seconds - batch_timeout: Timeout for batching in seconds - init_timeout: Timeout for initializing stages in seconds - shm_threshold_bytes: Threshold for using shared memory - log_stats: Enable detailed statistics logging - stage_configs_path: Optional path to YAML stage config file - **kwargs: Additional arguments passed to Omni - """ - cleanup_dist_env_and_memory() - _run_pre_test_cleanup(enable_force=True) - _run_post_test_cleanup(enable_force=True) - self.model_name = model_name - self.seed = seed - - self.omni = Omni( - model=model_name, - log_stats=log_stats, - stage_init_timeout=stage_init_timeout, - batch_timeout=batch_timeout, - init_timeout=init_timeout, - shm_threshold_bytes=shm_threshold_bytes, - stage_configs_path=stage_configs_path, - **kwargs, - ) - - def _estimate_prompt_len( - self, - additional_information: dict[str, Any], - model_name: str, - _cache: dict[str, Any] = {}, - ) -> int: - """Estimate prompt_token_ids placeholder length for the Talker stage. - - The AR Talker replaces all input embeddings via ``preprocess``, so the - placeholder values are irrelevant but the **length** must match the - embeddings that ``preprocess`` will produce. - """ - try: - from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig - from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( - Qwen3TTSTalkerForConditionalGeneration, - ) - - if model_name not in _cache: - from transformers import AutoTokenizer - - tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left") - cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True) - _cache[model_name] = (tok, getattr(cfg, "talker_config", None)) - - tok, tcfg = _cache[model_name] - task_type = (additional_information.get("task_type") or ["CustomVoice"])[0] - return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( - additional_information=additional_information, - task_type=task_type, - tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"], - codec_language_id=getattr(tcfg, "codec_language_id", None), - spk_is_dialect=getattr(tcfg, "spk_is_dialect", None), - ) - except Exception as exc: - logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc) - return 2048 - - def get_default_sampling_params_list(self) -> list[OmniSamplingParams]: - """ - Get a list of default sampling parameters for all stages. - - Returns: - List of SamplingParams with default decoding for each stage - """ - if not hasattr(self.omni, "default_sampling_params_list"): - raise AttributeError("Omni.default_sampling_params_list is not available") - return list(self.omni.default_sampling_params_list) - - def get_omni_inputs( - self, - prompts: list[str] | str, - system_prompt: str | None = None, - audios: PromptAudioInput = None, - images: PromptImageInput = None, - videos: PromptVideoInput = None, - mm_processor_kwargs: dict[str, Any] | None = None, - modalities: list[str] | None = None, - ) -> list[TextPrompt]: - """ - Construct Omni input format from prompts and multimodal data. - - Args: - prompts: Text prompt(s) - either a single string or list of strings - system_prompt: Optional system prompt (defaults to Qwen system prompt) - audios: Audio input(s) - tuple of (audio_array, sample_rate) or list of tuples - images: Image input(s) - PIL Image or list of PIL Images - videos: Video input(s) - numpy array or list of numpy arrays - mm_processor_kwargs: Optional processor kwargs (e.g., use_audio_in_video) - - Returns: - List of prompt dictionaries suitable for Omni.generate() - """ - if system_prompt is None: - system_prompt = ( - "You are Qwen, a virtual human developed by the Qwen Team, Alibaba " - "Group, capable of perceiving auditory and visual inputs, as well as " - "generating text and speech." - ) - - video_padding_token = "<|VIDEO|>" - image_padding_token = "<|IMAGE|>" - audio_padding_token = "<|AUDIO|>" - - if "Qwen3-Omni-30B-A3B-Instruct" in self.model_name: - video_padding_token = "<|video_pad|>" - image_padding_token = "<|image_pad|>" - audio_padding_token = "<|audio_pad|>" - elif "Ming-flash-omni" in self.model_name: - video_padding_token = "