diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 6dbafdfc1..b17a32761 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -5,10 +5,11 @@ This example shows how to use vLLM-Omni for running offline inference with the correct prompt format on Qwen2.5-Omni """ +import argparse import json import os import time -from typing import NamedTuple +from typing import Any, NamedTuple import numpy as np import soundfile as sf @@ -31,6 +32,16 @@ class QueryResult(NamedTuple): limit_mm_per_prompt: dict[str, int] +def parse_profiler_config(value: str) -> dict[str, Any]: + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError("--profiler-config must be a JSON object") + return config + + # NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on # lower-end GPUs. # Unless specified, these settings have been tested to work on a single L4. @@ -378,11 +389,11 @@ def main(args): for i, prompt in enumerate(prompts): prompt["modalities"] = output_modalities - profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) + profiler_enabled = args.profiler_config is not None if profiler_enabled and hasattr(omni, "start_profile"): omni.start_profile(stages=[0]) elif profiler_enabled: - print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.") + print("[Warn] --profiler-config was provided, but current engine does not support profiler controls.") omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator) # Determine output directory: prefer --output-dir; fallback to --output-wav @@ -458,6 +469,12 @@ def parse_args(): default=False, help="Enable writing detailed statistics (default: disabled)", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) parser.add_argument( "--stage-init-timeout", type=int, diff --git a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py index 62917de07..42d03b66a 100644 --- a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py +++ b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import json import os import urllib.request from pathlib import Path +from typing import Any import numpy as np import soundfile as sf @@ -19,6 +21,16 @@ from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav" +def parse_profiler_config(value: str) -> dict[str, Any]: + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError("--profiler-config must be a JSON object") + return config + + def _default_ref_audio() -> str: # Download the upstream zero_shot_prompt.wav into the current dir dest = Path("zero_shot_prompt.wav") @@ -45,6 +57,12 @@ def run_e2e(): help="Override the deploy config path. If unset, auto-loads " "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.") parser.add_argument( "--prompt-text", @@ -79,6 +97,7 @@ def run_e2e(): deploy_config=args.deploy_config, tokenizer=args.tokenizer, log_stats=True, + profiler_config=args.profiler_config, ) sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1} @@ -149,8 +168,8 @@ def run_e2e(): sampling_params_list = [gpt_sampling, s2mel_sampling] - # Start profiling (requires VLLM_TORCH_PROFILER_DIR env var) - if os.environ.get("VLLM_TORCH_PROFILER_DIR"): + profiler_enabled = args.profiler_config is not None + if profiler_enabled: print("Starting profiler...") omni.start_profile() @@ -158,7 +177,7 @@ def run_e2e(): outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2])) # Stop profiling and get results - if os.environ.get("VLLM_TORCH_PROFILER_DIR"): + if profiler_enabled: print("Stopping profiler...") profile_results = omni.stop_profile() print(f"Profile traces saved to: {profile_results}") diff --git a/tests/docs/test_offline_profiler_config.py b/tests/docs/test_offline_profiler_config.py new file mode 100644 index 000000000..f36a02b1b --- /dev/null +++ b/tests/docs/test_offline_profiler_config.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_EXAMPLES = ( + "examples/offline_inference/text_to_speech/cosyvoice3/end2end.py", + "examples/offline_inference/qwen2_5_omni/end2end.py", +) + + +@pytest.mark.parametrize("relative_path", _EXAMPLES) +def test_offline_examples_use_typed_profiler_config(relative_path: str): + source = (_REPO_ROOT / relative_path).read_text(encoding="utf-8") + + assert "--profiler-config" in source + assert "profiler_enabled = args.profiler_config is not None" in source + assert "VLLM_TORCH_PROFILER_DIR" not in source