[Refactor] Split vllm_omni/entrypoints/cli/benchmark/serve.py (#5206)

tests/benchmarks/test_serve_cli.py

@@ -1,3 +1,4 @@
+import argparse
 import json
 import os
 import subprocess
@@ -6,10 +7,148 @@ 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]
+
+
+@pytest.mark.parametrize(
+    "dataset_name",
+    [
+        "daily-omni",
+        "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"),
+        (["--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.seed_tts_locale == "en"
+    assert args.explicit_keys == set()
+
+
+def test_update_omni_help_updates_upstream_actions() -> None:
+    parser = TrackingArgumentParser()
+    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 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(
+    ("extra_body", "bot_task", "expected"),
+    [
+        (None, "think", {"bot_task": "think"}),
+        ({"bot_task": "recaption"}, "think", {"bot_task": "recaption"}),
+        ({"custom": True}, None, {"custom": True}),
+    ],
+)
+def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body(
+    extra_body: dict | None,
+    bot_task: str | None,
+    expected: dict,
+) -> None:
+    args = argparse.Namespace(extra_body=extra_body, bot_task=bot_task)
+
+    preprocess_serve_args(args)
+
+    assert args.extra_body == expected
+
+
+@pytest.mark.parametrize(
+    ("argv", "expected_extra_body", "expected_explicit"),
+    [
+        (
+            ["--print-stage", "--image-edits-bot-task", "recaption"],
+            {"bot_task": "recaption"},
+            {"print_stage", "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)
+    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
+
 
-@pytest.mark.core_model
-@pytest.mark.benchmark
-@pytest.mark.cpu
 def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
     num_prompts = 5
     port = 18000

vllm_omni/entrypoints/cli/benchmark/base.py

@@ -1,3 +1,11 @@
+"""Shared interface for vLLM-Omni benchmark CLI subcommands.
+
+``OmniBenchmarkSubcommandBase`` defines the argument-registration and command
+execution hooks required by benchmark implementations. The ``bench`` command
+discovers its concrete subclasses and uses these hooks to build and dispatch
+nested benchmark subcommands.
+"""
+
 import argparse
 
 from vllm.entrypoints.cli.types import CLISubcommand

vllm_omni/entrypoints/cli/benchmark/cli_args.py

@@ -0,0 +1,244 @@
+"""vLLM-Omni extensions for the ``vllm bench serve`` CLI.
+
+Core functions:
+    add_omni_args: Register all Omni-specific argument groups by calling the
+        feature-specific ``add_*_cli_args`` helpers in this module.
+    extend_omni_choices: Add Omni datasets and backends to choices defined by
+        the upstream vLLM parser, including its shadow parser.
+    update_omni_help: Extend upstream help text with Omni-specific behavior.
+    preprocess_serve_args: Apply transformations that require parsed values
+        before the serving benchmark starts.
+
+``OmniBenchmarkServingSubcommand.add_cli_args`` invokes the first three after
+upstream vLLM registers its arguments. New feature arguments should normally be
+added to the corresponding ``add_*_cli_args`` helper, or to a new helper called
+by ``add_omni_args``.
+"""
+
+import argparse
+
+
+def add_multi_stage_cli_args(parser: argparse.ArgumentParser) -> None:
+    """Add CLI arguments for vLLM-Omni multi-stage benchmarks."""
+    group = parser.add_argument_group("vLLM-Omni Multi-stage Benchmark Options")
+    group.add_argument(
+        "--print-stage",
+        action="store_true",
+        default=False,
+        help=(
+            "Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. "
+            "Disabled by default. The latency sections follow --percentile-metrics by modality: "
+            "ttft/tpot/itl control text stages, ttfc/tpoc/icl control internal stream stages, "
+            "and tpop controls both text TPOP and internal stream TPOP."
+        ),
+    )
+
+
+def add_diffusion_cli_args(parser: argparse.ArgumentParser) -> None:
+    """Add CLI arguments for diffusion model benchmarks."""
+    group = parser.add_argument_group("Diffusion Models Options")
+    group.add_argument(
+        "--image-edits-bot-task",
+        dest="bot_task",
+        type=str,
+        default="think",
+        help=(
+            "Default bot_task form field for --backend openai-image-edits-omni "
+            "(/v1/images/edits). "
+            'Use --extra-body \'{"bot_task":"..."}\' to override per run.'
+        ),
+    )
+
+
+def add_daily_omni_cli_args(parser: argparse.ArgumentParser) -> None:
+    """Add CLI arguments specific to the Daily-Omni dataset."""
+    group = parser.add_argument_group("Daily-Omni Dataset Options")
+    group.add_argument(
+        "--daily-omni-qa-json",
+        type=str,
+        default=None,
+        help="Path to local upstream qa.json. When set, QA rows are read from this file and "
+        "the HuggingFace dataset is not loaded (no network). Use with --daily-omni-video-dir "
+        "for fully offline runs. --dataset-path / Hub split flags are then ignored for QA loading.",
+    )
+    group.add_argument(
+        "--daily-omni-video-dir",
+        type=str,
+        default=None,
+        help="Root directory of extracted Daily-Omni videos (contents of Videos.tar: "
+        "each video_id in its own subdir with {video_id}_video.mp4). "
+        "If omitted, Videos.tar is downloaded from the Hugging Face dataset repo on first multimodal "
+        "request. "
+        "When using file URLs, you MUST start the vLLM server with "
+        "--allowed-local-media-path set to this same directory (or a parent), "
+        "otherwise requests fail with 'Cannot load local files without "
+        "--allowed-local-media-path'.",
+    )
+    group.add_argument(
+        "--daily-omni-inline-local-video",
+        action="store_true",
+        default=False,
+        help="For local videos only: embed MP4 as base64 data URLs in benchmark "
+        "requests so the server does not need --allowed-local-media-path. "
+        "Increases request size and client memory; use for small --num-prompts. "
+        "When using --daily-omni-input-mode audio or all, local WAV files are "
+        "embedded the same way.",
+    )
+    group.add_argument(
+        "--daily-omni-input-mode",
+        type=str,
+        choices=["all", "visual", "audio"],
+        default="all",
+        help="Daily-Omni input protocol (mirrors upstream Lliar-liar/Daily-Omni "
+        "--input_mode). 'visual': video only (default). 'audio': WAV only, "
+        "requires {video_id}/{video_id}_audio.wav under --daily-omni-video-dir. "
+        "'all': video + WAV together. Sets mm_processor_kwargs.use_audio_in_video=false "
+        "and matches official separate video/audio streams.",
+    )
+    group.add_argument(
+        "--daily-omni-save-eval-items",
+        action="store_true",
+        default=False,
+        help="Include per-request Daily-Omni accuracy rows (gold/predicted/correct) "
+        "in the saved JSON under key daily_omni_eval_items. "
+        "Alternatively set env DAILY_OMNI_SAVE_EVAL_ITEMS=1.",
+    )
+
+
+def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None:
+    """Add CLI arguments for Seed-TTS benchmarks."""
+    group = parser.add_argument_group("Seed-TTS Dataset Options")
+    group.add_argument(
+        "--seed-tts-locale",
+        type=str,
+        choices=["en", "zh"],
+        default="en",
+        help="Which Seed-TTS split to load: en/meta.lst or zh/meta.lst under the dataset root.",
+    )
+    group.add_argument(
+        "--seed-tts-root",
+        type=str,
+        default=None,
+        help="Override root directory that contains en/ and zh/ (meta.lst + prompt-wavs). "
+        "If set, --dataset-path can still name the HF repo for logging; this path is used for files.",
+    )
+    group.add_argument(
+        "--seed-tts-file-ref-audio",
+        action="store_true",
+        default=False,
+        help="Send ref_audio as file:// URIs (smaller HTTP bodies). Requires the API server "
+        "to be started with --allowed-local-media-path covering the Seed-TTS dataset root. "
+        "Default is inline data:audio/wav;base64 so Qwen3-TTS works without that flag.",
+    )
+    group.add_argument(
+        "--seed-tts-inline-ref-audio",
+        action="store_true",
+        default=False,
+        help=argparse.SUPPRESS,
+    )
+    group.add_argument(
+        "--seed-tts-system-prompt",
+        type=str,
+        default=None,
+        help="Override chat system message for --backend openai-chat-omni (Qwen3-Omni TTS). "
+        "Default follows official Qwen3-Omni identity + zero-shot voice-clone instructions.",
+    )
+    group.add_argument(
+        "--seed-tts-wer-eval",
+        action="store_true",
+        default=False,
+        help="Keep synthesized audio as 24 kHz mono PCM for WER (works with "
+        "--backend openai-audio-speech or openai-chat-omni). Scoring follows "
+        "zhaochenyang20/seed-tts-eval (Whisper-large-v3 / Paraformer-zh + jiwer). "
+        "Sets SEED_TTS_WER_EVAL=1. Install: pip install 'vllm-omni[dev]'. "
+        "Optional: SEED_TTS_EVAL_DEVICE, SEED_TTS_HF_WHISPER_MODEL.",
+    )
+    group.add_argument(
+        "--seed-tts-wer-save-items",
+        action="store_true",
+        default=False,
+        help="Include per-utterance ASR rows in the saved JSON under key seed_tts_wer_eval_items. "
+        "Or set SEED_TTS_WER_SAVE_ITEMS=1.",
+    )
+
+
+_OMNI_BENCH_DATASET_CHOICES = (
+    "daily-omni",
+    "seed-tts",
+    "seed-tts-text",
+    "seed-tts-design",
+    "ttsd",
+    "sound-effect",
+)
+
+
+def extend_omni_choices(parser: argparse.ArgumentParser) -> None:
+    """Extend upstream argument choices with Omni-specific values."""
+    parsers = [parser]
+    shadow = getattr(parser, "_shadow", None)
+    if shadow is not None:
+        parsers.append(shadow)
+
+    for current_parser in parsers:
+        for action in current_parser._actions:
+            if action.dest == "dataset_name" and action.choices is not None:
+                extra = [choice for choice in _OMNI_BENCH_DATASET_CHOICES if choice not in action.choices]
+                if extra:
+                    action.choices = list(action.choices) + extra
+            if action.dest == "backend" and action.choices is not None:
+                extra = [choice for choice in ("openai-image-edits-omni",) if choice not in action.choices]
+                if extra:
+                    action.choices = list(action.choices) + extra
+
+
+def update_omni_help(parser: argparse.ArgumentParser) -> None:
+    """Update upstream argument help text to describe Omni-specific behavior."""
+    for action in parser._actions:
+        if action.dest == "percentile_metrics":
+            action.help = (
+                "Comma-separated list of selected metrics to report percentiles. "
+                'For text metrics, "ttft", "tpot", and "itl" affect the global benchmark and text '
+                'stage metrics. "tpop" also requests text TPOT/TPOP globally and per stage, and internal '
+                'stream TPOP. "ttfc", "tpoc", and "icl" only affect internal stream stage metrics. '
+                'Audio metrics include "audio_ttfp", "audio_rtf", "audio_duration", and "audio_underrun".'
+            )
+        if action.dest == "random_mm_limit_mm_per_prompt":
+            action.help = (
+                "Per-modality hard caps for items attached per request, e.g. "
+                '\'{"image": 3, "video": 0, "audio": 1}\'. The sampled per-request item '
+                "count is clamped to the sum of these limits. When a modality "
+                "reaches its cap, its buckets are excluded and probabilities are "
+                "renormalized."
+            )
+        if action.dest == "random_mm_bucket_config":
+            action.help = (
+                "The bucket config is a dictionary mapping a multimodal item"
+                "sampling configuration to a probability."
+                "Currently allows for 3 modalities: audio, images and videos. "
+                "A bucket key is a tuple of (height, width, num_frames)"
+                "The value is the probability of sampling that specific item. "
+                "Example: "
+                "--random-mm-bucket-config "
+                "{(256, 256, 1): 0.5, (720, 1280, 16): 0.4, (0, 1, 5): 0.10} "
+                "First item: images with resolution 256x256 w.p. 0.5"
+                "Second item: videos with resolution 720x1280 and 16 frames "
+                "Third item: audios with 1s duration and 5 channels w.p. 0.1"
+                "OBS.: If the probabilities do not sum to 1, they are normalized."
+            )
+
+
+def add_omni_args(parser: argparse.ArgumentParser) -> None:
+    """Register all vLLM-Omni serving benchmark arguments."""
+    add_daily_omni_cli_args(parser)
+    add_seed_tts_cli_args(parser)
+    add_multi_stage_cli_args(parser)
+    add_diffusion_cli_args(parser)
+
+
+def preprocess_serve_args(args: argparse.Namespace) -> None:
+    """Apply serving benchmark CLI transformations after parsing."""
+    extra_body = dict(getattr(args, "extra_body", None) or {})
+    bot_task = getattr(args, "bot_task", None)
+    if bot_task is not None:
+        extra_body.setdefault("bot_task", bot_task)
+    args.extra_body = extra_body

vllm_omni/entrypoints/cli/benchmark/main.py

@@ -1,3 +1,11 @@
+"""Top-level ``vllm bench`` command integration for vLLM-Omni.
+
+``OmniBenchmarkSubcommand`` creates the benchmark command group, discovers
+``OmniBenchmarkSubcommandBase`` implementations, builds their nested parsers,
+and dispatches parsed commands to the selected implementation. ``cmd_init``
+exposes the command to the surrounding vLLM-Omni CLI registration mechanism.
+"""
+
 from __future__ import annotations
 
 import argparse

vllm_omni/entrypoints/cli/benchmark/serve.py

@@ -1,190 +1,25 @@
+"""Online serving benchmark subcommand for vLLM-Omni.
+
+``OmniBenchmarkServingSubcommand`` starts from vLLM's serving benchmark
+arguments, adds Omni-specific options and adaptations, then preprocesses the
+parsed namespace and delegates execution to ``vllm_omni.benchmarks.serve``.
+This keeps CLI integration separate from the multimodal benchmark runtime.
+"""
+
 import argparse
-import os
 
 from vllm.benchmarks.serve import add_cli_args
 
 from vllm_omni.benchmarks.serve import main
 from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase
-
-
-def add_daily_omni_cli_args(parser: argparse.ArgumentParser) -> None:
-    """Add CLI arguments specific to Daily-Omni dataset.
-
-    This function should be called by the CLI entrypoint to add additional
-    arguments for daily-omni benchmark support.
-
-    Args:
-        parser: The ArgumentParser instance to extend
-    """
-    # Daily-Omni specific arguments
-    daily_omni_group = parser.add_argument_group("Daily-Omni Dataset Options")
-
-    daily_omni_group.add_argument(
-        "--daily-omni-qa-json",
-        type=str,
-        default=None,
-        help="Path to local upstream qa.json. When set, QA rows are read from this file and "
-        "the HuggingFace dataset is not loaded (no network). Use with --daily-omni-video-dir "
-        "for fully offline runs. --dataset-path / Hub split flags are then ignored for QA loading.",
-    )
-    daily_omni_group.add_argument(
-        "--daily-omni-video-dir",
-        type=str,
-        default=None,
-        help="Root directory of extracted Daily-Omni videos (contents of Videos.tar: "
-        "each video_id in its own subdir with {video_id}_video.mp4). "
-        "If omitted, Videos.tar is downloaded from the Hugging Face dataset repo on first multimodal "
-        "request. "
-        "When using file URLs, you MUST start the vLLM server with "
-        "--allowed-local-media-path set to this same directory (or a parent), "
-        "otherwise requests fail with 'Cannot load local files without "
-        "--allowed-local-media-path'.",
-    )
-    daily_omni_group.add_argument(
-        "--daily-omni-inline-local-video",
-        action="store_true",
-        default=False,
-        help="For local videos only: embed MP4 as base64 data URLs in benchmark "
-        "requests so the server does not need --allowed-local-media-path. "
-        "Increases request size and client memory; use for small --num-prompts. "
-        "When using --daily-omni-input-mode audio or all, local WAV files are "
-        "embedded the same way.",
-    )
-    daily_omni_group.add_argument(
-        "--daily-omni-input-mode",
-        type=str,
-        choices=["all", "visual", "audio"],
-        default="all",
-        help="Daily-Omni input protocol (mirrors upstream Lliar-liar/Daily-Omni "
-        "--input_mode). 'visual': video only (default). 'audio': WAV only, "
-        "requires {video_id}/{video_id}_audio.wav under --daily-omni-video-dir. "
-        "'all': video + WAV together. Sets mm_processor_kwargs.use_audio_in_video=false "
-        "and matches official separate video/audio streams.",
-    )
-    daily_omni_group.add_argument(
-        "--daily-omni-save-eval-items",
-        action="store_true",
-        default=False,
-        help="Include per-request Daily-Omni accuracy rows (gold/predicted/correct) "
-        "in the saved JSON under key daily_omni_eval_items. "
-        "Alternatively set env DAILY_OMNI_SAVE_EVAL_ITEMS=1.",
-    )
-
-    # Note: --dataset-name daily-omni via get_samples patch; use either Hub (--dataset-path
-    # liarliar/Daily-Omni) or local --daily-omni-qa-json (offline).
-
-
-def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None:
-    """CLI for Seed-TTS zero-shot TTS benchmark (``--dataset-name seed-tts``)."""
-    g = parser.add_argument_group("Seed-TTS Dataset Options")
-    g.add_argument(
-        "--seed-tts-locale",
-        type=str,
-        choices=["en", "zh"],
-        default="en",
-        help="Which Seed-TTS split to load: en/meta.lst or zh/meta.lst under the dataset root.",
-    )
-    g.add_argument(
-        "--seed-tts-root",
-        type=str,
-        default=None,
-        help="Override root directory that contains en/ and zh/ (meta.lst + prompt-wavs). "
-        "If set, --dataset-path can still name the HF repo for logging; this path is used for files.",
-    )
-    g.add_argument(
-        "--seed-tts-file-ref-audio",
-        action="store_true",
-        default=False,
-        help="Send ref_audio as file:// URIs (smaller HTTP bodies). Requires the API server "
-        "to be started with --allowed-local-media-path covering the Seed-TTS dataset root. "
-        "Default is inline data:audio/wav;base64 so Qwen3-TTS works without that flag.",
-    )
-    g.add_argument(
-        "--seed-tts-inline-ref-audio",
-        action="store_true",
-        default=False,
-        help=argparse.SUPPRESS,
-    )
-    g.add_argument(
-        "--seed-tts-system-prompt",
-        type=str,
-        default=None,
-        help="Override chat system message for --backend openai-chat-omni (Qwen3-Omni TTS). "
-        "Default follows official Qwen3-Omni identity + zero-shot voice-clone instructions.",
-    )
-    g.add_argument(
-        "--seed-tts-wer-eval",
-        action="store_true",
-        default=False,
-        help="Keep synthesized audio as 24 kHz mono PCM for WER (works with "
-        "--backend openai-audio-speech or openai-chat-omni). Scoring follows "
-        "zhaochenyang20/seed-tts-eval (Whisper-large-v3 / Paraformer-zh + jiwer). "
-        "Sets SEED_TTS_WER_EVAL=1. Install: pip install 'vllm-omni[dev]'. "
-        "Optional: SEED_TTS_EVAL_DEVICE, SEED_TTS_HF_WHISPER_MODEL.",
-    )
-    g.add_argument(
-        "--seed-tts-wer-save-items",
-        action="store_true",
-        default=False,
-        help="Include per-utterance ASR rows in the saved JSON under key seed_tts_wer_eval_items. "
-        "Or set SEED_TTS_WER_SAVE_ITEMS=1.",
-    )
-
-
-def add_omni_benchmark_cli_args(parser: argparse.ArgumentParser) -> None:
-    """Add vLLM-Omni specific serving benchmark options."""
-    group = parser.add_argument_group("vLLM-Omni Multi-stage Benchmark Options")
-    group.add_argument(
-        "--print-stage",
-        action="store_true",
-        default=False,
-        help=(
-            "Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. "
-            "Disabled by default. The latency sections follow --percentile-metrics by modality: "
-            "ttft/tpot/itl control text stages, ttfc/tpoc/icl control internal stream stages, "
-            "and tpop controls both text TPOP and internal stream TPOP."
-        ),
-    )
-    group.add_argument(
-        "--image-edits-bot-task",
-        type=str,
-        default="think",
-        help=(
-            "Default bot_task form field for --backend openai-image-edits-omni "
-            '(/v1/images/edits). Use --extra-body \'{"bot_task":"..."}\' to override per run.'
-        ),
-    )
-
-
-_OMNI_BENCH_DATASET_CHOICES = (
-    "daily-omni",
-    "seed-tts",
-    "seed-tts-text",
-    "seed-tts-design",
-    "ttsd",
-    "sound-effect",
+from vllm_omni.entrypoints.cli.benchmark.cli_args import (
+    add_omni_args,
+    extend_omni_choices,
+    preprocess_serve_args,
+    update_omni_help,
 )
 
 
-def _extend_omni_dataset_name_choices(parser: argparse.ArgumentParser) -> None:
-    """Append omni benchmark dataset names to --dataset-name choices.
-
-    TrackingArgumentParser keeps a shadow parser for explicit-arg tracking; both
-    parsers must list the same choices or parse_args rejects valid omni values.
-    """
-    parsers = [parser]
-    shadow = getattr(parser, "_shadow", None)
-    if shadow is not None:
-        parsers.append(shadow)
-
-    for p in parsers:
-        for action in p._actions:
-            if action.dest == "dataset_name" and action.choices is not None:
-                extra = [c for c in _OMNI_BENCH_DATASET_CHOICES if c not in action.choices]
-                if extra:
-                    action.choices = list(action.choices) + extra
-
-
 class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
     """The `serve` subcommand for vllm bench."""
 
@@ -194,79 +29,11 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
     @classmethod
     def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
         add_cli_args(parser)
-
-        # Add Daily-Omni specific arguments
-        add_daily_omni_cli_args(parser)
-        add_seed_tts_cli_args(parser)
-        add_omni_benchmark_cli_args(parser)
-
-        for action in parser._actions:
-            if action.dest == "dataset_name" and action.choices is not None:
-                extra = [
-                    c
-                    for c in (
-                        "daily-omni",
-                        "seed-tts",
-                        "seed-tts-text",
-                        "seed-tts-design",
-                        "ttsd",
-                        "sound-effect",
-                    )
-                    if c not in action.choices
-                ]
-                if extra:
-                    action.choices = list(action.choices) + extra
-            if action.dest == "backend" and action.choices is not None:
-                extra = [c for c in ("openai-image-edits-omni",) if c not in action.choices]
-                if extra:
-                    action.choices = list(action.choices) + extra
-        _extend_omni_dataset_name_choices(parser)
-
-        # Update help messages for omni-specific features
-        for action in parser._actions:
-            if action.dest == "percentile_metrics":
-                action.help = (
-                    "Comma-separated list of selected metrics to report percentiles. "
-                    'For text metrics, "ttft", "tpot", and "itl" affect the global benchmark and text '
-                    'stage metrics. "tpop" also requests text TPOT/TPOP globally and per stage, and internal '
-                    'stream TPOP. "ttfc", "tpoc", and "icl" only affect internal stream stage metrics. '
-                    'Audio metrics include "audio_ttfp", "audio_rtf", "audio_duration", and "audio_underrun".'
-                )
-            if action.dest == "random_mm_limit_mm_per_prompt":
-                action.help = (
-                    "Per-modality hard caps for items attached per request, e.g. "
-                    '\'{"image": 3, "video": 0, "audio": 1}\'. The sampled per-request item '
-                    "count is clamped to the sum of these limits. When a modality "
-                    "reaches its cap, its buckets are excluded and probabilities are "
-                    "renormalized."
-                )
-            if action.dest == "random_mm_bucket_config":
-                action.help = (
-                    "The bucket config is a dictionary mapping a multimodal item"
-                    "sampling configuration to a probability."
-                    "Currently allows for 3 modalities: audio, images and videos. "
-                    "A bucket key is a tuple of (height, width, num_frames)"
-                    "The value is the probability of sampling that specific item. "
-                    "Example: "
-                    "--random-mm-bucket-config "
-                    "{(256, 256, 1): 0.5, (720, 1280, 16): 0.4, (0, 1, 5): 0.10} "
-                    "First item: images with resolution 256x256 w.p. 0.5"
-                    "Second item: videos with resolution 720x1280 and 16 frames "
-                    "Third item: audios with 1s duration and 5 channels w.p. 0.1"
-                    "OBS.: If the probabilities do not sum to 1, they are normalized."
-                )
+        add_omni_args(parser)
+        extend_omni_choices(parser)
+        update_omni_help(parser)
 
     @staticmethod
     def cmd(args: argparse.Namespace) -> None:
-        if getattr(args, "daily_omni_save_eval_items", False):
-            os.environ["DAILY_OMNI_SAVE_EVAL_ITEMS"] = "1"
-        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"
-        image_edits_bot_task = getattr(args, "image_edits_bot_task", None)
-        if image_edits_bot_task is not None:
-            extra_body = dict(getattr(args, "extra_body", None) or {})
-            extra_body.setdefault("bot_task", image_edits_bot_task)
-            args.extra_body = extra_body
+        preprocess_serve_args(args)
         main(args)