# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import argparse import json import os import subprocess import textwrap import time from pathlib import Path import pytest from vllm_omni.entrypoints.cli.benchmark.cli_args import ( add_omni_args, extend_omni_choices, preprocess_serve_args, update_omni_help, ) from vllm_omni.utils.tracking_parser import TrackingArgumentParser pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark] def test_it2i_dataset_uses_upstream_warmups(monkeypatch, tmp_path): """Exercise real dataset/benchmark code; replace only the HTTP backend.""" from vllm_omni.benchmarks import serve from vllm_omni.benchmarks.patch import patch from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand dataset_path = Path(__file__).parents[1] / "assets/hunyuan_image3/it2i.jsonl" sample = json.loads(dataset_path.read_text(encoding="utf-8")) parser = argparse.ArgumentParser() OmniBenchmarkServingSubcommand.add_cli_args(parser) args = parser.parse_args( [ "--model", "test-model", "--endpoint", "/v1/images/edits", "--dataset-name", "custom_image", "--dataset-path", str(dataset_path), "--skip-tokenizer-init", "--disable-shuffle", "--disable-tqdm", "--num-prompts", "8", "--num-warmups", "2", "--max-concurrency", "1", "--ready-check-timeout-sec", "0", "--percentile-metrics", "e2el", "--extra-body", '{"seed":42,"guidance_scale":2.5,"bot_task":"think_recaption"}', "--save-result", "--result-dir", str(tmp_path), "--result-filename", "result.json", ] ) requests = [] async def request_func(request_func_input, session, pbar=None): requests.append(request_func_input) if pbar is not None: pbar.update(1) return patch.MixRequestFuncOutput( success=True, prompt_len=1, output_tokens=1, image_count=1, start_time=time.perf_counter(), latency=99.0 if len(requests) <= 2 else 0.1, ) monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "/v1/images/edits", request_func) result = serve.main(args) assert len(requests) == 10 for request in requests: assert request.prompt == sample["prompt"] assert [image["image_url"]["url"] for image in request.multi_modal_content] == sample["image_files"] assert request.extra_body == args.extra_body assert result["completed"] == 8 saved = json.loads((tmp_path / "result.json").read_text()) assert saved["completed"] == 8 assert saved["mean_e2el_ms"] == pytest.approx(100.0) # Excludes both 99-second warmups. @pytest.mark.parametrize( "dataset_name", [ "daily-omni", "omniinteract", "seed-tts", "seed-tts-text", "seed-tts-design", "ttsd", "sound-effect", ], ) def test_extend_omni_choices_updates_tracking_parser_shadow(dataset_name: str) -> None: parser = TrackingArgumentParser() parser.add_argument("--dataset-name", choices=["random"]) parser.add_argument("--backend", choices=["openai-chat-omni"]) extend_omni_choices(parser) args = parser.parse_args( [ "--dataset-name", dataset_name, "--backend", "openai-image-edits-omni", ] ) assert args.dataset_name == dataset_name assert args.backend == "openai-image-edits-omni" assert args.explicit_keys == {"dataset_name", "backend"} @pytest.mark.parametrize( ("argv", "dest", "expected"), [ (["--print-stage"], "print_stage", True), (["--daily-omni-input-mode", "audio"], "daily_omni_input_mode", "audio"), (["--omniinteract-subsets", "1qna"], "omniinteract_subsets", ["1qna"]), (["--seed-tts-locale", "zh"], "seed_tts_locale", "zh"), ], ) def test_add_omni_args_registers_arguments_on_tracking_parser( argv: list[str], dest: str, expected: object, ) -> None: parser = TrackingArgumentParser() add_omni_args(parser) args = parser.parse_args(argv) assert getattr(args, dest) == expected assert args.explicit_keys == {dest} def test_add_omni_args_preserves_implicit_defaults() -> None: parser = TrackingArgumentParser() add_omni_args(parser) args = parser.parse_args([]) assert args.print_stage is False assert args.daily_omni_input_mode == "all" assert args.omniinteract_subsets == ["1q1a", "1q1a_math", "1qna"] assert args.seed_tts_locale == "en" assert args.explicit_keys == set() def test_update_omni_help_updates_upstream_actions() -> None: parser = TrackingArgumentParser() parser.add_argument("--num-prompts", default=1000, help="Number of prompts.") parser.add_argument("--percentile-metrics", help="Upstream percentile help.") parser.add_argument("--random-mm-limit-mm-per-prompt", help="Upstream limit help.") parser.add_argument("--random-mm-bucket-config", help="Upstream bucket help.") update_omni_help(parser) actions = {action.dest: action for action in parser._actions} assert "OmniInteract uses 3" in actions["num_prompts"].help assert all(metric in actions["percentile_metrics"].help for metric in ("ttfc", "tpop", "audio_rtf")) assert "probabilities are renormalized" in actions["random_mm_limit_mm_per_prompt"].help assert "Currently allows for 3 modalities" in actions["random_mm_bucket_config"].help @pytest.mark.parametrize( ("backend", "extra_body", "bot_task", "expected"), [ ("openai-chat-omni", None, "think", {}), ("openai-image-edits-omni", None, "think", {"bot_task": "think"}), ("openai-image-edits-omni", {"bot_task": "recaption"}, "think", {"bot_task": "recaption"}), ("openai-image-edits-omni", {"custom": True}, None, {"custom": True}), ], ) def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body( backend: str, extra_body: dict | None, bot_task: str | None, expected: dict, ) -> None: args = argparse.Namespace(backend=backend, extra_body=extra_body, bot_task=bot_task) preprocess_serve_args(args) assert args.extra_body == expected @pytest.mark.parametrize( ("argv", "expected"), [ ([], 3), (["--num-prompts", "0"], 0), (["--num-prompts", "12"], 12), ], ) def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default( argv: list[str], expected: int, ) -> None: parser = TrackingArgumentParser() parser.add_argument("--dataset-name", default="omniinteract") parser.add_argument("--backend", default="openai-realtime-duplex") parser.add_argument("--endpoint", default="/v1/realtime") parser.add_argument("--omniinteract-ref-audio", default="reference.wav") parser.add_argument("--num-prompts", type=int, default=1000) args = parser.parse_args(argv) preprocess_serve_args(args) assert args.num_prompts == expected @pytest.mark.parametrize( ("extra", "match"), [ (["--seed", "7"], "cannot be combined with --seed"), (["--dataset-path", "/tmp/omniinteract-data"], "cannot be combined with --dataset-path"), (["--omniinteract-scenario-tags", "realtime"], "cannot be combined with --omniinteract-scenario-tags"), (["--omniinteract-scenario-focus"], "cannot be combined with --omniinteract-scenario-focus"), ], ) def test_preprocess_serve_args_rejects_omniinteract_video_list_conflicts( tmp_path: Path, extra: list[str], match: str, ) -> None: ref = tmp_path / "ref.wav" ref.touch() video_list = tmp_path / "list.jsonl" video_list.write_text("{}\n") parser = TrackingArgumentParser() parser.add_argument("--dataset-name", default="sharegpt") parser.add_argument("--backend", default="vllm") parser.add_argument("--endpoint", default="/v1/completions") parser.add_argument("--model", default="dummy") parser.add_argument("--seed", type=int, default=0) parser.add_argument("--dataset-path", default=None) parser.add_argument("--num-prompts", type=int, default=1000) parser.add_argument("--max-concurrency", type=int, default=None) add_omni_args(parser) args = parser.parse_args( [ "--dataset-name", "omniinteract", "--backend", "openai-realtime-duplex", "--endpoint", "/v1/realtime", "--omniinteract-ref-audio", str(ref), "--omniinteract-video-list", str(video_list), *extra, ] ) with pytest.raises(ValueError, match=match): preprocess_serve_args(args) @pytest.mark.parametrize( ("argv", "expected_extra_body", "expected_explicit"), [ ( [ "--backend", "openai-image-edits-omni", "--print-stage", "--image-edits-bot-task", "recaption", ], {"bot_task": "recaption"}, {"backend", "print_stage", "bot_task"}, ), ( [ "--endpoint", "/v1/images/edits", "--image-edits-bot-task", "think", ], {"bot_task": "think"}, {"endpoint", "bot_task"}, ), ( [ "--backend", "/v1/images/edits", "--image-edits-bot-task", "recaption", ], {"bot_task": "recaption"}, {"backend", "bot_task"}, ), ( ["--extra-body", '{"bot_task":"vanilla"}'], {"bot_task": "vanilla"}, {"extra_body"}, ), ], ) def test_omni_args_parse_and_preprocess( argv: list[str], expected_extra_body: dict, expected_explicit: set[str], ) -> None: parser = TrackingArgumentParser() parser.add_argument("--extra-body", type=json.loads, default=None) parser.add_argument("--backend", default="openai-chat-omni") parser.add_argument("--endpoint", default="/v1/chat/completions") add_omni_args(parser) args = parser.parse_args(argv) preprocess_serve_args(args) assert args.extra_body == expected_extra_body assert args.explicit_keys == expected_explicit def test_bench_serve_cli_mocks_http_request(tmp_path: Path): num_prompts = 5 port = 18000 result_filename = "bench-result.json" calls_filename = "http-post-calls.json" result_path = tmp_path / result_filename calls_path = tmp_path / calls_filename sitecustomize_path = tmp_path / "sitecustomize.py" sitecustomize_path.write_text( textwrap.dedent( """ import atexit import json import os POST_CALLS = [] SSE_CHUNKS = [ b'data: {"choices":[{"delta":{"content":"hi"}}],"modality":"text"}\\n\\n', b'data: {"choices":[{"delta":{"content":" there"}}],"modality":"text","metrics":{"num_tokens_out":4,"num_tokens_in":5}}\\n\\n', b"data: [DONE]\\n\\n", ] class _Content: def __init__(self, chunks): self._chunks = chunks async def iter_any(self): for chunk in self._chunks: yield chunk class MockResponse: def __init__(self, status=200, reason="OK", chunks=None): self.status = status self.reason = reason self.content = _Content(chunks or SSE_CHUNKS) self._json_data = None self._text_data = None async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): return False async def text(self): if self._text_data is not None: return self._text_data return json.dumps({"tokens": [1, 2, 3, 4, 5]}) async def json(self): if self._json_data is not None: return self._json_data return {"tokens": [1, 2, 3, 4, 5]} def raise_for_status(self): pass class MockClientSession: def __init__(self, *args, **kwargs): pass async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): return False def post(self, url=None, *args, **kwargs): if url is not None: POST_CALLS.append(url) return MockResponse(status=200) async def get(self, url=None, *args, **kwargs): if url is not None: POST_CALLS.append(url) return MockResponse(status=200) async def close(self): return None class MockTokenizer: vocab_size = 256 all_special_ids = [] def num_special_tokens_to_add(self, *args, **kwargs): return 0 def decode(self, token_ids, *args, **kwargs): return " ".join(str(token_id) for token_id in token_ids) def encode(self, text, *args, **kwargs): return [int(token_id) for token_id in text.split()] # Patch globally so modules loaded after this also see the mock import aiohttp aiohttp.ClientSession = MockClientSession aiohttp.TCPConnector = lambda *args, **kwargs: object() # Also patch in the patch module import vllm_omni.benchmarks.patch.patch as patch_mod patch_mod.aiohttp.ClientSession = MockClientSession patch_mod.aiohttp.TCPConnector = lambda *args, **kwargs: object() # This test exercises CLI and HTTP routing, not model tokenization. # Keep its subprocess independent of hub access and model caches. import vllm.benchmarks.serve as benchmark_serve benchmark_serve.get_tokenizer = lambda *args, **kwargs: MockTokenizer() calls_file = os.environ.get("VLLM_OMNI_TEST_POST_CALLS_FILE") if calls_file: def _write_calls(): with open(calls_file, "w", encoding="utf-8") as f: json.dump({"requested_urls": POST_CALLS}, f) atexit.register(_write_calls) """ ), encoding="utf-8", ) env = os.environ.copy() env["PYTHONPATH"] = str(tmp_path) + os.pathsep + env.get("PYTHONPATH", "") env["VLLM_OMNI_TEST_POST_CALLS_FILE"] = str(calls_path) env["HF_HOME"] = str(tmp_path / "hf-cache") env["HF_HUB_OFFLINE"] = "1" env["TRANSFORMERS_OFFLINE"] = "1" cmd = [ "vllm", "bench", "serve", "--omni", "--model", "Qwen/Qwen2.5-Omni-7B", "--port", str(port), "--dataset-name", "random", "--random-input-len", "32", "--random-output-len", "4", "--num-prompts", str(num_prompts), "--endpoint", "/v1/chat/completions", "--backend", "openai-chat-omni", "--disable-tqdm", "--num-warmups", "0", "--ready-check-timeout-sec", "0", "--save-result", "--result-dir", str(tmp_path), "--result-filename", result_filename, ] proc = subprocess.run( cmd, cwd=str(Path(__file__).resolve().parents[2]), env=env, capture_output=True, text=True, ) assert proc.returncode == 0, f"CLI failed: stdout={proc.stdout}\nstderr={proc.stderr}" print(f"CLI output: {proc.stdout}") assert result_path.exists() result = json.loads(result_path.read_text(encoding="utf-8")) assert calls_path.exists() calls = json.loads(calls_path.read_text(encoding="utf-8")) expected_url = f"http://127.0.0.1:{port}/v1/chat/completions" requested_urls = calls["requested_urls"] # Count only benchmark requests (to the chat completions endpoint), # excluding tokenize/detokenize alignment calls from the upstream. bench_requests = [url for url in requested_urls if url == expected_url] sent_bench_requests = len(bench_requests) assert result["completed"] == sent_bench_requests == num_prompts, ( f"completed={result['completed']}, bench_requests={sent_bench_requests}, all_requested_urls={requested_urls}" ) assert bench_requests assert all(url == expected_url for url in bench_requests), f"Unexpected target URLs: {bench_requests}" def test_omni_request_timeout_s_flag_defaults_and_parses() -> None: parser = TrackingArgumentParser() add_omni_args(parser) args = parser.parse_args([]) assert args.omni_request_timeout_s is None args = parser.parse_args(["--omni-request-timeout-s", "0"]) assert args.omni_request_timeout_s == 0.0 args = parser.parse_args(["--omni-request-timeout-s", "3600"]) assert args.omni_request_timeout_s == 3600.0