[BugFix] Fix the issue of dataset names not being resolved (#4094)

vllm_omni/entrypoints/cli/benchmark/serve.py

@@ -131,6 +131,35 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
+_OMNI_BENCH_DATASET_CHOICES = (
+    "daily-omni",
+    "seed-tts",
+    "seed-tts-text",
+    "seed-tts-design",
+    "ttsd",
+    "sound-effect",
+)
+
+
+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."""
 
@@ -145,22 +174,7 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
         add_daily_omni_cli_args(parser)
         add_seed_tts_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
+        _extend_omni_dataset_name_choices(parser)
 
         # Update help messages for omni-specific features
         for action in parser._actions: