anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Test] Unify L2/L3 test layout, Buildkite steps, and test helpers (#2556)

Browse Source 
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
This commit is contained in:
wangyuandHuang, Zeyu authored and GitHub committed 2026-05-08 14:12:23 +08:00
1 parent b85833eeec
commit 6f2ad7b403
134 files changed
+4904 -4185
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/helpers/runtime.py
+434 -142
View File
@@ -2,6 +2,7 @@
import base64
import concurrent.futures
import copy
import errno
import io
import json
@@ -13,7 +14,7 @@ import tempfile
import threading
import time
from collections.abc import Generator
from dataclasses import dataclass
from dataclasses import asdict, dataclass
from io import BytesIO
from pathlib import Path
from typing import Any, NamedTuple
@@ -23,7 +24,7 @@ import requests
import soundfile as sf
import torch
import yaml
from openai import OpenAI, omit
from openai import APIError, OpenAI, omit
from PIL import Image
from vllm import TextPrompt
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
@@ -41,10 +42,48 @@ from tests.helpers.media import (
decode_b64_image,
)
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 _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
@@ -150,89 +189,6 @@ class OmniServerParams(NamedTuple):
stage_init_timeout: int | None = None # None: fixture supplies default (600 s)
def run_omni_server(
request: Any,
run_level: str,
model_prefix: str,
omni_fixture_lock: threading.Lock,
) -> Generator["OmniServer", Any, None]:
from tests.helpers.stage_config import modify_stage_config
with omni_fixture_lock:
params: OmniServerParams = request.param
model = model_prefix + params.model
port = params.port
stage_config_path = params.stage_config_path
if run_level in {"advanced_model", "full_model"} and stage_config_path is not None:
with open(stage_config_path, encoding="utf-8") as f:
cfg = yaml.safe_load(f) or {}
# Strip ``load_format: dummy`` (CI overlay default) so advanced_model
# tests use real weights. New schema (``stages:``) writes the field
# flat at stage level; legacy schema (``stage_args:``) nests it as
# ``engine_args.load_format``. Handle both.
new_schema_stages = cfg.get("stages")
stage_key = "stages" if new_schema_stages is not None else "stage_args"
delete_path = "load_format" if new_schema_stages is not None else "engine_args.load_format"
stage_entries = cfg.get(stage_key, [])
stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage]
stage_config_path = modify_stage_config(
stage_config_path,
deletes={stage_key: {stage_id: [delete_path] for stage_id in stage_ids}},
)
server_args = params.server_args or []
if params.use_omni and params.stage_init_timeout is not None:
server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
else:
server_args = [*server_args, "--stage-init-timeout", "600"]
if params.init_timeout is not None:
server_args = [*server_args, "--init-timeout", str(params.init_timeout)]
else:
server_args = [*server_args, "--init-timeout", "900"]
if params.use_stage_cli:
if not params.use_omni:
raise ValueError("omni_server with use_stage_cli=True requires use_omni=True")
if stage_config_path is None:
raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path")
server_args += ["--stage-configs-path", stage_config_path]
with OmniServerStageCli(
model,
stage_config_path,
server_args,
port=port,
env_dict=params.env_dict,
) as server:
print("OmniServer started successfully")
yield server
print("OmniServer stopping...")
else:
if stage_config_path is not None:
server_args += ["--stage-configs-path", stage_config_path]
with (
OmniServer(
model,
server_args,
port=port,
env_dict=params.env_dict,
use_omni=params.use_omni,
)
if port
else OmniServer(
model,
server_args,
env_dict=params.env_dict,
use_omni=params.use_omni,
)
) as server:
print("OmniServer started successfully")
yield server
print("OmniServer stopping...")
print("OmniServer stopped")
class OmniServer:
"""Omniserver for vLLM-Omni tests."""
@@ -653,6 +609,8 @@ class OmniResponse:
error_message: str | None = None
prompt_tokens: int | None = None
cached_tokens: int | None = None
# Set for HTTP-level assertions (e.g. ``assert_audio_speech_response`` with ``status_code`` in config).
status_code: int | None = None
logprobs: list | None = None
@@ -667,6 +625,20 @@ class DiffusionResponse:
error_message: str | None = None
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):
if port is None:
@@ -946,6 +918,12 @@ class OpenAIClientHandler:
- 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)
"""
@@ -970,8 +948,56 @@ class OpenAIClientHandler:
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 stream:
req_start = time.perf_counter()
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:
with self.client.audio.speech.with_streaming_response.create(
model=model,
input=text_input,
response_format=response_format,
extra_body=extra_body or None,
timeout=timeout,
voice=voice,
) as resp:
omni_resp = self._process_stream_audio_speech_response(resp, response_format=speech_fmt)
else:
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)
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.
with self.client.audio.speech.with_streaming_response.create(
model=model,
@@ -995,6 +1021,8 @@ class OpenAIClientHandler:
omni_resp = self._process_non_stream_audio_speech_response(resp, response_format=speech_fmt)
assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
elapsed = time.perf_counter() - req_start
print(f"[audio.speech] request#1 success in {elapsed:.3f}s")
responses.append(omni_resp)
return responses
else:
@@ -1002,7 +1030,8 @@ class OpenAIClientHandler:
if stream:
def _stream_task():
def _stream_task(request_idx: int):
task_start = time.perf_counter()
with self.client.audio.speech.with_streaming_response.create(
model=model,
input=text_input,
@@ -1011,83 +1040,137 @@ class OpenAIClientHandler:
timeout=timeout,
voice=voice,
) as resp:
return self._process_stream_audio_speech_response(resp, response_format=speech_fmt)
result = self._process_stream_audio_speech_response(resp, response_format=speech_fmt)
elapsed = time.perf_counter() - task_start
print(f"[audio.speech] request#{request_idx} success in {elapsed:.3f}s")
return result
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
futures = [executor.submit(_stream_task) for _ in range(request_num)]
futures = {executor.submit(_stream_task, i + 1): i + 1 for i in range(request_num)}
for future in concurrent.futures.as_completed(futures):
omni_resp = future.result()
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):
task_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,
)
elapsed = time.perf_counter() - task_start
print(f"[audio.speech] request#{request_idx} success in {elapsed:.3f}s")
return resp
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
futures = []
for _ in range(request_num):
future = executor.submit(
self.client.audio.speech.create,
model=model,
input=text_input,
response_format=response_format,
extra_body=extra_body or None,
timeout=timeout,
voice=voice,
)
futures.append(future)
futures = {executor.submit(_non_stream_task, i + 1): i + 1 for i in range(request_num)}
for future in concurrent.futures.as_completed(futures):
resp = future.result()
request_idx = futures[future]
try:
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
omni_resp = self._process_non_stream_audio_speech_response(resp, response_format=speech_fmt)
assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
responses.append(omni_resp)
return responses
def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[DiffusionResponse]:
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: Request configuration dictionary containing parameters like model, messages
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] = []
stream = request_config.get("stream", False)
modalities = request_config.get("modalities", omit) # Most diffusion models don't require modalities param
extra_body = request_config.get("extra_body", None)
if stream:
raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test")
if request_num == 1:
# Send single request
chat_completion = self.client.chat.completions.create(
model=request_config.get("model"),
messages=request_config.get("messages"),
extra_body=extra_body,
def _create_from_config(cfg: dict[str, Any]):
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
return self.client.chat.completions.create(
model=cfg.get("model"),
messages=cfg.get("messages"),
extra_body=extra,
modalities=modalities,
)
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): cfg for cfg in request_config}
for future in concurrent.futures.as_completed(futures):
cfg = futures[future]
chat_completion = future.result()
response = self._process_diffusion_response(chat_completion)
assert_diffusion_response(response, cfg, run_level=self.run_level)
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) for cc 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)
return [merged]
if request_num == 1:
# Send single request
chat_completion = _create_from_config(request_config)
response = self._process_diffusion_response(chat_completion)
assert_diffusion_response(response, request_config, run_level=self.run_level)
responses.append(response)
else:
# Send concurrent requests
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
futures = []
# Submit all request tasks
for _ in range(request_num):
future = executor.submit(
self.client.chat.completions.create,
model=request_config.get("model"),
messages=request_config.get("messages"),
modalities=modalities,
extra_body=extra_body,
)
futures.append(future)
# Process completed tasks
for future in concurrent.futures.as_completed(futures):
chat_completion = future.result()
response = self._process_diffusion_response(chat_completion)
assert_diffusion_response(response, request_config, run_level=self.run_level)
responses.append(response)
return responses
# 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) for _ in range(request_num)]
for future in concurrent.futures.as_completed(futures):
chat_completion = future.result()
response = self._process_diffusion_response(chat_completion)
assert_diffusion_response(response, request_config, run_level=self.run_level)
responses.append(response)
return responses
def send_video_diffusion_request(
@@ -1380,7 +1463,7 @@ class OmniRunner:
self,
prompts: list[Any],
sampling_params_list: list[Any] | None = None,
) -> list[Any]:
) -> 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)
@@ -1395,7 +1478,7 @@ class OmniRunner:
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
modalities: list[str] | None = None,
) -> list[Any]:
) -> list[OmniRequestOutput]:
omni_inputs = self.get_omni_inputs(
prompts=prompts,
system_prompt=system_prompt,
@@ -1460,10 +1543,10 @@ class OmniRunner:
class OmniRunnerHandler:
def __init__(self, omni_runner):
def __init__(self, omni_runner: OmniRunner):
self.runner = omni_runner
def _process_output(self, outputs: list[Any]) -> OmniResponse:
def _process_omni_output(self, outputs: list[OmniRequestOutput]) -> OmniResponse:
result = OmniResponse()
try:
text_content = None
@@ -1482,7 +1565,20 @@ class OmniRunnerHandler:
print(f"Error: {result.error_message}")
return result
def send_request(self, request_config: dict[str, Any] | None = None) -> OmniResponse:
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")
@@ -1493,10 +1589,71 @@ class OmniRunnerHandler:
outputs = self.runner.generate_multimodal(
prompts=prompts, videos=videos, images=images, audios=audios, modalities=modalities
)
response = self._process_output(outputs)
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.
@@ -1574,6 +1731,141 @@ class OmniRunnerHandler:
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 _core_model_stage_config_path_with_dummy_load_format(stage_config_path: str | None, run_level: str) -> str | None:
"""For ``core_model`` runs, patch every stage in the deploy YAML to ``load_format: dummy``."""
if run_level != "core_model" or stage_config_path is None:
return stage_config_path
from tests.helpers.stage_config import modify_stage_config
with open(stage_config_path, encoding="utf-8") as f:
cfg = yaml.safe_load(f) or {}
new_schema_stages = cfg.get("stages")
stage_key = "stages" if new_schema_stages is not None else "stage_args"
update_path = "load_format" if new_schema_stages is not None else "engine_args.load_format"
stage_entries = cfg.get(stage_key, [])
stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage]
return modify_stage_config(
stage_config_path,
updates={stage_key: {stage_id: {update_path: "dummy"} for stage_id in stage_ids}},
)
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 modify_stage_config
with omni_fixture_lock:
params: OmniServerParams = request.param
model = model_prefix + params.model
port = params.port
stage_config_path = params.stage_config_path
if run_level in {"advanced_model", "full_model"} and stage_config_path is not None:
with open(stage_config_path, encoding="utf-8") as f:
cfg = yaml.safe_load(f) or {}
new_schema_stages = cfg.get("stages")
stage_key = "stages" if new_schema_stages is not None else "stage_args"
delete_path = "load_format" if new_schema_stages is not None else "engine_args.load_format"
stage_entries = cfg.get(stage_key, [])
stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage]
stage_config_path = modify_stage_config(
stage_config_path,
deletes={stage_key: {stage_id: [delete_path] for stage_id in stage_ids}},
)
server_args = params.server_args or []
if params.use_omni and params.stage_init_timeout is not None:
server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
else:
server_args = [*server_args, "--stage-init-timeout", "600"]
if params.init_timeout is not None:
server_args = [*server_args, "--init-timeout", str(params.init_timeout)]
else:
server_args = [*server_args, "--init-timeout", "900"]
if params.use_stage_cli:
if not params.use_omni:
raise ValueError("omni_server with use_stage_cli=True requires use_omni=True")
if stage_config_path is None:
raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path")
server_args += ["--stage-configs-path", stage_config_path]
with OmniServerStageCli(
model,
stage_config_path,
server_args,
port=port,
env_dict=params.env_dict,
) as server:
print("OmniServer started successfully")
yield server
print("OmniServer stopping...")
else:
if stage_config_path is not None:
server_args += ["--stage-configs-path", stage_config_path]
with (
OmniServer(
model,
server_args,
port=port,
env_dict=params.env_dict,
use_omni=params.use_omni,
)
if port
else OmniServer(
model,
server_args,
env_dict=params.env_dict,
use_omni=params.use_omni,
)
) as server:
print("OmniServer started successfully")
yield server
print("OmniServer stopping...")
print("OmniServer stopped")
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."""
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 = _core_model_stage_config_path_with_dummy_load_format(stage_config_path, run_level)
model = model_prefix + 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",
"OmniResponse",