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: