vllm-project--vllm-omni
wehub-resource-sync
2944 行
122 KiB
Python
2944 行
122 KiB
Python
"""Server/client/runner runtime primitives for tests."""
|
|
|
|
import asyncio
|
|
import base64
|
|
import concurrent.futures
|
|
import copy
|
|
import errno
|
|
import gc
|
|
import io
|
|
import json
|
|
import os
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from collections.abc import Generator
|
|
from dataclasses import asdict, dataclass
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
from typing import Any, NamedTuple, cast
|
|
from urllib.parse import quote
|
|
|
|
import psutil
|
|
import requests
|
|
import soundfile as sf
|
|
import torch
|
|
import yaml
|
|
from openai import APIError, OpenAI, omit
|
|
from PIL import Image
|
|
from vllm import TextPrompt, envs
|
|
from vllm.distributed.parallel_state import (
|
|
destroy_distributed_environment,
|
|
destroy_model_parallel,
|
|
)
|
|
from vllm.logger import init_logger
|
|
|
|
from tests.helpers.assertions import (
|
|
assert_audio_speech_response,
|
|
assert_diffusion_response,
|
|
assert_http_error,
|
|
assert_omni_response,
|
|
)
|
|
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
|
|
from tests.helpers.media import (
|
|
_merge_base64_audio_to_segment,
|
|
decode_b64_image,
|
|
)
|
|
from tests.model_tests.diffusion.utils import resolve_tiny_model_path
|
|
from vllm_omni.config.stage_config import resolve_deploy_yaml
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
|
|
# Reset environment variable cache
|
|
envs.disable_envs_cache()
|
|
|
|
# Reset rocm_aiter_ops class variables to match current os.environ.
|
|
# These are class-level attributes that persist across tests and are
|
|
# NOT restored by monkeypatch (which only restores os.environ).
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
if current_omni_platform.is_rocm():
|
|
from vllm._aiter_ops import rocm_aiter_ops
|
|
|
|
rocm_aiter_ops.refresh_env_variables()
|
|
|
|
# Ensure all objects are not frozen before cleanup
|
|
gc.unfreeze()
|
|
|
|
destroy_model_parallel()
|
|
destroy_distributed_environment()
|
|
if shutdown_ray:
|
|
import ray # Lazy import Ray
|
|
|
|
ray.shutdown()
|
|
gc.collect()
|
|
|
|
if not current_omni_platform.is_cpu():
|
|
current_omni_platform.empty_cache()
|
|
try:
|
|
torch._C._host_emptyCache()
|
|
except AttributeError:
|
|
logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5")
|
|
|
|
|
|
def _parse_response_json(r: requests.Response) -> dict[str, Any] | list[Any] | None:
|
|
try:
|
|
data = r.json()
|
|
if isinstance(data, (dict, list)):
|
|
return data
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[str, Any]] | None:
|
|
"""If ``extra_body`` has list ``height``/``width``, return one config per index (scalar h/w, ``num_outputs_per_prompt=1``)."""
|
|
eb = cfg.get("extra_body")
|
|
if not eb:
|
|
return None
|
|
h, w = eb.get("height"), eb.get("width")
|
|
if (isinstance(h, (list, tuple)) or isinstance(w, (list, tuple))) and not (
|
|
isinstance(h, (list, tuple)) and isinstance(w, (list, tuple))
|
|
):
|
|
raise ValueError("extra_body height and width must both be lists or both be scalars")
|
|
if not (isinstance(h, (list, tuple)) and isinstance(w, (list, tuple))):
|
|
return None
|
|
if len(h) != len(w):
|
|
raise ValueError(f"height and width lists must have equal length; got {len(h)=} {len(w)=}")
|
|
n = len(h)
|
|
n_out = eb.get("num_outputs_per_prompt")
|
|
if n_out is not None:
|
|
n_out = int(n_out)
|
|
if n_out != n:
|
|
raise ValueError(
|
|
"When height/width are lists, num_outputs_per_prompt must equal their length; "
|
|
f"got num_outputs_per_prompt={n_out}, len(lists)={n}"
|
|
)
|
|
splits: list[dict[str, Any]] = []
|
|
for i in range(n):
|
|
sub = copy.deepcopy(cfg)
|
|
sub_eb = dict(sub.get("extra_body") or {})
|
|
sub_eb["height"] = int(h[i])
|
|
sub_eb["width"] = int(w[i])
|
|
sub_eb["num_outputs_per_prompt"] = 1
|
|
sub["extra_body"] = sub_eb
|
|
splits.append(sub)
|
|
return splits
|
|
|
|
|
|
PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
|
|
PromptImageInput = list[Any] | Any | None
|
|
PromptVideoInput = list[Any] | Any | None
|
|
|
|
|
|
def get_open_port(host: str = "127.0.0.1", *, max_attempts: int = 128) -> int:
|
|
"""Return a local TCP port that is suitable for binding a new listener.
|
|
|
|
A single ``bind(host, 0)`` / close cycle leaves a race where another process can
|
|
take the same port number before PyTorch/vLLM bind it, yielding
|
|
``EADDRINUSE`` / ``DistNetworkError``. We therefore:
|
|
|
|
#. Allocate an ephemeral port on *host*.
|
|
#. Immediately attempt ``bind(host, port)`` again. If that fails with
|
|
``errno.EADDRINUSE``, retry from step 1.
|
|
|
|
Raises ``RuntimeError`` if no free port is found after *max_attempts* (e.g. port
|
|
exhaustion under heavy parallel tests).
|
|
"""
|
|
last_exc: OSError | None = None
|
|
for _ in range(max_attempts):
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.bind((host, 0))
|
|
port = int(s.getsockname()[1])
|
|
try:
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
|
|
probe.bind((host, port))
|
|
except OSError as exc:
|
|
last_exc = exc
|
|
if exc.errno == errno.EADDRINUSE:
|
|
continue
|
|
raise
|
|
return port
|
|
raise RuntimeError(
|
|
f"Could not obtain a free TCP port on {host!r} after {max_attempts} attempts (last error: {last_exc!r})"
|
|
) from last_exc
|
|
|
|
|
|
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 _omni_subprocess_cwd() -> str:
|
|
"""Repo root for ``python -m vllm_omni...`` (legacy conftest lived under ``tests/``; helpers under ``tests/helpers/``)."""
|
|
return os.path.normpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
|
|
|
|
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: fixture supplies default (600 s)
|
|
|
|
|
|
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()
|
|
run_post_test_cleanup()
|
|
cleanup_dist_env_and_memory()
|
|
self.model = model
|
|
args = list(serve_args)
|
|
self.serve_args = args
|
|
self.log_stats = "--disable-log-stats" not in args and "--log-stats" in args
|
|
self.env_dict = env_dict
|
|
self.use_omni = use_omni
|
|
self.proc: subprocess.Popen | None = None
|
|
self.host = "127.0.0.1"
|
|
self.port = get_open_port() if port is None else port
|
|
|
|
def _start_server(self) -> None:
|
|
env = os.environ.copy()
|
|
env.setdefault("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)}")
|
|
startup_t0 = time.perf_counter()
|
|
self.proc = subprocess.Popen(
|
|
cmd,
|
|
env=env,
|
|
cwd=_omni_subprocess_cwd(),
|
|
)
|
|
|
|
max_wait = 1200
|
|
start_time = time.time()
|
|
while time.time() - start_time < max_wait:
|
|
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)
|
|
if sock.connect_ex((self.host, self.port)) == 0:
|
|
startup_s = time.perf_counter() - startup_t0
|
|
if self.log_stats:
|
|
print(
|
|
f"Server ready on {self.host}:{self.port} (OmniServer startup took {startup_s:.3f}s)",
|
|
flush=True,
|
|
)
|
|
return
|
|
time.sleep(2)
|
|
raise RuntimeError(f"Server failed to start within {max_wait} seconds")
|
|
|
|
@staticmethod
|
|
def _reap_zombie(proc: "psutil.Process") -> bool:
|
|
"""Reap a zombie child process via ``os.waitpid``.
|
|
|
|
``psutil.Process.wait()`` uses ``pidfd_open`` + ``poll()``, which
|
|
never fires for a zombie (the zombie is already dead and will not
|
|
change state). Since the test process is the parent, we can reap
|
|
the zombie directly with ``os.waitpid(pid, os.WNOHANG)`` and
|
|
retrieve its exit code.
|
|
|
|
Returns True if the process was a zombie and was reaped.
|
|
"""
|
|
if proc.status() != psutil.STATUS_ZOMBIE:
|
|
return False
|
|
try:
|
|
os.waitpid(proc.pid, os.WNOHANG)
|
|
except ChildProcessError:
|
|
pass
|
|
return True
|
|
|
|
@staticmethod
|
|
def _wait_or_reap(proc: "psutil.Process", timeout: int) -> None:
|
|
"""Wait for *proc* to exit, handling zombie state transparently.
|
|
|
|
When the process is already a zombie (e.g. orchestrator thread
|
|
hung during shutdown and the main process exited), ``pidfd_open``
|
|
+ ``poll()`` inside ``psutil`` will never see a state change.
|
|
Fall back to ``os.waitpid`` to reap the zombie.
|
|
"""
|
|
if OmniServer._reap_zombie(proc):
|
|
return
|
|
try:
|
|
proc.wait(timeout=timeout)
|
|
except psutil.TimeoutExpired:
|
|
OmniServer._reap_zombie(proc)
|
|
|
|
def _kill_process_tree(self, pid):
|
|
"""Kill the process tree rooted at *pid*.
|
|
|
|
Terminate the parent **first** so the OmniServer can gracefully shut
|
|
down its stage-engine children through the orchestrator. This avoids
|
|
the ``subprocess died unexpectedly`` ERROR that the APIServer monitor
|
|
thread logs when children are killed before the parent, which in turn
|
|
can cause CI watchdogs to false-trigger on the upstream ``Shutdown
|
|
initiated`` message.
|
|
|
|
When the parent does not exit within the grace period (e.g. CPU-
|
|
offloaded workers stuck in CUDA D-state), the method falls back to
|
|
killing children first so the parent can be reaped cleanly.
|
|
"""
|
|
try:
|
|
parent = psutil.Process(pid)
|
|
children = parent.children(recursive=True)
|
|
|
|
# 1. Terminate the parent first — let it run its graceful
|
|
# shutdown cascade (orchestrator → stage pools → engine cores).
|
|
try:
|
|
parent.terminate()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
|
|
# 2. Give the parent time to shut down its children cleanly.
|
|
parent_exited = False
|
|
try:
|
|
parent.wait(timeout=15)
|
|
parent_exited = True
|
|
except psutil.NoSuchProcess:
|
|
parent_exited = True
|
|
except psutil.TimeoutExpired:
|
|
parent_exited = OmniServer._reap_zombie(parent)
|
|
|
|
if not parent_exited:
|
|
# Parent is stuck — children (e.g. CPU-offloaded CFG workers)
|
|
# are likely in uninterruptible sleep. Kill children first
|
|
# so the parent can be reaped without lingering as a zombie.
|
|
for child in children:
|
|
try:
|
|
child.kill()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
psutil.wait_procs(children, timeout=5)
|
|
try:
|
|
parent.kill()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
OmniServer._wait_or_reap(parent, timeout=5)
|
|
else:
|
|
# Parent exited cleanly — clean up any remaining children.
|
|
for child in children:
|
|
try:
|
|
if child.is_running():
|
|
child.terminate()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
|
|
gone, still_alive = psutil.wait_procs(children, timeout=10)
|
|
|
|
for child in still_alive:
|
|
try:
|
|
child.kill()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
|
|
try:
|
|
if parent.is_running() and not OmniServer._reap_zombie(parent):
|
|
parent.kill()
|
|
parent.wait(timeout=10)
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
|
|
# 3. Final sweep — ``kill -9`` anything that escaped.
|
|
time.sleep(1)
|
|
alive_processes: list[int] = []
|
|
for child in children:
|
|
try:
|
|
if child.is_running():
|
|
alive_processes.append(child.pid)
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
# Only count the parent as alive if it is NOT a zombie
|
|
# (zombies are already dead — just waiting to be reaped).
|
|
try:
|
|
if parent.is_running() and parent.status() != psutil.STATUS_ZOMBIE:
|
|
alive_processes.append(parent.pid)
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
|
|
if alive_processes:
|
|
print(f"Warning: Processes still alive: {alive_processes}")
|
|
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()
|
|
run_post_test_cleanup()
|
|
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_replica_counts = self._load_stage_replica_counts(resolved_cfg)
|
|
self.stage_procs: dict[tuple[int, 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
|
|
|
|
@staticmethod
|
|
def _load_stage_replica_counts(resolved_config: dict) -> dict[int, int]:
|
|
replica_counts: dict[int, int] = {}
|
|
for stage in OmniServerStageCli._stage_entries(resolved_config):
|
|
stage_id = stage.get("stage_id")
|
|
if stage_id is None:
|
|
continue
|
|
replica_counts[int(stage_id)] = max(
|
|
1,
|
|
int(stage.get("num_replicas") or stage.get("runtime", {}).get("num_replicas", 1)),
|
|
)
|
|
return replica_counts
|
|
|
|
@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 _devices_for_replica(self, stage_id: int, devices: str, replica_id: int) -> str:
|
|
replica_count = self.stage_replica_counts.get(stage_id, 1)
|
|
if replica_count == 1:
|
|
return devices
|
|
|
|
device_list = self._parse_device_list(devices)
|
|
if len(device_list) % replica_count != 0:
|
|
raise ValueError(
|
|
f"Stage {stage_id} has {len(device_list)} device(s) for {replica_count} replica(s); "
|
|
"device count must be divisible by replica count"
|
|
)
|
|
devices_per_replica = len(device_list) // replica_count
|
|
start = replica_id * devices_per_replica
|
|
return ",".join(device_list[start : start + devices_per_replica])
|
|
|
|
def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str, replica_id: int = 0) -> None:
|
|
replica_devices = self._devices_for_replica(stage_id, devices, replica_id)
|
|
mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, replica_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, replica_id: int = 0) -> 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),
|
|
"--replica-id",
|
|
str(replica_id),
|
|
]
|
|
|
|
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, replica_id: int = 0) -> None:
|
|
env = os.environ.copy()
|
|
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
|
if self.env_dict is not None:
|
|
env.update(self.env_dict)
|
|
|
|
cmd = self._build_stage_cmd(stage_id, headless=headless, replica_id=replica_id)
|
|
print(f"Launching OmniServerStageCli stage {stage_id} replica {replica_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}_replica_{replica_id}_{self.master_port}.log"
|
|
self._stage_log_paths = getattr(self, "_stage_log_paths", {})
|
|
stage_key = (stage_id, replica_id)
|
|
self._stage_log_paths[stage_key] = 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_key] = 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_key] = proc
|
|
if stage_id == 0 and replica_id == 0:
|
|
self.proc = proc
|
|
|
|
def _ensure_stage_processes_alive(self) -> None:
|
|
for (stage_id, replica_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, replica_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 {} replica {} log ({}) ===\n{}".format(
|
|
stage_id, replica_id, log_path, "".join(lines[-60:]) or "<empty>"
|
|
)
|
|
except Exception as exc: # pragma: no cover - diagnostic only
|
|
tail = f"\n<failed to read stage log {log_path}: {exc}>"
|
|
raise RuntimeError(
|
|
f"Stage {stage_id} replica {replica_id} exited with code {ret} before API server became ready.{tail}"
|
|
)
|
|
|
|
def _start_server(self) -> None:
|
|
startup_t0 = time.perf_counter()
|
|
ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]]
|
|
|
|
self._launch_stage(0, headless=False, replica_id=0)
|
|
time.sleep(2)
|
|
self._ensure_stage_processes_alive()
|
|
|
|
for stage_id in ordered_stage_ids[1:]:
|
|
for replica_id in range(self.stage_replica_counts.get(stage_id, 1)):
|
|
self._launch_stage(stage_id, headless=True, replica_id=replica_id)
|
|
|
|
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:
|
|
startup_s = time.perf_counter() - startup_t0
|
|
if self.log_stats:
|
|
print(
|
|
f"OmniServerStageCli ready on {self.host}:{self.port} "
|
|
f"(stage-CLI startup took {startup_s:.3f}s)",
|
|
flush=True,
|
|
)
|
|
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, replica_id in sorted(log_paths):
|
|
log_path = log_paths[(stage_id, replica_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} replica {replica_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} replica {replica_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} replica {replica_id} log TAIL ({log_path}) ===", flush=True)
|
|
print("".join(tail_chunk).rstrip("\n"), flush=True)
|
|
print(f"=== end stage {stage_id} replica {replica_id} log ===\n", flush=True)
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
self._dump_stage_logs_for_debug()
|
|
for stage_key in sorted(self.stage_procs, reverse=True):
|
|
proc = self.stage_procs[stage_key]
|
|
if proc.poll() is None:
|
|
self._kill_process_tree(proc.pid)
|
|
run_pre_test_cleanup()
|
|
run_post_test_cleanup()
|
|
cleanup_dist_env_and_memory()
|
|
|
|
|
|
@dataclass
|
|
class OmniResponse:
|
|
"""Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""
|
|
|
|
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
|
|
#: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before the
|
|
#: OpenAI client call through response parsing and local post-process (e.g. audio decode).
|
|
e2e_latency: float | None = None
|
|
success: bool = False
|
|
prompt_tokens: int | None = None
|
|
cached_tokens: int | None = None
|
|
logprobs: list | None = None
|
|
#: HTTP status + error text for the error-handling path (e.g. validator
|
|
#: rejections); populated when the OpenAI client raises an APIError.
|
|
status_code: int | None = None
|
|
error_message: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class DiffusionResponse:
|
|
"""Decoded diffusion output from chat completions or offline runner (not raw ``requests``)."""
|
|
|
|
text_content: str | None = None
|
|
images: list[Image.Image] | None = None
|
|
audios: list[Any] | None = None
|
|
videos: list[Any] | None = None
|
|
#: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before
|
|
#: ``chat.completions.create`` through local image / audio decode.
|
|
e2e_latency: float | None = None
|
|
success: bool = False
|
|
|
|
|
|
@dataclass
|
|
class HttpResponse:
|
|
"""Normalized view of a ``requests`` response from :class:`OpenAIClientHandler` HTTP helpers."""
|
|
|
|
status_code: int
|
|
success: bool
|
|
error_message: str | None = None
|
|
json_body: dict[str, Any] | list[Any] | None = None
|
|
|
|
|
|
@dataclass
|
|
class WebSocketJsonResponse:
|
|
"""First JSON object delivered as a text WebSocket frame (streaming endpoints)."""
|
|
|
|
json_body: dict[str, Any]
|
|
|
|
|
|
def _merge_http_expectation_kwargs(
|
|
base: dict[str, Any] | None,
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> dict[str, Any]:
|
|
cfg = dict(base or {})
|
|
if err_code is not None:
|
|
cfg["err_code"] = err_code
|
|
if err_message is not None:
|
|
cfg["err_message"] = err_message
|
|
return cfg
|
|
|
|
|
|
def _merge_ws_expectation_kwargs(
|
|
base: dict[str, Any] | None,
|
|
*,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
ws_json_type: str | None = None,
|
|
ws_error_code: str | None = None,
|
|
) -> dict[str, Any]:
|
|
cfg = dict(base or {})
|
|
if err_message is not None:
|
|
cfg["err_message"] = err_message
|
|
if ws_json_type is not None:
|
|
cfg["ws_json_type"] = ws_json_type
|
|
if ws_error_code is not None:
|
|
cfg["ws_error_code"] = ws_error_code
|
|
return cfg
|
|
|
|
|
|
def _run_ws_expectations_from_request_config(cfg: dict[str, Any], resp: WebSocketJsonResponse) -> None:
|
|
jb = resp.json_body
|
|
want_type = cfg.get("ws_json_type")
|
|
if want_type is not None:
|
|
assert jb.get("type") == want_type, (jb, want_type)
|
|
want_code = cfg.get("ws_error_code")
|
|
if want_code is not None:
|
|
assert jb.get("code") == want_code, (jb, want_code)
|
|
err_message = cfg.get("err_message")
|
|
if err_message is not None:
|
|
assert_http_error(resp, err_message=err_message, websocket_json_message=True)
|
|
|
|
|
|
def _merge_diffusion_responses(parts: list[DiffusionResponse]) -> DiffusionResponse:
|
|
"""Concatenate images in order; ``e2e_latency`` is wall-clock of the batch (set by caller) or max of parts."""
|
|
merged = DiffusionResponse()
|
|
merged.success = all(p.success for p in parts) and len(parts) > 0
|
|
imgs: list[Image.Image] = []
|
|
for p in parts:
|
|
if p.images:
|
|
imgs.extend(p.images)
|
|
merged.images = imgs if imgs else None
|
|
latencies = [p.e2e_latency for p in parts if p.e2e_latency is not None]
|
|
merged.e2e_latency = max(latencies) if latencies else None
|
|
return merged
|
|
|
|
|
|
class OpenAIClientHandler:
|
|
def __init__(
|
|
self,
|
|
host: str = "127.0.0.1",
|
|
port: int = None,
|
|
api_key: str = "EMPTY",
|
|
run_level: str = None,
|
|
*,
|
|
log_stats: bool = True,
|
|
):
|
|
if port is None:
|
|
port = get_open_port()
|
|
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
|
|
self.log_stats = log_stats
|
|
|
|
def _print_client_stat(self, message: str) -> None:
|
|
if self.log_stats:
|
|
print(message, flush=True)
|
|
|
|
def _process_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse:
|
|
"""Wall clock from *before* ``chat.completions.create`` through stream drain + local decode."""
|
|
result = OmniResponse()
|
|
try:
|
|
text_content = ""
|
|
audio_data = []
|
|
for chunk in chat_completion:
|
|
for choice in chunk.choices:
|
|
content = getattr(getattr(choice, "delta", None), "content", None)
|
|
modality = getattr(chunk, "modality", None)
|
|
if modality == "audio" and content:
|
|
audio_data.append(content)
|
|
elif modality == "text" and content:
|
|
text_content += content
|
|
# Usage is yielded after the last token
|
|
if chunk.usage:
|
|
result.prompt_tokens = chunk.usage.prompt_tokens
|
|
if details := getattr(chunk.usage, "prompt_tokens_details", None):
|
|
result.cached_tokens = details.cached_tokens
|
|
|
|
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()
|
|
result.text_content = text_content
|
|
result.audio_data = audio_data
|
|
result.e2e_latency = time.perf_counter() - wall_start
|
|
result.success = True
|
|
except Exception as e:
|
|
msg = f"Stream processing error: {str(e)}"
|
|
print(f"Error: {msg}")
|
|
return result
|
|
|
|
def _process_non_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse:
|
|
"""Wall clock from *before* ``chat.completions.create`` through response parse + local decode."""
|
|
result = OmniResponse()
|
|
try:
|
|
audio_data = None
|
|
text_content = None
|
|
for choice in chat_completion.choices:
|
|
if hasattr(choice.message, "audio") and choice.message.audio is not None:
|
|
audio_data = choice.message.audio.data
|
|
if hasattr(choice.message, "content") and choice.message.content is not None:
|
|
text_content = choice.message.content
|
|
# Extract cached & prompt token counts for prefix caching tests
|
|
usage = getattr(chat_completion, "usage", None)
|
|
if usage:
|
|
result.prompt_tokens = usage.prompt_tokens
|
|
if details := getattr(usage, "prompt_tokens_details", None):
|
|
result.cached_tokens = details.cached_tokens
|
|
if audio_data:
|
|
result.audio_bytes = base64.b64decode(audio_data)
|
|
result.text_content = text_content
|
|
result.e2e_latency = time.perf_counter() - wall_start
|
|
if chat_completion.choices and chat_completion.choices[0].logprobs is not None:
|
|
result.logprobs = chat_completion.choices[0].logprobs.content
|
|
result.success = True
|
|
except Exception as e:
|
|
msg = f"Non-stream processing error: {str(e)}"
|
|
print(f"Error: {msg}")
|
|
return result
|
|
|
|
def _process_diffusion_response(self, chat_completion, *, wall_start: float) -> DiffusionResponse:
|
|
"""Wall clock from *before* ``chat.completions.create`` through image decode."""
|
|
result = DiffusionResponse()
|
|
try:
|
|
images = []
|
|
audios = []
|
|
for choice in chat_completion.choices:
|
|
content = getattr(choice.message, "content", None)
|
|
if isinstance(content, list):
|
|
for item in content:
|
|
image_url = None
|
|
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]
|
|
images.append(decode_b64_image(b64_data))
|
|
|
|
# OpenAI audio responses (e.g. AudioX text-to-audio) populate `message.audio`.
|
|
audio_obj = getattr(choice.message, "audio", None)
|
|
audio_b64 = getattr(audio_obj, "data", None) if audio_obj is not None else None
|
|
if audio_b64:
|
|
audios.append(
|
|
{
|
|
"wav_bytes": base64.b64decode(audio_b64),
|
|
"id": getattr(audio_obj, "id", None),
|
|
"expires_at": getattr(audio_obj, "expires_at", None),
|
|
}
|
|
)
|
|
result.images = images if images else None
|
|
result.audios = audios if audios else None
|
|
result.e2e_latency = time.perf_counter() - wall_start
|
|
result.success = True
|
|
except Exception as e:
|
|
msg = f"Diffusion response processing error: {str(e)}"
|
|
print(f"Error: {msg}")
|
|
return result
|
|
|
|
def _http_response_from_requests(self, r: requests.Response) -> HttpResponse:
|
|
payload = _parse_response_json(r)
|
|
ok = 200 <= r.status_code < 300
|
|
return HttpResponse(
|
|
status_code=r.status_code,
|
|
success=ok,
|
|
error_message=None if ok else (r.text[:8000] if r.text else None),
|
|
json_body=payload,
|
|
)
|
|
|
|
def send_health_http_request(
|
|
self,
|
|
request_config: dict[str, Any] | None = None,
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""GET ``/health`` (raw ``requests``).
|
|
|
|
``request_config``: optional ``timeout`` plus optional ``err_code`` / ``err_message`` for
|
|
:func:`~tests.helpers.assertions.assert_http_error` (also as keyword-only args).
|
|
"""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = requests.get(self._build_url("/health"), timeout=float(cfg.get("timeout", 120.0)))
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_models_http_request(
|
|
self,
|
|
request_config: dict[str, Any] | None = None,
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""GET ``/v1/models``. Optional ``timeout`` and HTTP assertions (see :func:`~tests.helpers.assertions.assert_http_error`)."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = requests.get(
|
|
self._build_url("/v1/models"),
|
|
headers={"Accept": "application/json"},
|
|
timeout=float(cfg.get("timeout", 120.0)),
|
|
)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_chat_completions_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/chat/completions`` with ``json`` or ``raw_body`` (malformed-body / contract tests)."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_json_endpoint("/v1/chat/completions", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_completions_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/completions`` with ``json`` or ``raw_body``."""
|
|
# TODO (Alex): A lot of these helpers should be consolidated as they differ only by endpoint
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_json_endpoint("/v1/completions", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_omni_sleep_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/omni/sleep`` — ``json`` or ``raw_body``, ``timeout``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_json_endpoint("/v1/omni/sleep", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_omni_wakeup_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/omni/wakeup``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_json_endpoint("/v1/omni/wakeup", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_audio_voices_list_http_request(
|
|
self,
|
|
request_config: dict[str, Any] | None = None,
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""GET ``/v1/audio/voices``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = requests.get(
|
|
self._build_url("/v1/audio/voices"),
|
|
headers={"Accept": "application/json"},
|
|
timeout=float(cfg.get("timeout", 120.0)),
|
|
)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_audio_voices_create_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/audio/voices`` (multipart): ``data`` / ``files`` / ``timeout``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_form_endpoint("/v1/audio/voices", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_audio_voices_delete_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""DELETE ``/v1/audio/voices/{name}`` — requires ``name``, optional ``timeout``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
name = cfg["name"]
|
|
timeout = float(cfg.get("timeout", 120.0))
|
|
path = f"/v1/audio/voices/{quote(str(name), safe='')}"
|
|
r = requests.delete(
|
|
self._build_url(path),
|
|
headers={"Accept": "application/json"},
|
|
timeout=timeout,
|
|
)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_audio_speech_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/audio/speech`` with ``json`` or ``raw_body``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_json_endpoint("/v1/audio/speech", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_audio_speech_batch_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/audio/speech/batch``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_json_endpoint("/v1/audio/speech/batch", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_audio_generate_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/audio/generate``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_json_endpoint("/v1/audio/generate", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_images_generations_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/images/generations`` — ``json`` or ``raw_body``, ``timeout``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_json_endpoint("/v1/images/generations", cfg, default_timeout=300.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_images_edits_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/images/edits`` — ``data`` / ``files`` / ``timeout``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_form_endpoint("/v1/images/edits", cfg, default_timeout=300.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_videos_create_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/videos`` (async job) — multipart ``data`` / ``files``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_form_endpoint("/v1/videos", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_videos_sync_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""POST ``/v1/videos/sync``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = self._post_form_endpoint("/v1/videos/sync", cfg, default_timeout=120.0)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_videos_list_http_request(
|
|
self,
|
|
request_config: dict[str, Any] | None = None,
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""GET ``/v1/videos`` — optional ``params``, ``timeout``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
r = requests.get(
|
|
self._build_url("/v1/videos"),
|
|
params=cfg.get("params"),
|
|
headers={"Accept": "application/json"},
|
|
timeout=float(cfg.get("timeout", 120.0)),
|
|
)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_video_retrieve_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""GET ``/v1/videos/{video_id}``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
video_id = cfg["video_id"]
|
|
timeout = float(cfg.get("timeout", 120.0))
|
|
r = requests.get(
|
|
self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
|
|
headers={"Accept": "application/json"},
|
|
timeout=timeout,
|
|
)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_video_delete_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""DELETE ``/v1/videos/{video_id}``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
video_id = cfg["video_id"]
|
|
timeout = float(cfg.get("timeout", 120.0))
|
|
r = requests.delete(
|
|
self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
|
|
headers={"Accept": "application/json"},
|
|
timeout=timeout,
|
|
)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def send_video_content_http_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_code: int | tuple[int, ...] | list[int] | None = None,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
) -> list[HttpResponse]:
|
|
"""GET ``/v1/videos/{video_id}/content``."""
|
|
cfg = _merge_http_expectation_kwargs(
|
|
request_config,
|
|
err_code=err_code,
|
|
err_message=err_message,
|
|
)
|
|
video_id = cfg["video_id"]
|
|
timeout = float(cfg.get("timeout", 120.0))
|
|
r = requests.get(
|
|
self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}/content"),
|
|
timeout=timeout,
|
|
)
|
|
resp = self._http_response_from_requests(r)
|
|
assert_http_error(
|
|
resp,
|
|
err_code=cfg.get("err_code"),
|
|
err_message=cfg.get("err_message"),
|
|
)
|
|
return [resp]
|
|
|
|
def _build_ws_url(self, path: str) -> str:
|
|
"""Turn HTTP ``base_url`` into ``ws`` / ``wss`` for WebSocket helpers."""
|
|
base = self.base_url.rstrip("/")
|
|
suffix = "/" + path.lstrip("/")
|
|
if base.startswith("http://"):
|
|
return "ws://" + base.removeprefix("http://") + suffix
|
|
if base.startswith("https://"):
|
|
return "wss://" + base.removeprefix("https://") + suffix
|
|
raise ValueError(f"Unsupported base_url for WebSocket: {base!r}")
|
|
|
|
def _send_websocket_first_json_request(
|
|
self,
|
|
path: str,
|
|
cfg: dict[str, Any],
|
|
) -> list[WebSocketJsonResponse]:
|
|
"""Connect, optionally send text frames, return first JSON text frame as :class:`WebSocketJsonResponse`.
|
|
|
|
``request_config`` keys:
|
|
|
|
- ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server
|
|
speaks first, e.g. ``/v1/realtime`` rejection path).
|
|
- ``ws_skip_types``: optional event ``type`` strings to ignore while waiting for the first matching frame
|
|
(e.g. ``["session.created"]`` on ``/v1/realtime``).
|
|
- ``timeout``: seconds to wait for the first inbound text frame (default ``120``).
|
|
- ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present.
|
|
"""
|
|
send_frames_raw = cfg.get("send_frames")
|
|
if send_frames_raw is None:
|
|
frames: list[str] = []
|
|
elif isinstance(send_frames_raw, str):
|
|
frames = [send_frames_raw]
|
|
else:
|
|
frames = list(send_frames_raw)
|
|
|
|
timeout = float(cfg.get("timeout", 120.0))
|
|
uri = self._build_ws_url(path)
|
|
skip_types = set(cfg.get("ws_skip_types") or [])
|
|
|
|
connect_kw: dict[str, Any] = {}
|
|
if "ws_max_size" in cfg:
|
|
connect_kw["max_size"] = cfg["ws_max_size"]
|
|
|
|
async def _recv_first_json_object() -> WebSocketJsonResponse:
|
|
import websockets
|
|
|
|
async with websockets.connect(uri, **connect_kw) as ws:
|
|
for frame in frames:
|
|
await ws.send(frame)
|
|
while True:
|
|
raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
|
|
if not isinstance(raw, str):
|
|
raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}")
|
|
try:
|
|
data = json.loads(raw)
|
|
except json.JSONDecodeError as exc:
|
|
raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc
|
|
if not isinstance(data, dict):
|
|
raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}")
|
|
if skip_types and data.get("type") in skip_types:
|
|
continue
|
|
return WebSocketJsonResponse(json_body=data)
|
|
|
|
resp = asyncio.run(_recv_first_json_object())
|
|
_run_ws_expectations_from_request_config(cfg, resp)
|
|
return [resp]
|
|
|
|
def send_audio_speech_stream_ws_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
ws_json_type: str | None = None,
|
|
ws_error_code: str | None = None,
|
|
) -> list[WebSocketJsonResponse]:
|
|
"""WebSocket ``/v1/audio/speech/stream`` — send ``send_frames`` then read first JSON text frame."""
|
|
cfg = _merge_ws_expectation_kwargs(
|
|
request_config,
|
|
err_message=err_message,
|
|
ws_json_type=ws_json_type,
|
|
ws_error_code=ws_error_code,
|
|
)
|
|
return self._send_websocket_first_json_request("/v1/audio/speech/stream", cfg)
|
|
|
|
def send_video_chat_stream_ws_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
ws_json_type: str | None = None,
|
|
ws_error_code: str | None = None,
|
|
) -> list[WebSocketJsonResponse]:
|
|
"""WebSocket ``/v1/video/chat/stream`` — send ``send_frames`` then read first JSON text frame."""
|
|
cfg = _merge_ws_expectation_kwargs(
|
|
request_config,
|
|
err_message=err_message,
|
|
ws_json_type=ws_json_type,
|
|
ws_error_code=ws_error_code,
|
|
)
|
|
return self._send_websocket_first_json_request("/v1/video/chat/stream", cfg)
|
|
|
|
def send_realtime_ws_request(
|
|
self,
|
|
request_config: dict[str, Any] | None = None,
|
|
*,
|
|
err_message: str | tuple[str, ...] | list[str] | None = None,
|
|
ws_json_type: str | None = None,
|
|
ws_error_code: str | None = None,
|
|
) -> list[WebSocketJsonResponse]:
|
|
"""WebSocket ``/v1/realtime`` — optional outbound frames, then first JSON text frame (often server-initiated)."""
|
|
cfg = _merge_ws_expectation_kwargs(
|
|
request_config,
|
|
err_message=err_message,
|
|
ws_json_type=ws_json_type,
|
|
ws_error_code=ws_error_code,
|
|
)
|
|
return self._send_websocket_first_json_request("/v1/realtime", cfg)
|
|
|
|
def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
|
|
"""Chat completions via the OpenAI Python SDK (not raw HTTP)."""
|
|
responses: list[OmniResponse] = []
|
|
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
|
|
if "sampling_params_list" in request_config:
|
|
extra_body["sampling_params_list"] = request_config["sampling_params_list"]
|
|
if request_config.get("extra_body"):
|
|
extra_body.update(request_config["extra_body"])
|
|
|
|
create_kwargs: dict[str, Any] = {
|
|
"model": request_config.get("model"),
|
|
"messages": request_config.get("messages"),
|
|
"stream": stream,
|
|
"modalities": modalities,
|
|
}
|
|
if "logprobs" in request_config:
|
|
create_kwargs["logprobs"] = request_config["logprobs"]
|
|
if "top_logprobs" in request_config:
|
|
create_kwargs["top_logprobs"] = request_config["top_logprobs"]
|
|
if "stream_options" in request_config:
|
|
create_kwargs["stream_options"] = request_config["stream_options"]
|
|
if extra_body:
|
|
create_kwargs["extra_body"] = extra_body
|
|
|
|
if request_num == 1:
|
|
wall_start = time.perf_counter()
|
|
chat_completion = self.client.chat.completions.create(**create_kwargs)
|
|
resp = (
|
|
self._process_stream_omni_response(chat_completion, wall_start=wall_start)
|
|
if stream
|
|
else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start)
|
|
)
|
|
assert_omni_response(resp, request_config, run_level=self.run_level)
|
|
if resp.e2e_latency is not None:
|
|
self._print_client_stat(f"[omni] request#1 success in {resp.e2e_latency:.3f}s")
|
|
else:
|
|
self._print_client_stat("[omni] request#1 completed")
|
|
responses.append(resp)
|
|
return responses
|
|
|
|
def _one():
|
|
wall_start = time.perf_counter()
|
|
chat_completion = self.client.chat.completions.create(**create_kwargs)
|
|
return (
|
|
self._process_stream_omni_response(chat_completion, wall_start=wall_start)
|
|
if stream
|
|
else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start)
|
|
)
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
|
|
futures = {executor.submit(_one): i + 1 for i in range(request_num)}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
request_idx = futures[future]
|
|
resp = future.result()
|
|
assert_omni_response(resp, request_config, run_level=self.run_level)
|
|
if resp.e2e_latency is not None:
|
|
self._print_client_stat(f"[omni] request#{request_idx} success in {resp.e2e_latency:.3f}s")
|
|
else:
|
|
self._print_client_stat(f"[omni] request#{request_idx} completed")
|
|
responses.append(resp)
|
|
return responses
|
|
|
|
def _process_stream_audio_speech_response(
|
|
self, response, *, response_format: str | None = None, wall_start: float
|
|
) -> OmniResponse:
|
|
"""
|
|
Process streaming /v1/audio/speech responses into an OmniResponse.
|
|
|
|
This mirrors _process_stream_omni_response but operates on low-level
|
|
audio bytes. Whisper transcription runs in assert_audio_speech_response
|
|
when the run_level requires it.
|
|
"""
|
|
result = OmniResponse()
|
|
|
|
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)
|
|
|
|
# Populate OmniResponse.
|
|
result.audio_bytes = raw_bytes
|
|
result.e2e_latency = time.perf_counter() - wall_start
|
|
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:
|
|
msg = f"Audio speech stream processing error: {str(e)}"
|
|
print(f"Error: {msg}")
|
|
|
|
return result
|
|
|
|
def _process_non_stream_audio_speech_response(
|
|
self, response, *, response_format: str | None = None, wall_start: float
|
|
) -> 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()
|
|
|
|
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)}")
|
|
|
|
result.audio_bytes = raw_bytes
|
|
result.e2e_latency = time.perf_counter() - wall_start
|
|
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:
|
|
msg = f"Audio speech non-stream processing error: {str(e)}"
|
|
print(f"Error: {msg}")
|
|
|
|
return result
|
|
|
|
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)
|
|
- min_audio_bytes: optional minimum ``len(audio_bytes)`` checked in ``assert_audio_speech_response``
|
|
- status_code: if set, HTTP status is asserted (int or e.g. ``(400, 422)``); uses APIError handling
|
|
- err_message: optional substring(s) to match against error text (``str`` or list/tuple of alternatives;
|
|
see ``assert_audio_speech_response``). If set, uses the same APIError path as ``status_code``.
|
|
- When both ``status_code`` and ``err_message`` are absent (or each is ``None``), the normal request path
|
|
is used (no try/except around ``APIError``).
|
|
- 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",
|
|
"seed",
|
|
"instructions",
|
|
"speed",
|
|
"stream_format",
|
|
):
|
|
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()
|
|
|
|
print(f"[audio.speech] start model={model}, stream={stream}, request_num={request_num}, timeout={timeout:.1f}s")
|
|
|
|
# Error validation path: only when at least one of these is set to a non-``None`` value.
|
|
expect_error_handling = (request_config.get("status_code") is not None) or (
|
|
request_config.get("err_message") is not None
|
|
)
|
|
if expect_error_handling and request_num != 1:
|
|
raise ValueError(
|
|
"request_config error validation (status_code / err_message) is only supported when request_num=1"
|
|
)
|
|
|
|
if request_num == 1:
|
|
if expect_error_handling:
|
|
# ``status`` and/or ``err_message`` requested: catch APIError (4xx) and (optionally) assert body text;
|
|
# HTTP 200 with JSON error body is handled in ``assert_audio_speech_response``.
|
|
try:
|
|
if stream:
|
|
wall_start = time.perf_counter()
|
|
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, wall_start=wall_start
|
|
)
|
|
else:
|
|
wall_start = time.perf_counter()
|
|
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, wall_start=wall_start
|
|
)
|
|
except APIError as e:
|
|
sc = getattr(e, "status_code", None)
|
|
if sc is None:
|
|
raise
|
|
omni_resp = OmniResponse(
|
|
success=False,
|
|
status_code=sc,
|
|
error_message=str(e),
|
|
)
|
|
else:
|
|
if getattr(omni_resp, "status_code", None) is None:
|
|
omni_resp.status_code = 200
|
|
elif stream:
|
|
# Use streaming response helper.
|
|
wall_start = time.perf_counter()
|
|
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, wall_start=wall_start
|
|
)
|
|
else:
|
|
# Non-streaming response.
|
|
wall_start = time.perf_counter()
|
|
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, wall_start=wall_start
|
|
)
|
|
|
|
assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
|
|
if omni_resp.e2e_latency is not None:
|
|
self._print_client_stat(f"[audio.speech] request#1 success in {omni_resp.e2e_latency:.3f}s")
|
|
else:
|
|
self._print_client_stat("[audio.speech] request#1 completed")
|
|
responses.append(omni_resp)
|
|
return responses
|
|
else:
|
|
# request_num > 1: concurrent requests (use same params as single-request path)
|
|
|
|
if stream:
|
|
|
|
def _stream_task(request_idx: int):
|
|
wall_start = time.perf_counter()
|
|
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:
|
|
result = self._process_stream_audio_speech_response(
|
|
resp, response_format=speech_fmt, wall_start=wall_start
|
|
)
|
|
if result.e2e_latency is not None:
|
|
self._print_client_stat(
|
|
f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s"
|
|
)
|
|
else:
|
|
self._print_client_stat(f"[audio.speech] request#{request_idx} completed")
|
|
return result
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
|
|
futures = {executor.submit(_stream_task, i + 1): i + 1 for i in range(request_num)}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
request_idx = futures[future]
|
|
try:
|
|
omni_resp = future.result()
|
|
except Exception as e:
|
|
print(
|
|
f"[audio.speech] request#{request_idx} failed "
|
|
f"(stream={stream}, timeout={timeout:.1f}s): {e!r}"
|
|
)
|
|
raise
|
|
assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
|
|
responses.append(omni_resp)
|
|
else:
|
|
|
|
def _non_stream_task(request_idx: int):
|
|
wall_start = time.perf_counter()
|
|
r = self.client.audio.speech.create(
|
|
model=model,
|
|
input=text_input,
|
|
response_format=response_format,
|
|
extra_body=extra_body or None,
|
|
timeout=timeout,
|
|
voice=voice,
|
|
)
|
|
result = self._process_non_stream_audio_speech_response(
|
|
r, response_format=speech_fmt, wall_start=wall_start
|
|
)
|
|
if result.e2e_latency is not None:
|
|
self._print_client_stat(
|
|
f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s"
|
|
)
|
|
else:
|
|
self._print_client_stat(f"[audio.speech] request#{request_idx} completed")
|
|
return result
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
|
|
futures = {executor.submit(_non_stream_task, i + 1): i + 1 for i in range(request_num)}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
request_idx = futures[future]
|
|
try:
|
|
omni_resp = future.result()
|
|
except Exception as e:
|
|
print(
|
|
f"[audio.speech] request#{request_idx} failed "
|
|
f"(stream={stream}, timeout={timeout:.1f}s): {e!r}"
|
|
)
|
|
raise
|
|
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] | list[dict[str, Any]], request_num: int = 1
|
|
) -> list[DiffusionResponse]:
|
|
"""
|
|
Send OpenAI requests for diffusion models.
|
|
If ``extra_body`` has list ``height``/``width``, sends one chat completion per index in parallel
|
|
(scalar h/w, ``num_outputs_per_prompt=1`` each) and merges images in list order.
|
|
|
|
Args:
|
|
request_config: A single request configuration dict, or a list of
|
|
request configuration dicts (one request per element)
|
|
request_num: Number of requests to send concurrently, defaults to 1 (single request)
|
|
Returns:
|
|
list[DiffusionResponse]: List of DiffusionResponse objects containing the response data
|
|
"""
|
|
responses: list[DiffusionResponse] = []
|
|
|
|
def _create_from_config(cfg: dict[str, Any]) -> tuple[Any, float]:
|
|
stream = cfg.get("stream", False)
|
|
if stream:
|
|
raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test")
|
|
modalities = cfg.get("modalities", omit) # Most diffusion models don't require modalities param
|
|
eb = cfg.get("extra_body")
|
|
extra = copy.deepcopy(eb) if eb else None
|
|
wall_start = time.perf_counter()
|
|
chat_completion = self.client.chat.completions.create(
|
|
model=cfg.get("model"),
|
|
messages=cfg.get("messages"),
|
|
extra_body=extra,
|
|
modalities=modalities,
|
|
)
|
|
return chat_completion, wall_start
|
|
|
|
if isinstance(request_config, list):
|
|
if not request_config:
|
|
raise ValueError("request_config list must not be empty")
|
|
if request_num != 1:
|
|
raise ValueError("request_num is not supported when request_config is a list")
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=len(request_config)) as executor:
|
|
futures = {
|
|
executor.submit(_create_from_config, cfg): (i + 1, cfg) for i, cfg in enumerate(request_config)
|
|
}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
request_idx, cfg = futures[future]
|
|
chat_completion, wall_start = future.result()
|
|
response = self._process_diffusion_response(chat_completion, wall_start=wall_start)
|
|
assert_diffusion_response(response, cfg, run_level=self.run_level)
|
|
if response.e2e_latency is not None:
|
|
self._print_client_stat(
|
|
f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s"
|
|
)
|
|
else:
|
|
self._print_client_stat(f"[diffusion] request#{request_idx} completed")
|
|
responses.append(response)
|
|
return responses
|
|
|
|
size_splits = _split_request_config_by_per_output_sizes(request_config)
|
|
if size_splits is not None:
|
|
if request_num != 1:
|
|
raise ValueError(
|
|
"request_num must be 1 when extra_body height/width are lists (split into concurrent per-size calls)"
|
|
)
|
|
t0 = time.perf_counter()
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=len(size_splits)) as executor:
|
|
futures = [executor.submit(_create_from_config, sub) for sub in size_splits]
|
|
chat_completions = [f.result() for f in futures]
|
|
parts = [self._process_diffusion_response(cc, wall_start=ws) for cc, ws in chat_completions]
|
|
merged = _merge_diffusion_responses(parts)
|
|
merged.e2e_latency = time.perf_counter() - t0
|
|
assert_diffusion_response(merged, request_config, run_level=self.run_level)
|
|
if merged.e2e_latency is not None:
|
|
self._print_client_stat(f"[diffusion] request#1 success in {merged.e2e_latency:.3f}s")
|
|
else:
|
|
self._print_client_stat("[diffusion] request#1 completed")
|
|
return [merged]
|
|
|
|
if request_num == 1:
|
|
# Send single request
|
|
chat_completion, wall_start = _create_from_config(request_config)
|
|
response = self._process_diffusion_response(chat_completion, wall_start=wall_start)
|
|
assert_diffusion_response(response, request_config, run_level=self.run_level)
|
|
if response.e2e_latency is not None:
|
|
self._print_client_stat(f"[diffusion] request#1 success in {response.e2e_latency:.3f}s")
|
|
else:
|
|
self._print_client_stat("[diffusion] request#1 completed")
|
|
responses.append(response)
|
|
return responses
|
|
|
|
# Send concurrent requests for the same request_config
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
|
|
futures = {executor.submit(_create_from_config, request_config): i + 1 for i in range(request_num)}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
request_idx = futures[future]
|
|
chat_completion, wall_start = future.result()
|
|
response = self._process_diffusion_response(chat_completion, wall_start=wall_start)
|
|
assert_diffusion_response(response, request_config, run_level=self.run_level)
|
|
if response.e2e_latency is not None:
|
|
self._print_client_stat(f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s")
|
|
else:
|
|
self._print_client_stat(f"[diffusion] request#{request_idx} completed")
|
|
responses.append(response)
|
|
return responses
|
|
|
|
def send_video_diffusion_request(
|
|
self, request_config: dict[str, Any], request_num: int = 1
|
|
) -> list[DiffusionResponse]:
|
|
"""
|
|
Send native /v1/videos requests: multipart ``form_data`` job create, poll until done, download content.
|
|
|
|
For raw HTTP to video routes without polling, use ``send_videos_create_http_request``, etc.
|
|
"""
|
|
if request_num != 1:
|
|
raise NotImplementedError("Concurrent video diffusion requests are not currently implemented")
|
|
|
|
form_data = request_config.get("form_data")
|
|
if not isinstance(form_data, dict):
|
|
raise ValueError("Video request_config must contain 'form_data'")
|
|
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")
|
|
video_reference = request_config.get("video_reference")
|
|
if image_reference and video_reference:
|
|
raise ValueError("Only one of image_reference or video_reference can be provided")
|
|
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})
|
|
if video_reference:
|
|
if video_reference.startswith("data:video"):
|
|
header, encoded = video_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["video_reference"] = json.dumps({"video_url": video_reference})
|
|
|
|
result = DiffusionResponse()
|
|
create_url = self._build_url("/v1/videos")
|
|
response = requests.post(
|
|
create_url,
|
|
data=normalized_form_data,
|
|
files=files,
|
|
headers={"Accept": "application/json"},
|
|
timeout=60,
|
|
)
|
|
start_time = time.perf_counter()
|
|
response.raise_for_status()
|
|
job_data = response.json()
|
|
video_id = job_data["id"]
|
|
self._wait_until_video_completed(video_id)
|
|
end_time = time.perf_counter()
|
|
video_content = self._download_video_content(video_id)
|
|
result.success = True
|
|
result.videos = [video_content]
|
|
result.e2e_latency = end_time - start_time
|
|
assert_diffusion_response(result, request_config, run_level=self.run_level)
|
|
if result.e2e_latency is not None:
|
|
self._print_client_stat(f"[diffusion] request#1 success in {result.e2e_latency:.3f}s")
|
|
else:
|
|
self._print_client_stat("[diffusion] request#1 completed")
|
|
return [result]
|
|
|
|
def _post_json_endpoint(
|
|
self,
|
|
path: str,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
default_timeout: float,
|
|
) -> requests.Response:
|
|
url = self._build_url(path)
|
|
timeout = float(request_config.get("timeout", default_timeout))
|
|
if "raw_body" in request_config:
|
|
raw = request_config["raw_body"]
|
|
payload = raw.encode("utf-8") if isinstance(raw, str) else raw
|
|
return requests.post(
|
|
url,
|
|
data=payload,
|
|
headers={"Content-Type": "application/json", "Accept": "application/json"},
|
|
timeout=timeout,
|
|
)
|
|
if "json" not in request_config:
|
|
raise ValueError(f"{path} request_config must include 'json' or 'raw_body'")
|
|
return requests.post(
|
|
url,
|
|
json=request_config["json"],
|
|
headers={"Accept": "application/json", "Content-Type": "application/json"},
|
|
timeout=timeout,
|
|
)
|
|
|
|
def _post_form_endpoint(
|
|
self,
|
|
path: str,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
default_timeout: float = 120.0,
|
|
) -> requests.Response:
|
|
url = self._build_url(path)
|
|
timeout = float(request_config.get("timeout", default_timeout))
|
|
data = request_config.get("data")
|
|
files = request_config.get("files")
|
|
if data is None and not files:
|
|
data = {}
|
|
return requests.post(
|
|
url,
|
|
data=data,
|
|
files=files,
|
|
headers={"Accept": "application/json"} if not files else {"Accept": "application/json"},
|
|
timeout=timeout,
|
|
)
|
|
|
|
def send_streaming_video_diffusion_request(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
request_num: int = 1,
|
|
*,
|
|
timeout_seconds: float = 600.0,
|
|
) -> list[DiffusionResponse]:
|
|
"""
|
|
Send a native ``/v1/realtime/video`` WebSocket request and return one
|
|
finalized MP4 artifact assembled from the streamed binary fragments.
|
|
"""
|
|
if request_num != 1:
|
|
raise NotImplementedError("Concurrent streaming video diffusion requests are not currently implemented")
|
|
|
|
response = asyncio.run(
|
|
self._send_streaming_video_diffusion_request_once(
|
|
request_config,
|
|
timeout_seconds=timeout_seconds,
|
|
)
|
|
)
|
|
assert_diffusion_response(response, request_config, run_level=self.run_level)
|
|
if response.e2e_latency is not None:
|
|
self._print_client_stat(f"[diffusion.stream] request#1 success in {response.e2e_latency:.3f}s")
|
|
else:
|
|
self._print_client_stat("[diffusion.stream] request#1 completed")
|
|
return [response]
|
|
|
|
async def _send_streaming_video_diffusion_request_once(
|
|
self,
|
|
request_config: dict[str, Any],
|
|
*,
|
|
timeout_seconds: float,
|
|
) -> DiffusionResponse:
|
|
form_data = request_config.get("form_data")
|
|
if not isinstance(form_data, dict):
|
|
raise ValueError("Video request_config must contain 'form_data'")
|
|
payload: dict[str, Any] = {
|
|
"type": "session.start",
|
|
**{key: value for key, value in form_data.items() if value is not None},
|
|
}
|
|
model = request_config.get("model")
|
|
if model is not None:
|
|
payload["model"] = model
|
|
payload.setdefault("format", "m4s")
|
|
|
|
fps = float(payload.get("fps") or 16)
|
|
stream_format = payload["format"]
|
|
url = self._build_ws_url("/v1/realtime/video")
|
|
|
|
result = DiffusionResponse()
|
|
chunks: list[bytes] = []
|
|
start_time = time.perf_counter()
|
|
deadline = start_time + timeout_seconds
|
|
|
|
import websockets
|
|
|
|
async with websockets.connect(url, max_size=None) as websocket:
|
|
await websocket.send(json.dumps(payload))
|
|
|
|
while True:
|
|
remaining = deadline - time.perf_counter()
|
|
if remaining <= 0:
|
|
raise TimeoutError(f"Streaming video request did not complete within {timeout_seconds}s")
|
|
|
|
message = await asyncio.wait_for(websocket.recv(), timeout=remaining)
|
|
if isinstance(message, bytes):
|
|
chunks.append(message)
|
|
continue
|
|
|
|
msg = json.loads(message)
|
|
msg_type = msg.get("type")
|
|
if msg_type == "video.start":
|
|
stream_format = msg.get("format") or stream_format
|
|
continue
|
|
if msg_type == "session.done":
|
|
break
|
|
if msg_type == "error":
|
|
raise RuntimeError(str(msg.get("message", msg)))
|
|
|
|
from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes
|
|
from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat
|
|
|
|
streamed_bytes = b"".join(chunks)
|
|
if not streamed_bytes:
|
|
raise RuntimeError("Streaming video request completed without binary video chunks")
|
|
result.videos = [
|
|
finalize_streaming_video_bytes(
|
|
streamed_bytes,
|
|
input_format=cast(StreamingVideoFormat, stream_format),
|
|
fps=fps,
|
|
)
|
|
]
|
|
result.e2e_latency = time.perf_counter() - start_time
|
|
result.success = True
|
|
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('/')}"
|
|
|
|
|
|
class OmniRunner:
|
|
def __init__(
|
|
self,
|
|
model_name: str,
|
|
seed: int = 42,
|
|
stage_init_timeout: int = 600,
|
|
batch_timeout: int = 10,
|
|
# Bumped from 900s -> 1800s to give CI cold-cache loads of large
|
|
# diffusion models enough headroom (Buildkite #8418 hit a 6-second
|
|
# overrun loading Tongyi-MAI/Z-Image-Turbo: weights alone took 690s,
|
|
# the full stage was ready at ~896s, but the orchestrator wrapper
|
|
# finished at ~906s, just past the previous 900s ceiling). Engine
|
|
# production default in AsyncOmniEngine remains 600s; this only
|
|
# affects the test runner wrapper.
|
|
init_timeout: int = 1800,
|
|
shm_threshold_bytes: int = 65536,
|
|
log_stats: bool = False,
|
|
stage_configs_path: str | None = None,
|
|
**kwargs,
|
|
) -> None:
|
|
startup_t0 = time.perf_counter()
|
|
cleanup_dist_env_and_memory()
|
|
run_pre_test_cleanup()
|
|
run_post_test_cleanup()
|
|
self.model_name = model_name
|
|
self.seed = seed
|
|
self._prompt_len_estimate_cache: dict[str, Any] = {}
|
|
from vllm_omni.entrypoints.omni import Omni
|
|
|
|
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,
|
|
)
|
|
startup_s = time.perf_counter() - startup_t0
|
|
if log_stats:
|
|
print(f"OmniRunner startup took {startup_s:.3f}s (model={model_name})", flush=True)
|
|
|
|
def get_default_sampling_params_list(self) -> list[Any]:
|
|
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 _estimate_prompt_len(
|
|
self,
|
|
additional_information: dict[str, Any],
|
|
model_name: str,
|
|
) -> 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.
|
|
"""
|
|
_cache = self._prompt_len_estimate_cache
|
|
try:
|
|
from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
|
|
from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
|
|
Qwen3TTSPromptEmbedsBuilder,
|
|
)
|
|
|
|
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 Qwen3TTSPromptEmbedsBuilder.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_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]:
|
|
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 = "<VIDEO>"
|
|
image_padding_token = "<IMAGE>"
|
|
audio_padding_token = "<AUDIO>"
|
|
if isinstance(prompts, str):
|
|
prompts = [prompts]
|
|
|
|
# Qwen-TTS: follow examples/offline_inference/text_to_speech/qwen3_tts/end2end.py style.
|
|
# Stage 0 expects token placeholders + additional_information (text/speaker/task_type/...),
|
|
# and Talker replaces embeddings in preprocess based on additional_information only.
|
|
is_tts_model = "Qwen3-TTS" in self.model_name or "qwen3_tts" in self.model_name.lower()
|
|
if is_tts_model and modalities == ["audio"]:
|
|
tts_kw = mm_processor_kwargs or {}
|
|
task_type = tts_kw.get("task_type", "CustomVoice")
|
|
speaker = tts_kw.get("speaker", "Vivian")
|
|
language = tts_kw.get("language", "Auto")
|
|
max_new_tokens = int(tts_kw.get("max_new_tokens", 2048))
|
|
ref_audio = tts_kw.get("ref_audio", None)
|
|
ref_text = tts_kw.get("ref_text", None)
|
|
|
|
omni_inputs: list[TextPrompt] = []
|
|
for prompt_text in prompts:
|
|
text_str = str(prompt_text).strip() or " "
|
|
additional_information: dict[str, Any] = {
|
|
"task_type": [task_type],
|
|
"text": [text_str],
|
|
"language": [language],
|
|
"speaker": [speaker],
|
|
"max_new_tokens": [max_new_tokens],
|
|
}
|
|
if ref_audio is not None:
|
|
additional_information["ref_audio"] = [ref_audio]
|
|
if ref_text is not None:
|
|
additional_information["ref_text"] = [ref_text]
|
|
plen = self._estimate_prompt_len(additional_information, self.model_name)
|
|
input_dict: TextPrompt = {
|
|
"prompt_token_ids": [0] * plen,
|
|
"additional_information": additional_information,
|
|
}
|
|
omni_inputs.append(input_dict)
|
|
return omni_inputs
|
|
|
|
def _normalize(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("Multimodal input list length must match prompts length")
|
|
return mm_input
|
|
return [mm_input] * num_prompts
|
|
|
|
num_prompts = len(prompts)
|
|
audios_list = _normalize(audios, num_prompts)
|
|
images_list = _normalize(images, num_prompts)
|
|
videos_list = _normalize(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: dict[str, Any] = {"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[Any],
|
|
sampling_params_list: list[Any] | None = None,
|
|
) -> list[OmniRequestOutput]:
|
|
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[Any] | 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]:
|
|
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 start_profile(self, profile_prefix: str | None = None, stages: list[int] | None = None) -> list[Any]:
|
|
return self.omni.start_profile(profile_prefix=profile_prefix, stages=stages)
|
|
|
|
def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
|
|
return self.omni.stop_profile(stages=stages)
|
|
|
|
def _cleanup_process(self):
|
|
try:
|
|
keywords = ["enginecore"]
|
|
matched = []
|
|
for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
|
|
try:
|
|
cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else ""
|
|
name = proc.name().lower()
|
|
if any(k in cmdline for k in keywords) or any(k in name for k in keywords):
|
|
print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
|
|
matched.append(proc)
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
|
pass
|
|
for proc in matched:
|
|
try:
|
|
proc.terminate()
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
|
pass
|
|
_, still_alive = psutil.wait_procs(matched, timeout=5)
|
|
for proc in still_alive:
|
|
try:
|
|
proc.kill()
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
|
pass
|
|
if still_alive:
|
|
_, stubborn = psutil.wait_procs(still_alive, timeout=3)
|
|
if stubborn:
|
|
print(f"Warning: failed to kill residual vllm pids: {[p.pid for p in stubborn]}")
|
|
else:
|
|
print(f"Force-killed residual vllm pids: {[p.pid for p in still_alive]}")
|
|
elif matched:
|
|
print(f"Terminated vllm pids: {[p.pid for p in matched]}")
|
|
except Exception as e:
|
|
print(f"Error in psutil vllm cleanup: {e}")
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
if hasattr(self.omni, "close"):
|
|
self.omni.close()
|
|
self._cleanup_process()
|
|
run_pre_test_cleanup()
|
|
run_post_test_cleanup()
|
|
cleanup_dist_env_and_memory()
|
|
|
|
|
|
class OmniRunnerHandler:
|
|
def __init__(self, omni_runner: OmniRunner):
|
|
self.runner = omni_runner
|
|
|
|
def _process_omni_output(self, outputs: list[OmniRequestOutput]) -> 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.outputs[0].text
|
|
if getattr(stage_output, "final_output_type", None) == "audio":
|
|
audio_content = stage_output.request_output.outputs[0].multimodal_output["audio"]
|
|
result.audio_content = audio_content
|
|
result.text_content = text_content
|
|
result.success = True
|
|
except Exception as e:
|
|
msg = f"Output processing error: {str(e)}"
|
|
result.success = False
|
|
print(f"Error: {msg}")
|
|
return result
|
|
|
|
def _process_diffusion_output(self, outputs: list[OmniRequestOutput]) -> DiffusionResponse:
|
|
result = DiffusionResponse()
|
|
output = outputs[0]
|
|
if isinstance(output.images[0], list):
|
|
# Returning frames of images as a video
|
|
result.videos = output.images
|
|
else:
|
|
# Returning actual images
|
|
result.images = output.images
|
|
# [TODO] Add audio processing when tests are introduced
|
|
result.success = True
|
|
return result
|
|
|
|
def send_omni_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_omni_output(outputs)
|
|
assert_omni_response(response, request_config, run_level="core_model")
|
|
return response
|
|
|
|
def send_diffusion_request(self, request_config: dict[str, Any]) -> DiffusionResponse:
|
|
prompt = request_config.get("prompt")
|
|
if prompt is None:
|
|
prompts = request_config.get("prompts")
|
|
if not prompts:
|
|
raise ValueError("request_config must contain a prompt")
|
|
if len(prompts) > 1:
|
|
raise ValueError(
|
|
"In the current internal data structure, "
|
|
"only one prompt is supported for diffusion requests. "
|
|
"Because one prompt can contain multiple images or videos, "
|
|
"the current internal data structure is ambiguous when multiple prompts are provided."
|
|
)
|
|
prompt = prompts[0]
|
|
|
|
# Sync previous `extra_body` field with sampling params object
|
|
sampling_params: OmniDiffusionSamplingParams | None = request_config.get("sampling_params")
|
|
if sampling_params is None:
|
|
extra_body = request_config.get("extra_body", {})
|
|
if extra_body:
|
|
sampling_params = OmniDiffusionSamplingParams(**extra_body)
|
|
else:
|
|
extra_body = asdict(sampling_params)
|
|
request_config["extra_body"] = extra_body
|
|
if not extra_body:
|
|
logger.warning("No sampling params provided in request_config, will skip output assertion")
|
|
|
|
negative_prompt = extra_body.get("negative_prompt") or request_config.get("negative_prompt")
|
|
videos = request_config.get("videos")
|
|
images = request_config.get("images")
|
|
audios = request_config.get("audios")
|
|
# Full dict (e.g. image + mask_image for inpainting) or partial; merged with top-level image/video keys.
|
|
extra_multi_modal = request_config.get("multi_modal_data")
|
|
modalities = request_config.get("modalities") # only used by limited models. Do not add default value here
|
|
|
|
prompt_object = OmniTextPrompt(prompt=prompt)
|
|
if negative_prompt:
|
|
prompt_object["negative_prompt"] = negative_prompt
|
|
multi_modal: dict = {}
|
|
if extra_multi_modal is not None:
|
|
multi_modal.update(dict(extra_multi_modal))
|
|
if videos is not None:
|
|
multi_modal["video"] = videos
|
|
if images is not None:
|
|
multi_modal["image"] = images
|
|
if audios is not None:
|
|
multi_modal["audio"] = audios
|
|
if multi_modal:
|
|
prompt_object["multi_modal_data"] = multi_modal
|
|
if modalities:
|
|
prompt_object["modalities"] = modalities # pyright: ignore[reportGeneralTypeIssues]
|
|
|
|
start_time = time.perf_counter()
|
|
response = self.runner.generate([prompt_object], [sampling_params] if sampling_params else None)
|
|
end_time = time.perf_counter()
|
|
|
|
response = self._process_diffusion_output(response)
|
|
response.e2e_latency = end_time - start_time
|
|
assert_diffusion_response(response, request_config, run_level="core_model")
|
|
return response
|
|
|
|
def send_audio_speech_request(self, request_config: dict[str, Any]) -> OmniResponse:
|
|
"""
|
|
Offline TTS: text -> audio via generate_multimodal, then validate with assert_audio_speech_response.
|
|
|
|
request_config must contain:
|
|
- 'input' or 'prompts': text to synthesize.
|
|
Optional keys:
|
|
- 'voice' -> speaker (CustomVoice)
|
|
- 'task_type' -> task_type in additional_information (default: "CustomVoice")
|
|
- 'language' -> language in additional_information (default: "Auto")
|
|
- 'max_new_tokens' -> max_new_tokens in additional_information (default: 2048)
|
|
- 'response_format' -> desired audio format (used only for assertion)
|
|
"""
|
|
input_text = request_config.get("input") or request_config.get("prompts")
|
|
if input_text is None:
|
|
raise ValueError("request_config must contain 'input' or 'prompts' for TTS")
|
|
if isinstance(input_text, list):
|
|
input_text = input_text[0] if input_text else ""
|
|
|
|
mm_processor_kwargs: dict[str, Any] = {}
|
|
if "voice" in request_config:
|
|
mm_processor_kwargs["speaker"] = request_config["voice"]
|
|
if "task_type" in request_config:
|
|
mm_processor_kwargs["task_type"] = request_config["task_type"]
|
|
if "ref_audio" in request_config:
|
|
mm_processor_kwargs["ref_audio"] = request_config["ref_audio"]
|
|
if "ref_text" in request_config:
|
|
mm_processor_kwargs["ref_text"] = request_config["ref_text"]
|
|
if "language" in request_config:
|
|
mm_processor_kwargs["language"] = request_config["language"]
|
|
if "max_new_tokens" in request_config:
|
|
mm_processor_kwargs["max_new_tokens"] = request_config["max_new_tokens"]
|
|
|
|
outputs = self.runner.generate_multimodal(
|
|
prompts=input_text,
|
|
modalities=["audio"],
|
|
mm_processor_kwargs=mm_processor_kwargs or None,
|
|
)
|
|
mm_out: dict[str, Any] | None = None
|
|
for stage_out in outputs:
|
|
if getattr(stage_out, "final_output_type", None) == "audio":
|
|
mm_out = stage_out.request_output.outputs[0].multimodal_output
|
|
break
|
|
if mm_out is None:
|
|
raise AssertionError("No audio output from pipeline")
|
|
|
|
audio_data = mm_out.get("audio")
|
|
if audio_data is None:
|
|
raise AssertionError("No audio tensor in multimodal output")
|
|
|
|
sr_raw = mm_out.get("sr")
|
|
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
|
|
sr = int(sr_val.item() if hasattr(sr_val, "item") else sr_val)
|
|
wav_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
|
|
wav_buf = io.BytesIO()
|
|
sf.write(
|
|
wav_buf,
|
|
wav_tensor.float().cpu().numpy().reshape(-1),
|
|
samplerate=sr,
|
|
format="WAV",
|
|
subtype="PCM_16",
|
|
)
|
|
result = OmniResponse(success=True, audio_bytes=wav_buf.getvalue(), audio_format="audio/wav")
|
|
assert_audio_speech_response(result, request_config, run_level="core_model")
|
|
return result
|
|
|
|
def start_profile(self, profile_prefix: str | None = None, stages: list[int] | None = None) -> list[Any]:
|
|
return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages)
|
|
|
|
def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
|
|
return self.runner.stop_profile(stages=stages)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pytest fixture helpers (used from ``tests.helpers.fixtures.runtime``; live here
|
|
# to avoid importing ``tests.helpers.runtime`` from the plugin module at import time).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def iter_omni_server(
|
|
request: Any,
|
|
run_level: str,
|
|
model_prefix: str,
|
|
omni_fixture_lock: threading.Lock,
|
|
) -> Generator[Any, Any, None]:
|
|
"""Start/stop an Omni HTTP server; used by ``omni_server`` / ``omni_server_function`` fixtures."""
|
|
from tests.helpers.stage_config import stage_config_path_for_run_level
|
|
|
|
with omni_fixture_lock:
|
|
params: OmniServerParams = request.param
|
|
# For now, when a tiny model is substituted, we preserve the original model
|
|
# name via --served-model-name (so that the server still accepts requests with
|
|
# the original name). We also do the same for server.model so that tests reading
|
|
# server.model send the correct name in requests.
|
|
#
|
|
# TODO: core models on this path currently do not clean up tiny models, although
|
|
# tiny model paths are deterministic, so it's not a huge footprint. Still, it would
|
|
# be ideal to cleanup consistently everywhere.
|
|
original_model = model_prefix + params.model
|
|
model = original_model
|
|
if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
|
|
model = resolve_tiny_model_path(model)
|
|
port = params.port
|
|
stage_config_path = stage_config_path_for_run_level(params.stage_config_path, run_level)
|
|
|
|
server_args = params.server_args or []
|
|
if model != original_model:
|
|
server_args = [*server_args, "--served-model-name", original_model]
|
|
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"]
|
|
# ``omni_server`` / ``omni_server_function``: match ``serve`` (``--disable-log-stats`` wins).
|
|
if "--disable-log-stats" not in server_args and "--log-stats" not in server_args:
|
|
server_args = [*server_args, "--log-stats"]
|
|
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:
|
|
if model != original_model:
|
|
server.model = original_model
|
|
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:
|
|
if model != original_model:
|
|
server.model = original_model
|
|
print("OmniServer started successfully")
|
|
yield server
|
|
print("OmniServer stopping...")
|
|
|
|
print("OmniServer stopped")
|
|
|
|
|
|
def iter_omni_runner(
|
|
request: Any,
|
|
model_prefix: str,
|
|
run_level: str,
|
|
omni_fixture_lock: threading.Lock,
|
|
) -> Generator[Any, None, None]:
|
|
"""Yield an :class:`OmniRunner`; used by ``omni_runner`` / ``omni_runner_function`` fixtures."""
|
|
from tests.helpers.stage_config import stage_config_path_for_run_level
|
|
|
|
with omni_fixture_lock:
|
|
param = request.param
|
|
if not isinstance(param, (tuple, list)) or len(param) not in (2, 3):
|
|
raise ValueError(
|
|
"omni_runner param must be (model, stage_config_path) or "
|
|
"(model, stage_config_path, extra_omni_kwargs_dict)"
|
|
)
|
|
if len(param) == 2:
|
|
model, stage_config_path = param[0], param[1]
|
|
extra_omni_kwargs: dict = {}
|
|
else:
|
|
model, stage_config_path, extra = param[0], param[1], param[2]
|
|
extra_omni_kwargs = dict(extra) if extra is not None else {}
|
|
stage_config_path = stage_config_path_for_run_level(stage_config_path, run_level)
|
|
model = model_prefix + model
|
|
if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
|
|
model = resolve_tiny_model_path(model)
|
|
with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
|
|
print("OmniRunner started successfully")
|
|
yield runner
|
|
print("OmniRunner stopping...")
|
|
|
|
print("OmniRunner stopped")
|
|
|
|
|
|
__all__ = [
|
|
"DiffusionResponse",
|
|
"HttpResponse",
|
|
"WebSocketJsonResponse",
|
|
"OmniResponse",
|
|
"OmniRunner",
|
|
"OmniRunnerHandler",
|
|
"OmniServer",
|
|
"OmniServerParams",
|
|
"OmniServerStageCli",
|
|
"OpenAIClientHandler",
|
|
"get_open_port",
|
|
"dummy_messages_from_mix_data",
|
|
]
|