[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}}

tests/entrypoints/test_omni_entrypoints.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import argparse
 import queue
 from collections.abc import Callable
 from types import SimpleNamespace
@@ -176,83 +175,6 @@ def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine)
     )
 
 
-def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-    from vllm_omni.entrypoints.omni import Omni
-
-    captured: dict[str, Any] = {}
-
-    def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine:
-        captured.update(kwargs)
-        return FakeAsyncOmniEngine()
-
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine)
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    parser.add_argument("--batch-timeout", type=int, default=10)
-    nullify_stage_engine_defaults(parser)
-    args = parser.parse_args([])
-    args.model = "fake-model"
-
-    Omni(**vars(args))
-
-    assert captured["gpu_memory_utilization"] is None
-    assert captured["batch_timeout"] == 10
-    assert "_cli_explicit_keys" not in captured
-
-
-def test_from_cli_args_warns_and_forwards_without_internal_keys(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    captured: dict[str, Any] = {}
-
-    def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine:
-        captured.update(kwargs)
-        return FakeAsyncOmniEngine()
-
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine)
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
-
-    args = argparse.Namespace(model="fake-model", gpu_memory_utilization=0.9, _cli_explicit_keys={"model"})
-    with pytest.deprecated_call(match="from_cli_args"):
-        Omni.from_cli_args(args)
-
-    assert captured["gpu_memory_utilization"] == 0.9
-    assert "_cli_explicit_keys" not in captured
-
-
-def test_deprecated_from_cli_args_preserves_legacy_parser_nulling(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    from vllm_omni.entrypoints.omni import Omni
-
-    captured: dict[str, Any] = {}
-
-    def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine:
-        captured.update(kwargs)
-        return FakeAsyncOmniEngine()
-
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine)
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
-    monkeypatch.setattr("sys.argv", ["prog"])
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    parser.add_argument("--batch-timeout", type=int, default=10)
-    args = parser.parse_args([])
-    args.model = "fake-model"
-
-    with pytest.deprecated_call(match="from_cli_args"):
-        Omni.from_cli_args(args, parser=parser)
-
-    assert captured["gpu_memory_utilization"] is None
-    assert captured["batch_timeout"] == 10
-
-
 def _make_base():
     from vllm_omni.entrypoints.omni_base import OmniBase
 

tests/entrypoints/test_serve.py

@@ -4,65 +4,56 @@ from __future__ import annotations
 
 import argparse
 from types import SimpleNamespace
-from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
 
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def test_serve_parser_accepts_no_async_chunk() -> None:
-    """``--no-async-chunk`` should parse after deploy-overriding parser
-    defaults are nullified."""
-    try:
-        from vllm.utils.argparse_utils import FlexibleArgumentParser
-    except Exception as exc:
-        pytest.skip(f"Cannot build parser in this environment: {exc}")
-
-    root = FlexibleArgumentParser()
-    subparsers = root.add_subparsers(dest="subcommand")
+def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
+    """``--no-async-chunk`` should parse to ``async_chunk=False`` and mark the
+    shared deploy-level dest as explicitly provided by the user."""
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
     cmd = OmniServeCommand()
     cmd.subparser_init(subparsers)
 
     argv = ["serve", "fake-model", "--omni", "--no-async-chunk"]
-    args = root.parse_args(argv)
-
+    args = parser.parse_args(argv)
     assert args.async_chunk is False
 
-
-# ---------------------------------------------------------------------------
-# run_headless validation
-# ---------------------------------------------------------------------------
+    explicit = args.get_explicit_kwargs_dict()
+    assert args.get_explicit_kwargs_dict()
+    assert not explicit["async_chunk"]
 
 
-def _make_headless_args(**overrides: Any) -> argparse.Namespace:
-    """Build an argparse.Namespace shaped like the headless CLI passes in.
-
-    Defaults pass every validation gate so individual tests can mutate just
-    the field they're exercising.
-    """
-    defaults = dict(
-        model="fake-model",
-        stage_id=0,
-        replica_id=0,
-        omni_master_address="127.0.0.1",
-        omni_master_port=26000,
-        omni_replica_address=None,
-        omni_dp_size_local=1,
-        api_server_count=0,
-        worker_backend="multi_process",
-        stage_configs_path=None,
-        deploy_config=None,
-        log_stats=False,
-        disable_log_stats=False,
-        stage_init_timeout=600,
-        tokenizer=None,
+def _make_headless_args(**kwargs) -> TrackingNamespace:
+    defaults = {
+        "model": "fake-model",
+        "stage_id": 0,
+        "replica_id": 0,
+        "omni_master_address": "127.0.0.1",
+        "omni_master_port": 26000,
+        "omni_replica_address": None,
+        "omni_dp_size_local": 1,
+        "worker_backend": "multi_process",
+        "stage_configs_path": None,
+        "deploy_config": None,
+        "log_stats": False,
+        "disable_log_stats": False,
+        "stage_init_timeout": 600,
+        "tokenizer": None,
+    }
+    ns_kwargs = {**defaults, **kwargs}
+    ns = argparse.Namespace(**ns_kwargs)
+    return TrackingNamespace(
+        unfiltered_ns=ns,
+        explicit_keys=frozenset(ns.__dict__.keys()),
     )
-    defaults.update(overrides)
-    return argparse.Namespace(**defaults)
 
 
 def test_run_headless_requires_stage_id() -> None:
@@ -83,12 +74,6 @@ def test_run_headless_requires_master_port() -> None:
         run_headless(args)
 
 
-def test_run_headless_rejects_multi_api_server_count() -> None:
-    args = _make_headless_args(api_server_count=2)
-    with pytest.raises(ValueError, match="api_server_count can't be set"):
-        run_headless(args)
-
-
 def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
     args = _make_headless_args(worker_backend="ray")
     with pytest.raises(ValueError, match="worker_backend=multi_process"):