[Core] Integrate TrackingArgumentParser (#3369)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -4,12 +4,11 @@
 from types import SimpleNamespace
 
 import pytest
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.config.stage_config import deploy_override_field_names
 from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
-from vllm_omni.entrypoints.cli.serve import OmniServeCommand, _create_default_diffusion_stage_cfg
+from vllm_omni.entrypoints.cli.serve import OmniServeCommand
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -34,15 +33,6 @@ def test_default_stage_config_includes_cache_backend():
     assert engine_args["model_stage"] == "diffusion"
 
 
-def test_default_stage_config_ignores_none_deploy_overrides():
-    """Ensure nullified deploy override defaults do not alter diffusion defaults."""
-    baseline = AsyncOmniEngine._create_default_diffusion_stage_cfg({})[0]
-    nullified_overrides = {name: None for name in deploy_override_field_names()}
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(nullified_overrides)[0]
-
-    assert stage_cfg == baseline
-
-
 def test_default_cache_config_used_when_missing():
     """Ensure default cache_config is synthesized when only backend is given."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -221,7 +211,7 @@ def test_default_stage_config_whitelist_none_fallback():
 
 def test_serve_cli_accepts_ulysses_mode():
     """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
 
@@ -237,7 +227,8 @@ def test_serve_cli_accepts_ulysses_mode():
         ]
     )
 
-    stage_cfg = _create_default_diffusion_stage_cfg(args)[0]
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
 
     assert args.ulysses_mode == "advanced_uaa"
@@ -247,7 +238,7 @@ def test_serve_cli_accepts_ulysses_mode():
 
 def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
     """Ensure diffusion serve CLI exposes the profiler switch."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
 
@@ -260,7 +251,8 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
         ]
     )
 
-    stage_cfg = _create_default_diffusion_stage_cfg(args)[0]
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
 
     assert args.enable_diffusion_pipeline_profiler is True
     assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
@@ -268,7 +260,7 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
 
 def test_serve_cli_accepts_diffusion_attention_backend():
     """Ensure diffusion serve CLI exposes the shorthand backend flag."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
 
@@ -282,7 +274,8 @@ def test_serve_cli_accepts_diffusion_attention_backend():
         ]
     )
 
-    stage_cfg = _create_default_diffusion_stage_cfg(args)[0]
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
     diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
 
     assert args.diffusion_attention_backend == "FLASH_ATTN"
@@ -293,7 +286,7 @@ def test_serve_cli_accepts_diffusion_attention_backend():
 
 def test_serve_cli_accepts_additional_config():
     """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
 
@@ -307,7 +300,8 @@ def test_serve_cli_accepts_additional_config():
         ]
     )
 
-    stage_cfg = _create_default_diffusion_stage_cfg(args)[0]
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))[0]
+
     engine_args = stage_cfg["engine_args"]
 
     assert args.additional_config == {"torchair_graph_config": {"enabled": True}}