46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import contextlib
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from vllm.benchmarks.serve import main_async
|
|
|
|
# Import patch to register daily-omni dataset and omni backends
|
|
# This monkey-patches vllm.benchmarks.datasets.get_samples before it's used
|
|
# Must be imported before any vllm.benchmarks module usage
|
|
import vllm_omni.benchmarks.patch.patch # noqa: F401
|
|
from vllm_omni.benchmarks.omniinteract import omniinteract_output_lock
|
|
from vllm_omni.benchmarks.patch.patch import (
|
|
maybe_enable_stage_metrics,
|
|
set_print_stage,
|
|
should_request_stage_metrics,
|
|
)
|
|
|
|
|
|
def main(args: argparse.Namespace) -> dict[str, Any]:
|
|
if getattr(args, "seed_tts_wer_eval", False):
|
|
os.environ["SEED_TTS_WER_EVAL"] = "1"
|
|
if getattr(args, "seed_tts_wer_save_items", False):
|
|
os.environ["SEED_TTS_WER_SAVE_ITEMS"] = "1"
|
|
if getattr(args, "daily_omni_save_eval_items", False):
|
|
os.environ["DAILY_OMNI_SAVE_EVAL_ITEMS"] = "1"
|
|
set_print_stage(getattr(args, "print_stage", False))
|
|
args.extra_body = maybe_enable_stage_metrics(
|
|
getattr(args, "extra_body", None),
|
|
enabled=should_request_stage_metrics(args),
|
|
)
|
|
lock = (
|
|
omniinteract_output_lock(Path(args.omniinteract_output_dir))
|
|
if getattr(args, "dataset_name", None) == "omniinteract"
|
|
else contextlib.nullcontext()
|
|
)
|
|
with lock:
|
|
return asyncio.run(main_async(args))
|