2014 lines
71 KiB
Python
2014 lines
71 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""
|
|
Unit tests for patch.py
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import time
|
|
from argparse import Namespace
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from pytest_mock import MockerFixture
|
|
from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
|
|
|
|
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
|
|
SeedTTSSampleRequest,
|
|
SeedTTSTextSampleRequest,
|
|
)
|
|
from vllm_omni.benchmarks.patch.patch import (
|
|
MixRequestFuncOutput,
|
|
_add_combined_video_form_references,
|
|
_add_video_extra_body_to_form,
|
|
_add_video_reference_to_form,
|
|
_apply_image_metrics_from_payload,
|
|
_apply_stage0_token_timings,
|
|
_apply_video_metrics_from_payload,
|
|
_attach_seed_tts_to_request_func_input,
|
|
_build_benchmark_session,
|
|
_extract_stage_durations_from_payload,
|
|
_iter_image_reference_inputs,
|
|
_iter_video_reference_inputs,
|
|
_omni_request_timeout_s,
|
|
async_request_openai_chat_omni_completions,
|
|
async_request_openai_image_edits_omni,
|
|
async_request_openai_image_generations_omni,
|
|
async_request_openai_realtime_duplex,
|
|
should_request_stage_metrics,
|
|
)
|
|
from vllm_omni.clients.duplex import EventCollector
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
|
|
|
|
|
|
def _seed_tts_request_func_input() -> RequestFuncInput:
|
|
return RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="target text",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=2,
|
|
output_len=20,
|
|
extra_body={"modalities": ["text", "audio"]},
|
|
)
|
|
|
|
|
|
def test_seed_tts_chat_request_carries_reference_audio_once() -> None:
|
|
reference_audio_url = "data:audio/wav;base64,AAAA"
|
|
sample = SeedTTSSampleRequest(
|
|
prompt="target text",
|
|
prompt_len=2,
|
|
expected_output_len=20,
|
|
multi_modal_data=None,
|
|
seed_tts_system_prompt="Clone the supplied voice.",
|
|
seed_tts_speech_extra={
|
|
"ref_audio": reference_audio_url,
|
|
"ref_text": "reference text",
|
|
"task_type": "Base",
|
|
},
|
|
)
|
|
request_func_input = _seed_tts_request_func_input()
|
|
|
|
_attach_seed_tts_to_request_func_input(sample, request_func_input)
|
|
|
|
assert request_func_input.omni_chat_messages == [
|
|
{
|
|
"role": "system",
|
|
"content": [{"type": "text", "text": "Clone the supplied voice."}],
|
|
},
|
|
{"role": "user", "content": [{"type": "text", "text": "target text"}]},
|
|
]
|
|
assert request_func_input.extra_body["ref_audio"] == reference_audio_url
|
|
serialized_request = json.dumps(
|
|
{"messages": request_func_input.omni_chat_messages, **request_func_input.extra_body},
|
|
)
|
|
assert serialized_request.count(reference_audio_url) == 1
|
|
|
|
|
|
def test_seed_tts_text_chat_messages_do_not_add_reference_audio() -> None:
|
|
sample = SeedTTSTextSampleRequest(
|
|
prompt="target text",
|
|
prompt_len=2,
|
|
expected_output_len=20,
|
|
multi_modal_data=None,
|
|
seed_tts_system_prompt="Use the configured voice.",
|
|
seed_tts_speech_extra=None,
|
|
)
|
|
request_func_input = _seed_tts_request_func_input()
|
|
|
|
_attach_seed_tts_to_request_func_input(sample, request_func_input)
|
|
|
|
assert getattr(request_func_input, "seed_tts_row", False) is True
|
|
assert request_func_input.omni_chat_messages == [
|
|
{
|
|
"role": "system",
|
|
"content": [{"type": "text", "text": "Use the configured voice."}],
|
|
},
|
|
{"role": "user", "content": [{"type": "text", "text": "target text"}]},
|
|
]
|
|
assert "ref_audio" not in request_func_input.extra_body
|
|
|
|
|
|
class MockResponse:
|
|
"""Mock aiohttp response for testing"""
|
|
|
|
def __init__(self, status, chunks, delay_between_chunks=0):
|
|
self.status = status
|
|
self.reason = "OK" if status == 200 else "Error"
|
|
self._chunks = chunks
|
|
self._delay = delay_between_chunks
|
|
self.content = self
|
|
|
|
async def iter_any(self):
|
|
for chunk in self._chunks:
|
|
if self._delay > 0:
|
|
await asyncio.sleep(self._delay)
|
|
yield chunk
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch):
|
|
class FakeRealtimeClient:
|
|
last_instance = None
|
|
instances: list = []
|
|
|
|
def __init__(self, url):
|
|
assert url == "ws://localhost:8000/v1/realtime?duplex=1"
|
|
self.events = EventCollector()
|
|
self.configure_kwargs = None
|
|
self.sent = []
|
|
self.response_count = 0
|
|
self.ack_count = 0
|
|
self.silence_seconds = []
|
|
self.closed = 0
|
|
type(self).instances.append(self)
|
|
type(self).last_instance = self
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
return None
|
|
|
|
async def configure(self, model, **kwargs):
|
|
assert model == "openbmb/MiniCPM-o-4_5"
|
|
self.configure_kwargs = kwargs
|
|
|
|
async def stream_silence(self, *, seconds, chunk_ms=200, until=None):
|
|
# A model-native session speaks off its seeded context once audio
|
|
# units arrive; the silence itself carries no content. The probe
|
|
# stops the silence as soon as the turn settles.
|
|
self.silence_seconds.append(seconds)
|
|
self._emit_response()
|
|
assert until is not None and until()
|
|
return 1.0
|
|
|
|
async def send(self, event):
|
|
self.sent.append(event)
|
|
if event["type"] != "response.create":
|
|
return
|
|
self._emit_response()
|
|
|
|
def _emit_response(self):
|
|
self.response_count += 1
|
|
response_id = f"resp-{self.response_count}"
|
|
now = time.monotonic()
|
|
self.events.add(
|
|
{"type": "response.created", "response": {"id": response_id}},
|
|
received_at_s=now + 0.01,
|
|
)
|
|
self.events.add(
|
|
{
|
|
"type": "response.output_audio_transcript.delta",
|
|
"response_id": response_id,
|
|
"delta": f"turn {self.response_count}",
|
|
},
|
|
received_at_s=now + 0.02,
|
|
)
|
|
self.events.add(
|
|
{
|
|
"type": "response.output_audio.delta",
|
|
"response_id": response_id,
|
|
"delta": base64.b64encode(b"\x00\x00" * 2400).decode(),
|
|
"sample_rate_hz": 24_000,
|
|
"metadata": {"audio_duration_ms": 100},
|
|
},
|
|
received_at_s=now + 0.03,
|
|
)
|
|
self.events.add(
|
|
{
|
|
"type": "response.done",
|
|
"response": {
|
|
"id": response_id,
|
|
"metadata": {
|
|
"vllm_omni": {
|
|
"stage_metrics": {
|
|
"0": {
|
|
"num_tokens_out": 3,
|
|
"vllm_ttft_ms": 20.0,
|
|
"vllm_tpot_ms": 10.0,
|
|
"vllm_itls_ms": [10.0, 10.0],
|
|
}
|
|
}
|
|
}
|
|
},
|
|
},
|
|
},
|
|
received_at_s=now + 0.04,
|
|
)
|
|
|
|
async def acknowledge_playback(self):
|
|
self.ack_count += 1
|
|
|
|
async def close_session(self, **_kwargs):
|
|
self.closed += 1
|
|
|
|
monkeypatch.setattr(
|
|
"vllm_omni.benchmarks.patch.patch._RealtimeTTSProbe",
|
|
FakeRealtimeClient,
|
|
)
|
|
request_input = RequestFuncInput(
|
|
model="openbmb/MiniCPM-o-4_5",
|
|
model_name="openbmb/MiniCPM-o-4_5",
|
|
prompt="hello",
|
|
api_url="http://localhost:8000/v1/realtime",
|
|
prompt_len=1,
|
|
output_len=20,
|
|
logprobs=None,
|
|
multi_modal_content=None,
|
|
ignore_eos=False,
|
|
extra_body={"save_duplex_request_metrics": True},
|
|
)
|
|
request_input.seed_tts_speech_extra = {"ref_audio": "data:audio/wav;base64,AAAA"}
|
|
request_input.seed_tts_system_prompt = "Speak exactly."
|
|
request_input.seed_tts_turns = tuple(
|
|
SimpleNamespace(utterance_id=f"utt-{index}", target_text=f"text {index}") for index in range(4)
|
|
)
|
|
|
|
output = await async_request_openai_realtime_duplex(
|
|
request_input,
|
|
session=None,
|
|
)
|
|
|
|
# One session per utterance: a model-native duplex session takes its text
|
|
# once, in the session context, so it cannot be re-seeded for a second one.
|
|
sessions = FakeRealtimeClient.instances
|
|
assert len(sessions) == 4
|
|
for index, client in enumerate(sessions):
|
|
assert "native_duplex" not in client.configure_kwargs
|
|
assert client.configure_kwargs["extra_body"] == {
|
|
"ref_audio": "data:audio/wav;base64,AAAA",
|
|
"return_stage_metrics": True,
|
|
"duplex_initial_user_text": f"text {index}",
|
|
"force_listen_count": 0,
|
|
}
|
|
# The target text rides the session context, never a conversation item:
|
|
# a text-only response.create is rejected by a model-native session.
|
|
assert client.sent == []
|
|
assert client.silence_seconds and all(seconds > 0 for seconds in client.silence_seconds)
|
|
assert client.ack_count == 1
|
|
assert client.closed == 1
|
|
|
|
assert output.success is True
|
|
assert output.generated_text == "turn 1 turn 1 turn 1 turn 1"
|
|
assert output.audio_duration == pytest.approx(0.4)
|
|
assert output.ttft > 0
|
|
assert output.audio_ttfp > output.ttft
|
|
assert output.audio_rtf > 0
|
|
assert output.latency > 0
|
|
assert output.output_tokens == 12
|
|
assert output.itl == [0.01, 0.01] * 4
|
|
assert output.text_latency == pytest.approx(output.ttft + 0.08)
|
|
assert output.tpot_measured is True
|
|
assert output.tts_turn_pcm_bytes == [b"\x00\x00" * 2400] * 4
|
|
assert output.tts_output_pcm_bytes == b"\x00\x00" * 9600
|
|
session_id = output.duplex_request_metrics[0]["session_id"]
|
|
assert len(output.duplex_request_metrics) == 4
|
|
for index, metrics in enumerate(output.duplex_request_metrics):
|
|
assert metrics == {
|
|
"session_id": session_id,
|
|
"request_index": index,
|
|
"utterance_id": f"utt-{index}",
|
|
# Response ids are per session, and each utterance now gets its own.
|
|
"response_id": "resp-1",
|
|
"source": "client_monotonic_receive",
|
|
"measurement_origin": {
|
|
"ttft": "first silence append client send to first non-empty text delta",
|
|
"ttfp": "first silence append client send to first audio packet",
|
|
"rtf": "request-start-to-last-audio receive time divided by emitted audio duration",
|
|
"tpot": "Stage-0 engine mean time per output token",
|
|
},
|
|
"ttft_ms": pytest.approx(20.0, abs=2.0),
|
|
"tpot_ms": 10.0,
|
|
"ttfp_ms": pytest.approx(30.0, abs=2.0),
|
|
"rtf": pytest.approx(0.3, abs=0.03),
|
|
"audio_generation_ms": pytest.approx(30.0, abs=2.0),
|
|
"audio_duration_ms": 100.0,
|
|
}
|
|
session = output.duplex_session_metrics
|
|
assert session["session_id"] == session_id
|
|
assert session["audio_turn_count"] == 4
|
|
assert session["ttft_ms"]["count"] == 4
|
|
assert session["ttft_ms"]["mean"] == pytest.approx(20.0, abs=2.0)
|
|
assert session["tpot_ms"] == {"count": 4, "mean": 10.0, "p50": 10.0, "p99": 10.0}
|
|
assert session["ttfp_ms"]["count"] == 4
|
|
assert session["ttfp_ms"]["mean"] == pytest.approx(30.0, abs=2.0)
|
|
assert session["rtf"]["count"] == 4
|
|
assert session["rtf"]["mean"] == pytest.approx(0.3, abs=0.03)
|
|
|
|
|
|
def test_stage0_token_timings_use_weighted_tpot_when_itls_are_incomplete():
|
|
output = MixRequestFuncOutput(ttft=0.1)
|
|
|
|
measured = _apply_stage0_token_timings(
|
|
output,
|
|
[
|
|
{"output_token_count": 3, "itls_ms": [], "tpot_ms": 10.0},
|
|
{"output_token_count": 2, "itls_ms": [], "tpot_ms": 40.0},
|
|
],
|
|
expected_output_tokens=5,
|
|
)
|
|
|
|
assert measured is True
|
|
assert output.itl == []
|
|
assert output.tpot_measured is True
|
|
assert (output.text_latency - output.ttft) / 4 == pytest.approx(0.02)
|
|
|
|
|
|
def test_stage0_zero_itls_fall_back_to_positive_tpot():
|
|
output = MixRequestFuncOutput(ttft=0.1)
|
|
|
|
measured = _apply_stage0_token_timings(
|
|
output,
|
|
[{"output_token_count": 3, "itls_ms": [0.0, 0.0], "tpot_ms": 25.0}],
|
|
expected_output_tokens=3,
|
|
)
|
|
|
|
assert measured is True
|
|
assert output.itl == []
|
|
assert output.tpot_measured is True
|
|
assert (output.text_latency - output.ttft) / 2 == pytest.approx(0.025)
|
|
|
|
|
|
def test_stage0_token_count_without_timing_is_not_measured():
|
|
output = MixRequestFuncOutput(ttft=0.1)
|
|
|
|
measured = _apply_stage0_token_timings(
|
|
output,
|
|
[{"output_token_count": 3, "itls_ms": [], "tpot_ms": None}],
|
|
expected_output_tokens=3,
|
|
)
|
|
|
|
assert measured is False
|
|
assert output.itl == []
|
|
assert output.text_latency == output.ttft
|
|
assert output.tpot_measured is False
|
|
|
|
|
|
def create_sse_chunk(data_dict):
|
|
"""Helper to create SSE formatted chunk"""
|
|
return f"data: {json.dumps(data_dict)}\n\n".encode()
|
|
|
|
|
|
def test_chat_text_timing_metrics_request_stage_metrics():
|
|
args = Namespace(
|
|
backend="openai-chat-omni",
|
|
percentile_metrics="ttft,tpot,itl,e2el",
|
|
print_stage=False,
|
|
extra_body={},
|
|
)
|
|
|
|
assert should_request_stage_metrics(args) is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bundled_first_text_chunk_uses_stage0_token_timings(mocker: MockerFixture):
|
|
"""Engine timings recover TPOT when every text token arrives together."""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "ABCD"}}],
|
|
"modality": "text",
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
|
|
"metrics": {
|
|
"stage_metrics": {
|
|
"0": {
|
|
"num_tokens_out": 4,
|
|
"vllm_itls_ms": [10.0, 11.0, 12.0],
|
|
"vllm_tpot_ms": 11.0,
|
|
}
|
|
}
|
|
},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
assert output.success is True
|
|
assert output.output_tokens == 4
|
|
assert output.itl == pytest.approx([0.010, 0.011, 0.012])
|
|
assert output.text_latency - output.ttft == pytest.approx(0.033)
|
|
assert output.tpot_measured is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_error_chunk_marks_request_failed(mocker: MockerFixture):
|
|
"""HTTP 200 streams can still terminate with an OpenAI error event."""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "partial response"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"error": {
|
|
"message": "EngineCore encountered an issue",
|
|
"type": "BadRequestError",
|
|
"code": 400,
|
|
}
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
assert output.success is False
|
|
assert output.error == "EngineCore encountered an issue"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_positive_client_text_timings_take_precedence_over_stage0(mocker: MockerFixture):
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "A"}}],
|
|
"modality": "text",
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "B"}}],
|
|
"modality": "text",
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
|
|
"metrics": {
|
|
"stage_metrics": {
|
|
"0": {
|
|
"num_tokens_out": 2,
|
|
"vllm_itls_ms": [1000.0],
|
|
"vllm_tpot_ms": 1000.0,
|
|
}
|
|
}
|
|
},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
mock_response = MockResponse(200, chunks, delay_between_chunks=0.01)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
assert output.success is True
|
|
assert len(output.itl) == 1
|
|
assert 0.0 < output.itl[0] < 0.1
|
|
|
|
|
|
# ============================================================================
|
|
# output_tokens Tests
|
|
# ============================================================================
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_tokens_assigned_with_metrics(mocker: MockerFixture):
|
|
"""Test that output.output_tokens is assigned when metrics are present"""
|
|
# Arrange
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response chunks with metrics
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Hello"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " world"}}],
|
|
"modality": "text",
|
|
"metrics": {"num_tokens_out": 42, "num_tokens_in": 10},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert output.output_tokens == 42, "output_tokens should be assigned from metrics"
|
|
assert output.generated_text == "Hello world"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_tokens_not_assigned_without_metrics(mocker: MockerFixture):
|
|
"""Test that output.output_tokens defaults to 0 when no metrics present"""
|
|
# Arrange
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response chunks without metrics
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Hello"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " world"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
# Without server metrics we no longer infer count from SSE chunks because
|
|
# one chunk may carry multiple tokens. output_tokens stays 0 here.
|
|
assert output.output_tokens == 0, "output_tokens should be 0 when no server metrics are present"
|
|
assert output.generated_text == "Hello world"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_tokens_assigned_multiple_metrics(mocker: MockerFixture):
|
|
"""Test that output.output_tokens is updated with the latest metrics value"""
|
|
# Arrange
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response chunks with multiple metrics updates
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Hello"}}],
|
|
"modality": "text",
|
|
"metrics": {"num_tokens_out": 5},
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " world"}}],
|
|
"modality": "text",
|
|
"metrics": {"num_tokens_out": 10},
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "!"}}],
|
|
"modality": "text",
|
|
"metrics": {"num_tokens_out": 15},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert output.output_tokens == 15, "output_tokens should be updated to the latest value"
|
|
assert output.generated_text == "Hello world!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_tokens_with_audio_and_text(mocker: MockerFixture):
|
|
"""Test output_tokens assignment in mixed audio and text response"""
|
|
# Arrange
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response chunks with both audio and text, with metrics
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Text response"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": ""}}],
|
|
"modality": "audio",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"modality": "text",
|
|
"metrics": {"num_tokens_out": 25, "num_tokens_in": 10},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert output.output_tokens == 25, "output_tokens should be assigned even with audio modality"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_tokens_with_missing_num_tokens_out(mocker: MockerFixture):
|
|
"""Test that output_tokens defaults to 0 when num_tokens_out is missing in metrics"""
|
|
# Arrange
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response chunks with metrics but without num_tokens_out
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Hello"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " world"}}],
|
|
"modality": "text",
|
|
"metrics": {"num_tokens_in": 10}, # Missing num_tokens_out
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert output.output_tokens == 0, "output_tokens should be 0 when num_tokens_out is missing"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_tokens_initialization():
|
|
"""Test that MixRequestFuncOutput initializes output_tokens correctly"""
|
|
# Arrange & Act
|
|
output = MixRequestFuncOutput()
|
|
|
|
# Assert
|
|
assert hasattr(output, "output_tokens"), "MixRequestFuncOutput should have output_tokens attribute"
|
|
assert output.output_tokens == 0, "output_tokens should be initialized to 0"
|
|
|
|
|
|
# ============================================================================
|
|
# text_latency Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestTextLatencyAttribute:
|
|
"""Tests for text_latency attribute existence and assignment"""
|
|
|
|
def test_mix_request_func_output_has_text_latency(self):
|
|
"""Test that MixRequestFuncOutput has text_latency attribute"""
|
|
output = MixRequestFuncOutput()
|
|
assert hasattr(output, "text_latency"), "MixRequestFuncOutput should have text_latency attribute"
|
|
|
|
def test_text_latency_initial_value(self):
|
|
"""Test that text_latency initializes to a default value"""
|
|
output = MixRequestFuncOutput()
|
|
# Check if attribute exists and has a value (should be 0.0 or similar default)
|
|
text_latency = getattr(output, "text_latency", None)
|
|
assert text_latency is not None or hasattr(output, "text_latency"), "text_latency attribute should exist"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_latency_assigned_with_text_response(self, mocker: MockerFixture):
|
|
"""Test that text_latency is assigned when text response is received"""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response chunks with text modality
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Hello"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " world"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert hasattr(output, "text_latency"), "Output should have text_latency attribute"
|
|
assert output.text_latency > 0, "text_latency should be greater than 0 for text response"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_latency_updated_with_multiple_text_chunks(self, mocker: MockerFixture):
|
|
"""Test that text_latency is updated as more text chunks arrive"""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response chunks with small delays to simulate streaming
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "First"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " second"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " third"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert hasattr(output, "text_latency"), "Output should have text_latency attribute"
|
|
assert output.text_latency > 0, "text_latency should accumulate"
|
|
assert output.generated_text == "First second third"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_latency_with_only_audio_response(self, mocker: MockerFixture):
|
|
"""Test text_latency behavior when only audio is received"""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response chunks with only audio modality (no text)
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": ""}}],
|
|
"modality": "audio",
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert hasattr(output, "text_latency"), "Output should have text_latency attribute even with audio-only"
|
|
# text_latency should either be 0 or the initial value since no text was processed
|
|
assert output.text_latency >= 0, "text_latency should be non-negative"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_latency_not_affected_by_metrics(self, mocker: MockerFixture):
|
|
"""Test that text_latency is independent of metrics data"""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response with text and metrics
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Response text"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": ""}}],
|
|
"modality": "text",
|
|
"metrics": {"num_tokens_out": 100, "num_tokens_in": 20},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert hasattr(output, "text_latency"), "text_latency should exist"
|
|
assert output.text_latency > 0, "text_latency should be set for text response"
|
|
# Verify metrics are also present
|
|
assert output.output_tokens == 100, "metrics should not affect text_latency"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_latency_mixed_modalities(self, mocker: MockerFixture):
|
|
"""Test text_latency with mixed text and audio modalities"""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
# Create response with both text and audio
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Text"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": ""}}],
|
|
"modality": "audio",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " more text"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert hasattr(output, "text_latency"), "text_latency should exist with mixed modalities"
|
|
assert output.text_latency > 0, "text_latency should be set when text is present"
|
|
assert output.generated_text == "Text more text"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_latency_value_consistency(self, mocker: MockerFixture):
|
|
"""Test that text_latency matches latency minus ttft relationship"""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Hello"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " world"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert hasattr(output, "text_latency"), "text_latency should exist"
|
|
assert hasattr(output, "ttft"), "ttft should exist"
|
|
assert hasattr(output, "latency"), "latency should exist"
|
|
# text_latency should be between ttft and total latency
|
|
assert output.ttft <= output.text_latency <= output.latency, (
|
|
"text_latency should be between ttft and total latency"
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# prompt_len Tests
|
|
# ============================================================================
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skips_empty_role_chunk_for_ttft_and_itl(mocker: MockerFixture):
|
|
"""Role/empty text chunks must not count toward TTFT or ITL."""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"role": "assistant", "content": ""}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Hello"}}],
|
|
"modality": "text",
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " world"}}],
|
|
"modality": "text",
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks, delay_between_chunks=0.01)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
assert output.success is True
|
|
assert output.output_tokens == 2
|
|
assert len(output.itl) == 1
|
|
assert output.ttft > 0
|
|
assert output.text_latency - output.ttft == pytest.approx(sum(output.itl), rel=1e-6, abs=1e-6)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bundled_tokens_split_itl_from_usage(mocker: MockerFixture):
|
|
"""Multiple tokens in one SSE chunk should expand ITL via usage deltas."""
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "A"}}],
|
|
"modality": "text",
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "BCD"}}],
|
|
"modality": "text",
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks, delay_between_chunks=0.02)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
assert output.success is True
|
|
assert output.output_tokens == 4
|
|
assert len(output.itl) == 3
|
|
assert output.text_latency - output.ttft == pytest.approx(sum(output.itl), rel=1e-6, abs=1e-6)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_tokens_prefers_usage_over_metrics(mocker: MockerFixture):
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Hello"}}],
|
|
"modality": "text",
|
|
"metrics": {"num_tokens_out": 5},
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
assert output.success is True
|
|
assert output.output_tokens == 8
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_metrics_only_chunk_updates_output_tokens(mocker: MockerFixture):
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=10,
|
|
output_len=20,
|
|
)
|
|
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Text response"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"modality": "text",
|
|
"metrics": {"num_tokens_out": 25, "num_tokens_in": 10},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
assert output.success is True
|
|
assert output.output_tokens == 25
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prompt_len_assigned_from_usage(mocker: MockerFixture):
|
|
# Arrange: request claims prompt_len=100, but server reports 4992 (multimodal).
|
|
request_input = RequestFuncInput(
|
|
model="test-model",
|
|
model_name="test-model",
|
|
prompt="test prompt",
|
|
api_url="http://test.com/v1/chat/completions",
|
|
prompt_len=100,
|
|
output_len=20,
|
|
)
|
|
|
|
chunks = [
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": "Hello"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [{"delta": {"content": " world"}}],
|
|
"modality": "text",
|
|
}
|
|
),
|
|
# Final usage chunk emitted because stream_options.include_usage=True.
|
|
create_sse_chunk(
|
|
{
|
|
"choices": [],
|
|
"usage": {"prompt_tokens": 4992, "completion_tokens": 2, "total_tokens": 4994},
|
|
}
|
|
),
|
|
b"data: [DONE]\n\n",
|
|
]
|
|
|
|
mock_response = MockResponse(200, chunks)
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
# Act
|
|
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
|
|
|
# Assert
|
|
assert output.success is True
|
|
assert output.prompt_len == 4992, (
|
|
"prompt_len should be overridden by usage.prompt_tokens to reflect the true multimodal input token count"
|
|
)
|
|
|
|
|
|
class TestOmniRequestTimeout:
|
|
"""``--omni-request-timeout-s`` precedence: explicit value > 900 s default."""
|
|
|
|
_OVERRIDE = "vllm_omni.benchmarks.patch.patch._REQUEST_TIMEOUT_OVERRIDE_S"
|
|
|
|
def test_default_timeout_is_900s_when_not_set(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(self._OVERRIDE, None)
|
|
assert _omni_request_timeout_s() == 900.0
|
|
|
|
def test_explicit_value_wins_over_default(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(self._OVERRIDE, 123.5)
|
|
assert _omni_request_timeout_s() == 123.5
|
|
|
|
def test_non_positive_restores_legacy_6h_cap(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(self._OVERRIDE, 0)
|
|
assert _omni_request_timeout_s() == 6 * 60 * 60.0
|
|
|
|
async def test_benchmark_session_uses_configured_timeout(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(self._OVERRIDE, 42.0)
|
|
session = _build_benchmark_session(max_concurrency=8, ssl_setting=False)
|
|
try:
|
|
assert session.timeout.total == 42.0
|
|
assert session.connector.limit == 8
|
|
assert session.connector.limit_per_host == 8
|
|
finally:
|
|
await session.close()
|
|
|
|
async def test_hung_server_request_times_out_as_failed(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A server that accepts the request but never responds must surface as ``failed``."""
|
|
|
|
async def handler(reader, writer):
|
|
# Drain the request but never respond — the original hang failure mode.
|
|
try:
|
|
while await reader.read(4096):
|
|
pass
|
|
except ConnectionResetError:
|
|
pass
|
|
finally:
|
|
# Close our side of the socket: the client aborts after the
|
|
# timeout and only sends FIN, so without writer.close() the
|
|
# half-closed connection keeps Server.wait_closed() (which on
|
|
# Python 3.12+ waits for every accepted connection to drop)
|
|
# blocked forever.
|
|
writer.close()
|
|
try:
|
|
await writer.wait_closed()
|
|
except (ConnectionResetError, BrokenPipeError):
|
|
pass
|
|
|
|
monkeypatch.setattr(self._OVERRIDE, 1.0)
|
|
server = await asyncio.start_server(handler, "127.0.0.1", 0)
|
|
port = server.sockets[0].getsockname()[1]
|
|
session = _build_benchmark_session(max_concurrency=1, ssl_setting=False)
|
|
try:
|
|
request_input = _seed_tts_request_func_input()
|
|
request_input.api_url = f"http://127.0.0.1:{port}/v1/chat/completions"
|
|
output = await async_request_openai_chat_omni_completions(request_input, session)
|
|
assert output.success is False
|
|
assert output.error
|
|
finally:
|
|
await session.close()
|
|
server.close()
|
|
await server.wait_closed()
|
|
|
|
|
|
def test_video_rtf_prefers_generation_time_over_poll_latency():
|
|
"""RTF must not grow with client poll_interval overshoot in output.latency."""
|
|
payload = {
|
|
"duration_s": 2.0,
|
|
"num_frames": 48,
|
|
"fps": 24.0,
|
|
"stage_durations": {"stage_0_gen_ms": 4000.0},
|
|
}
|
|
request_body: dict[str, object] = {}
|
|
|
|
short_poll = MixRequestFuncOutput()
|
|
short_poll.latency = 4.2 # ~generation + small poll overshoot
|
|
_apply_video_metrics_from_payload(short_poll, payload, request_body)
|
|
|
|
long_poll = MixRequestFuncOutput()
|
|
long_poll.latency = 8.0 # same job, larger poll_interval_s overshoot
|
|
_apply_video_metrics_from_payload(long_poll, payload, request_body)
|
|
|
|
assert short_poll.video_generation_time_ms == pytest.approx(4000.0)
|
|
assert long_poll.video_generation_time_ms == pytest.approx(4000.0)
|
|
assert short_poll.video_rtf == pytest.approx(2.0)
|
|
assert long_poll.video_rtf == pytest.approx(2.0)
|
|
assert short_poll.video_rtf == long_poll.video_rtf
|
|
|
|
|
|
def test_video_rtf_falls_back_to_e2e_latency_without_generation_time():
|
|
output = MixRequestFuncOutput()
|
|
output.latency = 6.0
|
|
_apply_video_metrics_from_payload(
|
|
output,
|
|
{"duration_s": 2.0, "num_frames": 48, "fps": 24.0},
|
|
{},
|
|
)
|
|
assert output.video_generation_time_ms == 0.0
|
|
assert output.video_rtf == pytest.approx(3.0)
|
|
|
|
|
|
# 1x1 PNG used as a valid b64_json image payload in edit-client tests.
|
|
_MIN_PNG_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_edits_defaults_to_non_streaming_json(mocker: MockerFixture) -> None:
|
|
"""Single-stage servers reject stream=true; default path must send stream=false + JSON."""
|
|
|
|
class MockJsonResponse:
|
|
status = 200
|
|
|
|
def __init__(self):
|
|
self._payload = {
|
|
"created": 1,
|
|
"data": [
|
|
{
|
|
"b64_json": _MIN_PNG_B64,
|
|
"stage_durations": {"stage_0_gen_ms": 12.5},
|
|
}
|
|
],
|
|
}
|
|
|
|
async def json(self):
|
|
return self._payload
|
|
|
|
async def text(self):
|
|
return json.dumps(self._payload)
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
return None
|
|
|
|
captured_stream: list[str] = []
|
|
import aiohttp
|
|
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
if name == "stream":
|
|
captured_stream.append(str(value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=MockJsonResponse())
|
|
|
|
request = RequestFuncInput(
|
|
model="single-stage-edit",
|
|
model_name="single-stage-edit",
|
|
prompt="make it sunny",
|
|
api_url="http://test.com/v1/images/edits",
|
|
prompt_len=4,
|
|
output_len=1,
|
|
multi_modal_content=[{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"}}],
|
|
extra_body={"num_inference_steps": 4},
|
|
)
|
|
output = await async_request_openai_image_edits_omni(request, mock_session, pbar=None)
|
|
|
|
assert captured_stream == ["false"]
|
|
assert output.success is True
|
|
assert not output.error
|
|
assert output.image_count == 1
|
|
assert output.image_generation_time_ms == pytest.approx(12.5)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_edits_stream_true_uses_sse_path(mocker: MockerFixture) -> None:
|
|
"""Explicit stream=true keeps the multi-stage SSE client path."""
|
|
sse_chunk = (
|
|
b'data: {"type":"image","data":[{"b64_json":"' + _MIN_PNG_B64.encode() + b'"}]}\n\n' + b"data: [DONE]\n\n"
|
|
)
|
|
mock_response = MockResponse(200, [sse_chunk])
|
|
captured_stream: list[str] = []
|
|
|
|
import aiohttp
|
|
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
if name == "stream":
|
|
captured_stream.append(str(value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
request = RequestFuncInput(
|
|
model="multi-stage-edit",
|
|
model_name="multi-stage-edit",
|
|
prompt="edit",
|
|
api_url="http://test.com/v1/images/edits",
|
|
prompt_len=2,
|
|
output_len=1,
|
|
multi_modal_content=[{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"}}],
|
|
extra_body={"stream": True},
|
|
)
|
|
output = await async_request_openai_image_edits_omni(request, mock_session, pbar=None)
|
|
|
|
assert captured_stream == ["true"]
|
|
assert output.success is True
|
|
assert output.image_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_edits_stream_error_marks_request_failed(mocker: MockerFixture) -> None:
|
|
"""HTTP 200 image-edit streams can terminate with an error event."""
|
|
sse_chunk = (
|
|
b'data: {"object":"image.edit.chunk","type":"ar_delta","delta":"partial"}\n\n'
|
|
b'data: {"object":"error","error":{"message":"image generation failed",'
|
|
b'"type":"server_error","code":500}}\n\n'
|
|
b"data: [DONE]\n\n"
|
|
)
|
|
mock_response = MockResponse(200, [sse_chunk])
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
|
|
|
request = RequestFuncInput(
|
|
model="multi-stage-edit",
|
|
model_name="multi-stage-edit",
|
|
prompt="edit",
|
|
api_url="http://test.com/v1/images/edits",
|
|
prompt_len=2,
|
|
output_len=1,
|
|
multi_modal_content=[
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"},
|
|
}
|
|
],
|
|
extra_body={"stream": True},
|
|
)
|
|
|
|
output = await async_request_openai_image_edits_omni(request, mock_session, pbar=None)
|
|
|
|
assert output.success is False
|
|
assert output.error == "image generation failed"
|
|
assert output.generated_text == "partial"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_generations_e2el_includes_json_body_consume(mocker: MockerFixture) -> None:
|
|
"""E2EL must include body transfer/decode, not stop at HTTP headers."""
|
|
|
|
class SlowJsonResponse:
|
|
status = 200
|
|
|
|
def __init__(self):
|
|
self._payload = {
|
|
"created": 1,
|
|
"data": [{"b64_json": _MIN_PNG_B64, "stage_durations": {"stage_0_gen_ms": 1.0}}],
|
|
}
|
|
|
|
async def json(self):
|
|
await asyncio.sleep(0.05)
|
|
return self._payload
|
|
|
|
async def text(self):
|
|
return json.dumps(self._payload)
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
return None
|
|
|
|
mock_session = mocker.AsyncMock()
|
|
mock_session.post = mocker.MagicMock(return_value=SlowJsonResponse())
|
|
request = RequestFuncInput(
|
|
model="img-gen",
|
|
model_name="img-gen",
|
|
prompt="a cat",
|
|
api_url="http://test.com/v1/images/generations",
|
|
prompt_len=2,
|
|
output_len=1,
|
|
extra_body={},
|
|
)
|
|
output = await async_request_openai_image_generations_omni(request, mock_session, pbar=None)
|
|
assert output.success is True
|
|
assert output.latency >= 0.05
|
|
|
|
|
|
def test_video_local_image_reference_not_forwarded_as_raw_extra_field(tmp_path, mocker: MockerFixture) -> None:
|
|
"""Local image_reference must only be uploaded via dedicated serializer."""
|
|
import aiohttp
|
|
|
|
ref_path = tmp_path / "ref.png"
|
|
ref_path.write_bytes(base64.b64decode(_MIN_PNG_B64))
|
|
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
|
|
form = aiohttp.FormData()
|
|
extra_body = {
|
|
"num_inference_steps": 2,
|
|
"image_reference": str(ref_path),
|
|
}
|
|
request_body = {"model": "vid", "prompt": "p", **extra_body}
|
|
_add_video_extra_body_to_form(form, extra_body, request_body)
|
|
assert _add_video_reference_to_form(form, extra_body["image_reference"]) is True
|
|
|
|
field_names = [name for name, _ in captured]
|
|
assert "image_reference" not in field_names
|
|
assert field_names.count("input_reference") == 1
|
|
# Dedicated uploader sends file bytes, not the raw local path string.
|
|
uploaded = next(value for name, value in captured if name == "input_reference")
|
|
assert uploaded == base64.b64decode(_MIN_PNG_B64)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"reference",
|
|
[
|
|
{"image_url": "https://example.com/ref.png"},
|
|
[{"image_url": "https://example.com/a.png"}, {"image_url": "https://example.com/b.png"}],
|
|
],
|
|
)
|
|
def test_video_structured_image_reference_serialized_to_form(reference: object, mocker: MockerFixture) -> None:
|
|
"""Object-form image_reference must be JSON-serialized, not silently dropped."""
|
|
import aiohttp
|
|
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
|
|
form = aiohttp.FormData()
|
|
extra_body = {"image_reference": reference}
|
|
request_body = {"model": "vid", "prompt": "p", **extra_body}
|
|
_add_video_extra_body_to_form(form, extra_body, request_body)
|
|
assert _add_video_reference_to_form(form, reference) is True
|
|
|
|
field_names = [name for name, _ in captured]
|
|
# Reserved key must not be double-forwarded as a generic extra field;
|
|
# exactly one dedicated image_reference field.
|
|
assert field_names.count("image_reference") == 1
|
|
assert "input_reference" not in field_names
|
|
payload = next(value for name, value in captured if name == "image_reference")
|
|
assert isinstance(payload, (str, bytes, bytearray))
|
|
assert json.loads(payload) == reference
|
|
|
|
|
|
def test_video_file_id_reference_is_rejected() -> None:
|
|
"""file_id is unsupported on the server and must not be sent as image_reference."""
|
|
import aiohttp
|
|
|
|
form = aiohttp.FormData()
|
|
with pytest.raises(ValueError, match="file_id is not supported yet"):
|
|
_add_video_reference_to_form(form, {"file_id": "file-abc"})
|
|
with pytest.raises(ValueError, match="file_id is not supported yet"):
|
|
_add_video_reference_to_form(form, [{"image_url": "https://example.com/a.png"}, {"file_id": "file-xyz"}])
|
|
with pytest.raises(ValueError, match="file_id is not supported yet"):
|
|
_add_combined_video_form_references(form, None, {"video_reference": {"file_id": "file-vid"}})
|
|
|
|
|
|
def test_image_reference_urls_from_random_mm_content(mocker: MockerFixture) -> None:
|
|
"""random-mm image_url parts keep an explicit image_reference type."""
|
|
import aiohttp
|
|
|
|
content = [
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "https://example.com/ref.png"},
|
|
},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"},
|
|
},
|
|
]
|
|
refs = list(_iter_image_reference_inputs(content))
|
|
assert refs == [
|
|
{"image_url": "https://example.com/ref.png"},
|
|
{"image_url": f"data:image/png;base64,{_MIN_PNG_B64}"},
|
|
]
|
|
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
form = aiohttp.FormData()
|
|
assert _add_video_reference_to_form(form, refs[0]) is True
|
|
assert _add_video_reference_to_form(form, refs[1]) is True
|
|
|
|
field_names = [name for name, _ in captured]
|
|
assert field_names.count("image_reference") == 2
|
|
assert "video_reference" not in field_names
|
|
payloads = []
|
|
for name, value in captured:
|
|
if name != "image_reference":
|
|
continue
|
|
assert isinstance(value, (str, bytes, bytearray))
|
|
payloads.append(json.loads(value))
|
|
assert payloads == refs
|
|
|
|
|
|
def test_video_reference_urls_from_random_mm_content(mocker: MockerFixture) -> None:
|
|
"""random-mm data:video_url parts stay on video_reference."""
|
|
import aiohttp
|
|
|
|
content = [
|
|
{
|
|
"type": "video_url",
|
|
"video_url": {"url": "data:video/mp4;base64,AAAA"},
|
|
}
|
|
]
|
|
refs = list(_iter_video_reference_inputs(content))
|
|
assert refs == [{"video_url": "data:video/mp4;base64,AAAA"}]
|
|
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
form = aiohttp.FormData()
|
|
assert _add_video_reference_to_form(form, refs[0]) is True
|
|
|
|
field_names = [name for name, _ in captured]
|
|
assert field_names.count("video_reference") == 1
|
|
assert "input_references" not in field_names
|
|
assert "input_reference" not in field_names
|
|
assert "image_reference" not in field_names
|
|
payload = next(value for name, value in captured if name == "video_reference")
|
|
assert isinstance(payload, (str, bytes, bytearray))
|
|
assert json.loads(payload) == {"video_url": "data:video/mp4;base64,AAAA"}
|
|
|
|
|
|
def test_video_reference_https_url_from_random_mm_content(mocker: MockerFixture) -> None:
|
|
"""HTTP(S) video_url parts must stay on video_reference, not image_reference."""
|
|
import aiohttp
|
|
|
|
content = [
|
|
{
|
|
"type": "video_url",
|
|
"video_url": {"url": "https://example.com/ref.mp4"},
|
|
}
|
|
]
|
|
refs = list(_iter_video_reference_inputs(content))
|
|
assert refs == [{"video_url": "https://example.com/ref.mp4"}]
|
|
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
form = aiohttp.FormData()
|
|
assert _add_video_reference_to_form(form, refs[0]) is True
|
|
|
|
field_names = [name for name, _ in captured]
|
|
assert field_names.count("video_reference") == 1
|
|
assert "image_reference" not in field_names
|
|
assert "input_references" not in field_names
|
|
payload = next(value for name, value in captured if name == "video_reference")
|
|
assert isinstance(payload, (str, bytes, bytearray))
|
|
assert json.loads(payload) == {"video_url": "https://example.com/ref.mp4"}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("reference", "field_name", "expected"),
|
|
[
|
|
("data:image/png;base64,AAAA", "image_reference", {"image_url": "data:image/png;base64,AAAA"}),
|
|
("data:video/mp4;base64,AAAA", "video_reference", {"video_url": "data:video/mp4;base64,AAAA"}),
|
|
("https://example.com/ref.png", "image_reference", {"image_url": "https://example.com/ref.png"}),
|
|
("https://example.com/ref.mp4", "video_reference", {"video_url": "https://example.com/ref.mp4"}),
|
|
],
|
|
)
|
|
def test_bare_reference_string_uses_matching_json_field(
|
|
reference: str,
|
|
field_name: str,
|
|
expected: dict[str, str],
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
"""Image and video URLs use matching JSON fields; neither is a file upload."""
|
|
import aiohttp
|
|
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
assert _add_video_reference_to_form(aiohttp.FormData(), reference) is True
|
|
assert [name for name, _ in captured] == [field_name]
|
|
payload = captured[0][1]
|
|
assert isinstance(payload, (str, bytes, bytearray))
|
|
assert json.loads(payload) == expected
|
|
|
|
|
|
def test_bare_video_string_is_not_collected_as_image() -> None:
|
|
content = ["https://example.com/ref.mp4", "data:video/mp4;base64,AAAA"]
|
|
assert list(_iter_image_reference_inputs(content)) == []
|
|
assert list(_iter_video_reference_inputs(content)) == [
|
|
{"video_url": "https://example.com/ref.mp4"},
|
|
{"video_url": "data:video/mp4;base64,AAAA"},
|
|
]
|
|
|
|
|
|
def test_bare_http_without_media_extension_is_rejected() -> None:
|
|
import aiohttp
|
|
|
|
with pytest.raises(ValueError, match="image or video extension"):
|
|
_add_video_reference_to_form(aiohttp.FormData(), "https://example.com/ref")
|
|
|
|
|
|
def test_image_and_inline_video_use_combined_reference_fields(mocker: MockerFixture) -> None:
|
|
"""Image plus data:video must be image_reference + video_reference, not input_references."""
|
|
import aiohttp
|
|
|
|
content = [
|
|
{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}},
|
|
{"type": "video_url", "video_url": {"url": "data:video/mp4;base64,AAAA"}},
|
|
]
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
form = aiohttp.FormData()
|
|
_add_combined_video_form_references(form, content)
|
|
|
|
field_names = [name for name, _ in captured]
|
|
assert field_names.count("image_reference") == 1
|
|
assert field_names.count("video_reference") == 1
|
|
assert "input_references" not in field_names
|
|
assert "input_reference" not in field_names
|
|
image_payload = next(value for name, value in captured if name == "image_reference")
|
|
video_payload = next(value for name, value in captured if name == "video_reference")
|
|
assert isinstance(image_payload, (str, bytes, bytearray))
|
|
assert isinstance(video_payload, (str, bytes, bytearray))
|
|
assert json.loads(image_payload) == {"image_url": "https://example.com/ref.png"}
|
|
assert json.loads(video_payload) == {"video_url": "data:video/mp4;base64,AAAA"}
|
|
|
|
|
|
def test_image_url_and_bare_data_video_string_use_json_fields(mocker: MockerFixture) -> None:
|
|
"""Image URL plus a bare data:video string must not upload input_references."""
|
|
import aiohttp
|
|
|
|
content = [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}]
|
|
extra_body = {"video_reference": "data:video/mp4;base64,AAAA"}
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
form = aiohttp.FormData()
|
|
_add_combined_video_form_references(form, content, extra_body)
|
|
|
|
field_names = [name for name, _ in captured]
|
|
assert field_names.count("image_reference") == 1
|
|
assert field_names.count("video_reference") == 1
|
|
assert "input_references" not in field_names
|
|
assert "input_reference" not in field_names
|
|
image_payload = next(value for name, value in captured if name == "image_reference")
|
|
video_payload = next(value for name, value in captured if name == "video_reference")
|
|
assert isinstance(image_payload, (str, bytes, bytearray))
|
|
assert isinstance(video_payload, (str, bytes, bytearray))
|
|
assert json.loads(image_payload) == {"image_url": "https://example.com/ref.png"}
|
|
assert json.loads(video_payload) == {"video_url": "data:video/mp4;base64,AAAA"}
|
|
|
|
captured.clear()
|
|
assert _add_video_reference_to_form(form, "data:video/mp4;base64,AAAA") is True
|
|
assert [name for name, _ in captured] == ["video_reference"]
|
|
|
|
|
|
def _oversized_data_video_url() -> str:
|
|
"""data:video whose JSON text part is larger than 1MB."""
|
|
return "data:video/mp4;base64," + ("A" * (1024 * 1024))
|
|
|
|
|
|
def test_oversized_data_video_uploads_as_input_references(mocker: MockerFixture) -> None:
|
|
"""A lone data:video over the 1MB text limit is uploaded, not sent as JSON."""
|
|
import aiohttp
|
|
|
|
video_url = _oversized_data_video_url()
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
assert _add_video_reference_to_form(aiohttp.FormData(), {"video_url": video_url}) is True
|
|
uploaded = next(value for name, value in captured if name == "input_references")
|
|
assert uploaded == base64.b64decode("A" * (1024 * 1024))
|
|
field_names = [name for name, _ in captured]
|
|
assert "video_reference" not in field_names
|
|
assert "image_reference" not in field_names
|
|
|
|
|
|
def test_oversized_data_video_with_image_stays_json(mocker: MockerFixture) -> None:
|
|
"""input_references cannot be combined with image_reference, even for a large video."""
|
|
import aiohttp
|
|
|
|
content = [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}]
|
|
extra_body = {"video_reference": _oversized_data_video_url()}
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
_add_combined_video_form_references(aiohttp.FormData(), content, extra_body)
|
|
field_names = [name for name, _ in captured]
|
|
assert "image_reference" in field_names
|
|
assert "video_reference" in field_names
|
|
assert "input_references" not in field_names
|
|
|
|
|
|
def test_image_upload_bytes_and_structured_video_use_json_fields(mocker: MockerFixture) -> None:
|
|
"""Image upload bytes plus a video URL must not use singular input_reference."""
|
|
import aiohttp
|
|
|
|
image_bytes = base64.b64decode(_MIN_PNG_B64)
|
|
extra_body = {
|
|
"image_reference": {"bytes": image_bytes, "content_type": "image/png"},
|
|
"video_reference": {"video_url": "https://example.com/ref.mp4"},
|
|
}
|
|
captured: list[tuple[str, object]] = []
|
|
real_add_field = aiohttp.FormData.add_field
|
|
|
|
def tracking_add_field(self, name, value=None, **kwargs):
|
|
captured.append((str(name), value))
|
|
return real_add_field(self, name, value, **kwargs)
|
|
|
|
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
|
form = aiohttp.FormData()
|
|
_add_combined_video_form_references(form, None, extra_body)
|
|
|
|
field_names = [name for name, _ in captured]
|
|
assert field_names.count("image_reference") == 1
|
|
assert field_names.count("video_reference") == 1
|
|
assert "input_reference" not in field_names
|
|
assert "input_references" not in field_names
|
|
image_payload = next(value for name, value in captured if name == "image_reference")
|
|
video_payload = next(value for name, value in captured if name == "video_reference")
|
|
assert isinstance(image_payload, (str, bytes, bytearray))
|
|
assert isinstance(video_payload, (str, bytes, bytearray))
|
|
assert json.loads(image_payload) == {"image_url": f"data:image/png;base64,{_MIN_PNG_B64}"}
|
|
assert json.loads(video_payload) == {"video_url": "https://example.com/ref.mp4"}
|
|
|
|
|
|
def test_video_unsupported_image_reference_raises() -> None:
|
|
import aiohttp
|
|
|
|
form = aiohttp.FormData()
|
|
with pytest.raises(ValueError, match="Unsupported reference"):
|
|
_add_video_reference_to_form(form, {"not_a_supported_key": "x"})
|
|
with pytest.raises(ValueError, match="Unsupported reference"):
|
|
_add_video_reference_to_form(form, "/tmp/does-not-exist-ref.png")
|
|
|
|
|
|
def test_extract_stage_durations_from_video_and_image_shapes() -> None:
|
|
video_payload = {
|
|
"stage_durations": {
|
|
"diffuse": 1.5,
|
|
"text_encoder.forward": 0.2,
|
|
"vae.decode": 0.1,
|
|
"stage_0_gen_ms": 4000.0,
|
|
}
|
|
}
|
|
assert _extract_stage_durations_from_payload(video_payload) == video_payload["stage_durations"]
|
|
|
|
image_stage_durations = {
|
|
"diffuse": 2.0,
|
|
"vae.decode": 0.3,
|
|
}
|
|
image_payload = {
|
|
"metrics": {"stage_durations": image_stage_durations},
|
|
"data": [{"b64_json": "x"}],
|
|
}
|
|
assert _extract_stage_durations_from_payload(image_payload) == image_stage_durations
|
|
|
|
|
|
def test_video_metrics_persist_full_stage_durations() -> None:
|
|
output = MixRequestFuncOutput()
|
|
output.latency = 4.2
|
|
stage_durations = {
|
|
"diffuse": 1.5,
|
|
"vae.decode": 0.25,
|
|
"stage_0_gen_ms": 4000.0,
|
|
}
|
|
payload = {
|
|
"duration_s": 2.0,
|
|
"num_frames": 48,
|
|
"fps": 24.0,
|
|
"stage_durations": stage_durations,
|
|
}
|
|
_apply_video_metrics_from_payload(output, payload, {})
|
|
assert output.stage_durations == stage_durations
|
|
assert output.video_generation_time_ms == pytest.approx(4000.0)
|
|
|
|
|
|
def test_image_metrics_persist_stage_durations_from_metrics() -> None:
|
|
output = MixRequestFuncOutput()
|
|
stage_durations = {
|
|
"diffuse": 1.1,
|
|
"text_encoder.forward": 0.4,
|
|
"vae.decode": 0.2,
|
|
}
|
|
payload = {
|
|
"created": 1,
|
|
"data": [{"b64_json": _MIN_PNG_B64}],
|
|
"metrics": {"stage_durations": stage_durations},
|
|
}
|
|
assert _apply_image_metrics_from_payload(output, payload) == 1
|
|
assert output.stage_durations == stage_durations
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v", "-s"])
|