Zeyu Huang | 黃澤宇andHongsheng Liu
a19098233b
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2011 lines
72 KiB
Python
2011 lines
72 KiB
Python
import base64
|
||
import datetime
|
||
import io
|
||
import math
|
||
import os
|
||
import random
|
||
|
||
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 gc
|
||
import multiprocessing
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
from collections.abc import Generator
|
||
from dataclasses import dataclass
|
||
from io import BytesIO
|
||
from pathlib import Path
|
||
from typing import Any, NamedTuple
|
||
|
||
import imageio.v3 as iio
|
||
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 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.entrypoints.omni import Omni
|
||
from vllm_omni.inputs.data import OmniSamplingParams
|
||
from vllm_omni.outputs import OmniRequestOutput
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
|
||
PromptImageInput = list[Any] | Any | None
|
||
PromptVideoInput = list[Any] | Any | None
|
||
|
||
|
||
class OmniServerParams(NamedTuple):
|
||
model: str
|
||
port: int | None = None
|
||
stage_config_path: str | None = None
|
||
server_args: list[str] | None = None
|
||
|
||
|
||
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} in {image.name}"
|
||
if height is not None:
|
||
assert image.height == height, f"Expected height={height}, got {image.height} in {image.name}"
|
||
return image
|
||
|
||
|
||
def assert_video_valid(frames: Path | np.ndarray, *, width: int, height: int, num_frames: int) -> None:
|
||
"""Assert the MP4 has the expected resolution and exact frame count."""
|
||
if isinstance(frames, Path):
|
||
assert frames.exists(), f"Video not found: {frames}"
|
||
frames = iio.imread(str(frames), plugin="pyav", index=None)
|
||
assert frames.shape[0] == num_frames, f"Expected {num_frames} frames, got {frames.shape[0]}"
|
||
assert frames.shape[1] == height, f"Expected height={height}, got {frames.shape[1]}"
|
||
assert frames.shape[2] == width, f"Expected width={width}, got {frames.shape[2]}"
|
||
|
||
|
||
def assert_audio_valid(path: Path, *, sample_rate: int, channels: int, duration_s: float) -> None:
|
||
"""Assert the WAV has the expected sample rate, channel count, and duration."""
|
||
|
||
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}"
|
||
expected_frames = int(duration_s * sample_rate)
|
||
assert info.frames == expected_frames, (
|
||
f"Expected {expected_frames} 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("GPU cleanup disabled")
|
||
return
|
||
|
||
print("Pre-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 tempfile
|
||
|
||
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
|
||
if save_to_file and output_path:
|
||
result["file_path"] = output_path
|
||
|
||
return result
|
||
|
||
|
||
def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_file: bool = False) -> dict[str, Any]:
|
||
"""Generate synthetic video with bouncing balls and return base64 string."""
|
||
|
||
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,
|
||
}
|
||
video_bytes = None
|
||
saved_file_path = None
|
||
|
||
buffer = io.BytesIO()
|
||
writer_kwargs = {
|
||
"format": "mp4",
|
||
"fps": 30,
|
||
"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}",
|
||
],
|
||
}
|
||
|
||
if save_to_file:
|
||
import datetime
|
||
|
||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
output_path = f"video_{width}x{height}_{timestamp}.mp4"
|
||
try:
|
||
with imageio.get_writer(output_path, **writer_kwargs) as writer:
|
||
for frame in video_frames:
|
||
writer.append_data(frame)
|
||
|
||
saved_file_path = output_path
|
||
print(f"Video saved to: {saved_file_path}")
|
||
with open(output_path, "rb") as f:
|
||
video_bytes = f.read()
|
||
|
||
except Exception as e:
|
||
print(f"Warning: Failed to save video to file {output_path}: {e}")
|
||
save_to_file = False
|
||
|
||
if not save_to_file or video_bytes is None:
|
||
with imageio.get_writer(buffer, **writer_kwargs) as writer:
|
||
for frame in video_frames:
|
||
writer.append_data(frame)
|
||
|
||
buffer.seek(0)
|
||
video_bytes = buffer.read()
|
||
|
||
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 re
|
||
|
||
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)
|
||
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)
|
||
|
||
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_{int(time.time())}"
|
||
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 _whisper_transcribe_in_current_process(output_path: str) -> str:
|
||
import whisper
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
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:
|
||
# Sync GPU so in-flight ops finish before we free the model; otherwise
|
||
# freed memory may not show up until those ops complete.
|
||
if torch is not None and torch.cuda.is_available():
|
||
torch.cuda.synchronize()
|
||
del model
|
||
gc.collect()
|
||
if torch is not None and torch.cuda.is_available():
|
||
torch.cuda.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."""
|
||
# Import locally to avoid impacting test module import time.
|
||
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 merge_base64_and_convert_to_text(base64_list):
|
||
"""
|
||
Merge a list of base64 encoded audio data and convert to text.
|
||
"""
|
||
from pydub import AudioSegment
|
||
|
||
merged_audio = None
|
||
for base64_data in base64_list:
|
||
audio_data = base64.b64decode(base64_data.split(",", 1)[-1])
|
||
seg = AudioSegment.from_file(io.BytesIO(audio_data))
|
||
if merged_audio is None:
|
||
merged_audio = seg
|
||
else:
|
||
merged_audio += seg
|
||
output_path = f"./test_{int(time.time())}"
|
||
merged_audio.export(output_path, format="wav")
|
||
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: {'runtime.max_batch_size': 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")
|
||
|
||
# Apply deletions first
|
||
if deletes:
|
||
for key, value in deletes.items():
|
||
if key == "stage_args":
|
||
if value and isinstance(value, dict):
|
||
stage_args = config.get("stage_args", [])
|
||
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:
|
||
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}")
|
||
|
||
# Delete specified paths in this stage
|
||
for path in delete_paths:
|
||
if path: # Skip empty paths
|
||
delete_by_path(target_stage, 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 == "stage_args":
|
||
if value and isinstance(value, dict):
|
||
stage_args = config.get("stage_args", [])
|
||
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 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 path:
|
||
# Direct key assignment (e.g., updating a list value)
|
||
target_stage[path] = val
|
||
else:
|
||
# Dot-separated path (e.g., nested dict access)
|
||
apply_update(target_stage, path, val)
|
||
elif "." in key:
|
||
# Apply using dot-separated path
|
||
apply_update(config, key, value)
|
||
else:
|
||
# Direct top-level key
|
||
config[key] = value
|
||
|
||
# Save to new file with timestamp
|
||
timestamp = int(time.time())
|
||
base_name = yaml_path.rsplit(".", 1)[0] if "." in yaml_path else yaml_path
|
||
output_path = f"{base_name}_{timestamp}.yaml"
|
||
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
yaml.dump(config, f, default_flow_style=None, sort_keys=False, allow_unicode=True, indent=2)
|
||
|
||
return 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,
|
||
) -> 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.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,
|
||
"--omni",
|
||
"--host",
|
||
self.host,
|
||
"--port",
|
||
str(self.port),
|
||
] + 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()
|
||
|
||
|
||
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 server as a subprocess with actual model weights.
|
||
Uses session scope so the server starts only once for the entire test session.
|
||
Multi-stage initialization can take 10-20+ minutes.
|
||
"""
|
||
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:
|
||
stage_config_path = modify_stage_config(
|
||
stage_config_path,
|
||
deletes={
|
||
"stage_args": {
|
||
0: ["engine_args.load_format"],
|
||
1: ["engine_args.load_format"],
|
||
2: ["engine_args.load_format"],
|
||
}
|
||
},
|
||
)
|
||
|
||
server_args = params.server_args or []
|
||
server_args = ["--stage-init-timeout", "120", *server_args]
|
||
if stage_config_path is not None:
|
||
server_args += ["--stage-configs-path", stage_config_path]
|
||
|
||
with OmniServer(model, server_args, port=port) if port else OmniServer(model, server_args) 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
|
||
similarity: float | None = None
|
||
e2e_latency: float | None = None
|
||
success: bool = False
|
||
error_message: str | 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 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 avg e2e latency is: {e2e_latency}")
|
||
|
||
modalities = request_config.get("modalities", ["text", "audio"])
|
||
|
||
if "audio" in modalities:
|
||
assert response.audio_content is not None, "No audio output is generated"
|
||
print(f"audio content is: {response.audio_content}")
|
||
|
||
if "text" in modalities:
|
||
assert response.text_content is not None, "No text output is generated"
|
||
print(f"text content is: {response.text_content}")
|
||
|
||
if run_level == "advanced_model":
|
||
# 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:
|
||
assert any(keyword in response.text_content.lower() for keyword in keywords), (
|
||
"The output does not contain any of the keywords."
|
||
)
|
||
else:
|
||
if keywords:
|
||
assert any(keyword in response.audio_content.lower() for keyword in keywords), (
|
||
"The output does not contain any of the keywords."
|
||
)
|
||
|
||
# Verify similarity
|
||
if "text" in modalities and "audio" in modalities:
|
||
assert response.similarity > 0.9, "The audio content is not same as the text"
|
||
print(f"similarity is: {response.similarity}")
|
||
|
||
|
||
def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None):
|
||
"""
|
||
Validate diffusion response results.
|
||
|
||
Args:
|
||
response: DiffusionResponse object. Any not-None content will be validated based on the request_config.
|
||
request_config: Request configuration dictionary containing parameters like model, messages, extra_body.
|
||
When validating a certain modality, the corresponding params in request_config['extra_body'] must present.
|
||
It will be used to check against the multimedia file in the response.
|
||
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}")
|
||
|
||
extra_body = request_config.get("extra_body", {})
|
||
|
||
num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1)
|
||
|
||
if response.images is not None:
|
||
assert len(response.images) > 0, "No images in response"
|
||
assert len(response.images) == num_outputs_per_prompt, (
|
||
f"Expected {num_outputs_per_prompt} images, got {len(response.images)}"
|
||
)
|
||
if run_level == "advanced_model":
|
||
expected_width = extra_body["width"] # intend to raise KeyError
|
||
expected_height = extra_body["height"] # intend to raise KeyError
|
||
for img in response.images:
|
||
assert_image_valid(img, width=expected_width, height=expected_height)
|
||
if response.videos is not None:
|
||
raise NotImplementedError(
|
||
"Video validation is not implemented yet"
|
||
) # consider using assert_video_valid defined above
|
||
if response.audios is not None:
|
||
raise NotImplementedError(
|
||
"Audio validation is not implemented yet"
|
||
) # consider using assert_audio_valid defined above
|
||
|
||
|
||
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.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:
|
||
audio_content = merge_base64_and_convert_to_text(audio_data)
|
||
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
|
||
|
||
# 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:
|
||
audio_content = convert_audio_to_text(audio_data)
|
||
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 hasattr(item, "image_url") and item.image_url is not None:
|
||
image_url = item.image_url.url
|
||
if 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 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
|
||
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"])
|
||
|
||
if request_num == 1:
|
||
# Send single request
|
||
chat_completion = self.client.chat.completions.create(
|
||
model=request_config.get("model"),
|
||
messages=request_config.get("messages"),
|
||
stream=stream,
|
||
modalities=modalities,
|
||
)
|
||
|
||
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
|
||
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,
|
||
stream=stream,
|
||
)
|
||
futures.append(future)
|
||
|
||
# Process completed tasks
|
||
for future in concurrent.futures.as_completed(futures):
|
||
chat_completion = future.result()
|
||
|
||
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)
|
||
|
||
return responses
|
||
|
||
def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
|
||
"""
|
||
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[OmniResponse]: List of response objects
|
||
"""
|
||
responses = []
|
||
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
|
||
|
||
|
||
@pytest.fixture
|
||
def openai_client(omni_server: OmniServer, run_level: str):
|
||
"""Create OpenAIClientHandler fixture to facilitate communication with OmniServer
|
||
with encapsulated request sending, concurrent requests, response handling, and validation."""
|
||
return OpenAIClientHandler(host=omni_server.host, port=omni_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 = 300,
|
||
batch_timeout: int = 10,
|
||
init_timeout: int = 300,
|
||
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 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
|
||
"""
|
||
return [st.default_sampling_params for st in self.omni.stage_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|>"
|
||
|
||
if isinstance(prompts, str):
|
||
prompts = [prompts]
|
||
|
||
def _normalize_mm_input(mm_input, num_prompts):
|
||
if mm_input is None:
|
||
return [None] * num_prompts
|
||
if isinstance(mm_input, list):
|
||
if len(mm_input) != num_prompts:
|
||
raise ValueError(
|
||
f"Multimodal input list length ({len(mm_input)}) must match prompts length ({num_prompts})"
|
||
)
|
||
return mm_input
|
||
return [mm_input] * num_prompts
|
||
|
||
num_prompts = len(prompts)
|
||
audios_list = _normalize_mm_input(audios, num_prompts)
|
||
images_list = _normalize_mm_input(images, num_prompts)
|
||
videos_list = _normalize_mm_input(videos, num_prompts)
|
||
|
||
omni_inputs = []
|
||
for i, prompt_text in enumerate(prompts):
|
||
user_content = ""
|
||
multi_modal_data = {}
|
||
|
||
audio = audios_list[i]
|
||
if audio is not None:
|
||
if isinstance(audio, list):
|
||
for _ in audio:
|
||
user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
|
||
multi_modal_data["audio"] = audio
|
||
else:
|
||
user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
|
||
multi_modal_data["audio"] = audio
|
||
|
||
image = images_list[i]
|
||
if image is not None:
|
||
if isinstance(image, list):
|
||
for _ in image:
|
||
user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
|
||
multi_modal_data["image"] = image
|
||
else:
|
||
user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
|
||
multi_modal_data["image"] = image
|
||
|
||
video = videos_list[i]
|
||
if video is not None:
|
||
if isinstance(video, list):
|
||
for _ in video:
|
||
user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
|
||
multi_modal_data["video"] = video
|
||
else:
|
||
user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
|
||
multi_modal_data["video"] = video
|
||
|
||
user_content += prompt_text
|
||
|
||
full_prompt = (
|
||
f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
|
||
f"<|im_start|>user\n{user_content}<|im_end|>\n"
|
||
f"<|im_start|>assistant\n"
|
||
)
|
||
|
||
input_dict: TextPrompt = {"prompt": full_prompt}
|
||
if multi_modal_data:
|
||
input_dict["multi_modal_data"] = multi_modal_data
|
||
if modalities:
|
||
input_dict["modalities"] = modalities
|
||
if mm_processor_kwargs:
|
||
input_dict["mm_processor_kwargs"] = mm_processor_kwargs
|
||
|
||
omni_inputs.append(input_dict)
|
||
|
||
return omni_inputs
|
||
|
||
def generate(
|
||
self,
|
||
prompts: list[TextPrompt],
|
||
sampling_params_list: list[OmniSamplingParams] | None = None,
|
||
) -> list[OmniRequestOutput]:
|
||
"""
|
||
Generate outputs for the given prompts.
|
||
|
||
Args:
|
||
prompts: List of prompt dictionaries with 'prompt' and optionally
|
||
'multi_modal_data' keys
|
||
sampling_params_list: List of sampling parameters for each stage.
|
||
If None, uses default parameters.
|
||
|
||
Returns:
|
||
List of OmniRequestOutput objects from stages with final_output=True
|
||
"""
|
||
if sampling_params_list is None:
|
||
sampling_params_list = self.get_default_sampling_params_list()
|
||
|
||
return self.omni.generate(prompts, sampling_params_list)
|
||
|
||
def generate_multimodal(
|
||
self,
|
||
prompts: list[str] | str,
|
||
sampling_params_list: list[OmniSamplingParams] | None = None,
|
||
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[OmniRequestOutput]:
|
||
"""
|
||
Convenience method to generate with multimodal inputs.
|
||
|
||
Args:
|
||
prompts: Text prompt(s)
|
||
sampling_params_list: List of sampling parameters for each stage
|
||
system_prompt: Optional system prompt
|
||
audios: Audio input(s)
|
||
images: Image input(s)
|
||
videos: Video input(s)
|
||
mm_processor_kwargs: Optional processor kwargs
|
||
|
||
Returns:
|
||
List of OmniRequestOutput objects from stages with final_output=True
|
||
"""
|
||
omni_inputs = self.get_omni_inputs(
|
||
prompts=prompts,
|
||
system_prompt=system_prompt,
|
||
audios=audios,
|
||
images=images,
|
||
videos=videos,
|
||
mm_processor_kwargs=mm_processor_kwargs,
|
||
modalities=modalities,
|
||
)
|
||
return self.generate(omni_inputs, sampling_params_list)
|
||
|
||
def _cleanup_process(self):
|
||
try:
|
||
keywords = ["enginecore"]
|
||
|
||
for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
|
||
try:
|
||
cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else ""
|
||
name = proc.name().lower()
|
||
|
||
is_process = any(keyword in cmdline for keyword in keywords) or any(
|
||
keyword in name for keyword in keywords
|
||
)
|
||
|
||
if is_process:
|
||
print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
|
||
|
||
try:
|
||
proc.terminate()
|
||
time.sleep(2)
|
||
except Exception:
|
||
proc.kill()
|
||
|
||
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
||
pass
|
||
|
||
except Exception as e:
|
||
print(f"Error in psutil vllm cleanup: {e}")
|
||
|
||
def __enter__(self):
|
||
"""Context manager entry."""
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
"""Context manager exit - cleanup resources."""
|
||
if hasattr(self.omni, "close"):
|
||
self.omni.close()
|
||
self._cleanup_process()
|
||
_run_pre_test_cleanup(enable_force=True)
|
||
_run_post_test_cleanup(enable_force=True)
|
||
cleanup_dist_env_and_memory()
|
||
|
||
|
||
@pytest.fixture(scope="module")
|
||
def omni_runner(request, model_prefix):
|
||
with _omni_server_lock:
|
||
model, stage_config_path = request.param
|
||
model = model_prefix + model
|
||
with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner:
|
||
print("OmniRunner started successfully")
|
||
yield runner
|
||
print("OmniRunner stopping...")
|
||
|
||
print("OmniRunner stopped")
|
||
|
||
|
||
class OmniRunnerHandler:
|
||
def __init__(self, omni_runner):
|
||
self.runner = omni_runner
|
||
|
||
def _process_output(self, outputs: list[Any]) -> OmniResponse:
|
||
result = OmniResponse()
|
||
try:
|
||
text_content = None
|
||
audio_content = None
|
||
for stage_output in outputs:
|
||
if getattr(stage_output, "final_output_type", None) == "text":
|
||
text_content = stage_output.request_output[0].outputs[0].text
|
||
if getattr(stage_output, "final_output_type", None) == "audio":
|
||
audio_content = stage_output.request_output[0].outputs[0].multimodal_output["audio"]
|
||
|
||
result.audio_content = audio_content
|
||
result.text_content = text_content
|
||
result.success = True
|
||
|
||
except Exception as e:
|
||
result.error_message = f"Output processing error: {str(e)}"
|
||
result.success = False
|
||
print(f"Error: {result.error_message}")
|
||
|
||
return result
|
||
|
||
def send_request(self, request_config: dict[str, Any] | None = None) -> OmniResponse:
|
||
if request_config is None:
|
||
request_config = {}
|
||
prompts = request_config.get("prompts")
|
||
videos = request_config.get("videos")
|
||
images = request_config.get("images")
|
||
audios = request_config.get("audios")
|
||
modalities = request_config.get("modalities", ["text", "audio"])
|
||
outputs = self.runner.generate_multimodal(
|
||
prompts=prompts, videos=videos, images=images, audios=audios, modalities=modalities
|
||
)
|
||
response = self._process_output(outputs)
|
||
assert_omni_response(response, request_config, run_level="L2")
|
||
return response
|
||
|
||
|
||
@pytest.fixture
|
||
def omni_runner_handler(omni_runner):
|
||
return OmniRunnerHandler(omni_runner)
|