[Refactor] Share offline profiler argument parser

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -5,11 +5,10 @@ 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 Any, NamedTuple
+from typing import NamedTuple
 
 import numpy as np
 import soundfile as sf
@@ -22,6 +21,7 @@ from vllm.multimodal.media.audio import load_audio
 from vllm.sampling_params import SamplingParams
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.profiler import add_profiler_config_arg
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
@@ -32,16 +32,6 @@ 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.
@@ -336,6 +326,7 @@ def main(args):
     else:
         query_result = query_func()
     args.quantization_config = quantization_config
+    # ``vars(args)`` below forwards profiler_config to Omni.
     omni_kwargs = vars(args).copy()
     # Override CLI --model with the derived model_name.
     omni_kwargs["model"] = model_name
@@ -469,12 +460,7 @@ 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"}\'.',
-    )
+    add_profiler_config_arg(parser)
     parser.add_argument(
         "--stage-init-timeout",
         type=int,

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -1,11 +1,9 @@
 # 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
@@ -16,21 +14,12 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
 from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
+from vllm_omni.utils.profiler import add_profiler_config_arg
 
 # Upstream zero-shot reference clip
 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")
@@ -57,12 +46,7 @@ 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"}\'.',
-    )
+    add_profiler_config_arg(parser)
     parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
     parser.add_argument(
         "--prompt-text",

tests/docs/test_offline_profiler_config.py

@@ -18,6 +18,6 @@ _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 "add_profiler_config_arg(parser)" in source
     assert "profiler_enabled = args.profiler_config is not None" in source
     assert "VLLM_TORCH_PROFILER_DIR" not in source

vllm_omni/utils/profiler.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import argparse
+import json
+from typing import Any
+
+
+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
+
+
+def parse_profiler_config(value: str) -> dict[str, Any]:
+    """Parse the JSON object passed to ``--profiler-config``."""
+    return parse_json_object(value, flag_name="--profiler-config")
+
+
+def add_profiler_config_arg(parser: argparse.ArgumentParser) -> None:
+    """Add the shared JSON profiler configuration option to ``parser``."""
+    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"}\'.'
+        ),
+    )