From 9bb497419885c667a1bc8e18c98151d93f371842 Mon Sep 17 00:00:00 2001
From: amy-why-3459 <wuhaiyan17@huawei.com>
Date: Wed, 3 Jun 2026 14:17:04 +0800
Subject: [PATCH] [BugFix] Fix the issue of dataset names not being resolved
 (#4094)

---
 vllm_omni/entrypoints/cli/benchmark/serve.py | 46 +++++++++++++-------
 1 file changed, 30 insertions(+), 16 deletions(-)

diff --git a/vllm_omni/entrypoints/cli/benchmark/serve.py b/vllm_omni/entrypoints/cli/benchmark/serve.py
index b969e5146..3ba03e297 100644
--- a/vllm_omni/entrypoints/cli/benchmark/serve.py
+++ b/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: