diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 6dbafdfc1..bbeef32b0 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -5,10 +5,12 @@ This example shows how to use vLLM-Omni for running offline inference with the correct prompt format on Qwen2.5-Omni """ +import argparse +import functools import json import os import time -from typing import NamedTuple +from typing import Any, NamedTuple import numpy as np import soundfile as sf @@ -26,6 +28,20 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser SEED = 42 +def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]: + """Parse a CLI value as a JSON object, attributing errors to ``flag_name``.""" + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object") + return config + + +parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config") + + class QueryResult(NamedTuple): inputs: dict limit_mm_per_prompt: dict[str, int] @@ -378,11 +394,10 @@ def main(args): for i, prompt in enumerate(prompts): prompt["modalities"] = output_modalities - profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) - if profiler_enabled and hasattr(omni, "start_profile"): + profiler_enabled = args.profiler_config is not None + if profiler_enabled: + print("[Profiler] Starting profiling...") omni.start_profile(stages=[0]) - elif profiler_enabled: - print("[Warn] VLLM_TORCH_PROFILER_DIR is set, 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 @@ -418,7 +433,7 @@ def main(args): print(f"Request ID: {request_id}, Saved audio to {output_wav}") processed_count += 1 - if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests: + if profiler_enabled and processed_count >= total_requests: print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...") # Stop the profiler while workers are still alive omni.stop_profile() @@ -476,6 +491,12 @@ def parse_args(): default=300, help="Timeout for initializing stages in seconds (default: 300)", ) + 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( "--shm-threshold-bytes", 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..40e194120 100644 --- a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py +++ b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import functools +import json import os import urllib.request from pathlib import Path +from typing import Any import numpy as np import soundfile as sf @@ -29,6 +32,20 @@ def _default_ref_audio() -> str: return str(dest) +def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]: + """Parse a CLI value as a JSON object, attributing errors to ``flag_name``.""" + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object") + return config + + +parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config") + + def run_e2e(): parser = argparse.ArgumentParser() # ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512 @@ -64,6 +81,12 @@ def run_e2e(): required=True, help="Path to tokenizer directory (e.g., /CosyVoice-BlankEN).", ) + 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"}\'.', + ) args = parser.parse_args() # Ensure tokenizer directory exists if not os.path.exists(args.tokenizer): @@ -79,6 +102,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 +173,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 +182,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}")