[Core] Integrate TrackingArgumentParser (#3369)

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

tests/comfyui/test_comfyui_integration.py

@@ -30,12 +30,12 @@ from PIL import Image
 from pytest_mock import MockerFixture
 from vllm import SamplingParams
 from vllm.outputs import CompletionOutput, RequestOutput
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -444,7 +444,7 @@ def mock_async_omni(
 @pytest.fixture
 def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni):
     """Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     cmd = OmniServeCommand()
     cmd.subparser_init(subparsers)

tests/engine/test_async_omni_engine_stage_init.py

@@ -640,38 +640,6 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand
     assert captured_timeout == 302
 
 
-def test_async_omni_engine_reads_tokenizer_from_engine_args(monkeypatch):
-    import vllm_omni.engine.async_omni_engine as engine_mod
-
-    class DummyThread:
-        def __init__(self, *args, **kwargs):
-            pass
-
-        def start(self):
-            pass
-
-        def is_alive(self):
-            return False
-
-    fake_engine_args = types.SimpleNamespace(
-        _explicit_fields=frozenset({"tokenizer"}),
-        explicit_kwargs=lambda: {"tokenizer": "/tokenizer/from-engine-args"},
-    )
-
-    monkeypatch.setattr(engine_mod.threading, "Thread", DummyThread)
-    monkeypatch.setattr(
-        AsyncOmniEngine,
-        "_resolve_stage_configs",
-        lambda self, model, kwargs: ("dummy-config", [types.SimpleNamespace(engine_args={})]),
-    )
-    monkeypatch.setattr(AsyncOmniEngine, "_wait_for_orchestrator_init", lambda *_, **__: None)
-
-    engine = AsyncOmniEngine("dummy-model", engine_args=fake_engine_args)
-
-    assert engine.tokenizer == "/tokenizer/from-engine-args"
-    engine.shutdown()
-
-
 def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_path):
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 

tests/engine/test_single_stage_mode.py

@@ -403,39 +403,6 @@ class TestSingleStageModeDetection:
         assert engine.single_stage_mode is True
         assert engine._single_stage_id_filter is None
 
-    def test_engine_args_create_only_forwards_explicit_fields(self, mocker: MockerFixture):
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-
-        captured: dict[str, Any] = {}
-
-        def fake_resolve(self, model: str, kwargs: dict[str, Any]):
-            captured.update(kwargs)
-            return "/fake/path", [_make_stage_cfg(0)]
-
-        mocker.patch.object(AsyncOmniEngine, "_resolve_stage_configs", fake_resolve)
-        mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator")
-        mock_thread_cls = mocker.patch("threading.Thread")
-        mock_future_cls = mocker.patch("concurrent.futures.Future")
-        mock_future = mocker.Mock()
-        mock_future.result.return_value = mocker.Mock()
-        mock_future_cls.return_value = mock_future
-        mock_thread = mocker.Mock()
-        mock_thread.is_alive.return_value = False
-        mock_thread_cls.return_value = mock_thread
-
-        ea = OmniEngineArgs.create(model="ignored", gpu_memory_utilization=0.5)
-        AsyncOmniEngine(model="fake-model", engine_args=ea)
-
-        assert captured["gpu_memory_utilization"] == 0.5
-        assert "model" not in captured
-        assert "max_num_seqs" not in captured
-
-    def test_bare_engine_args_rejected(self, mocker: MockerFixture):
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-
-        with pytest.raises(TypeError, match="OmniEngineArgs.create"):
-            self._make_engine_no_thread(mocker, engine_args=OmniEngineArgs(model="fake-model"))
-
     def test_master_address_and_port_stored(self, mocker: MockerFixture):
         engine = self._make_engine_no_thread(
             mocker,

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"):

tests/test_arg_utils.py

@@ -5,303 +5,41 @@ hold for the orchestrator/engine/server CLI flag partition."""
 
 from __future__ import annotations
 
-import logging
-from dataclasses import dataclass, fields
+from dataclasses import fields
 
 import pytest
 
 from vllm_omni.engine.arg_utils import (
     SHARED_FIELDS,
-    derive_server_dests_from_vllm_parser,
+    OmniEngineArgs,
     internal_blacklist_keys,
-    orchestrator_args_from_argparse,
     orchestrator_field_names,
-    split_kwargs,
 )
 
-# ---------------------------------------------------------------------------
-# Fake engine class for unit testing — avoids pulling in the full vllm
-# EngineArgs and its heavy __post_init__ at test time.
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class _FakeEngineArgs:
-    """Stand-in for OmniEngineArgs with a representative subset of fields."""
-
-    model: str = ""
-    stage_id: int = 0
-    max_num_seqs: int = 64
-    gpu_memory_utilization: float = 0.9
-    async_chunk: bool = False  # also in OrchestratorArgs → shared
-    log_stats: bool = False  # also in OrchestratorArgs → shared
-    stage_configs_path: str | None = None
-
-
-# ============================================================================
-# Invariant 1 — OrchestratorArgs and engine must not ambiguously overlap.
-# ============================================================================
-
-
-def test_no_ambiguous_overlap_with_fake_engine():
-    """OrchestratorArgs ∩ engine fields must be ⊆ SHARED_FIELDS."""
-    orch = orchestrator_field_names()
-    engine = {f.name for f in fields(_FakeEngineArgs)}
-    overlap = orch & engine
-    unexpected = overlap - SHARED_FIELDS
-    assert not unexpected, (
-        f"Fields declared in both OrchestratorArgs and the engine class "
-        f"but not in SHARED_FIELDS: {sorted(unexpected)}. These cause "
-        f"double-routing — either remove the duplicate declaration or add "
-        f"to SHARED_FIELDS if sharing is intentional."
-    )
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_no_ambiguous_overlap_with_real_engine():
-    """Same check, but against the real OmniEngineArgs."""
-    try:
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-    except Exception as exc:
-        pytest.skip(f"OmniEngineArgs not importable: {exc}")
+    """OrchestratorArgs ∩ OmniEngineArgs must be ⊆ SHARED_FIELDS ∪ orchestrator-captured.
+
+    Fields that the orchestrator intentionally captures to prevent uniform
+    propagation (async_chunk, tokenizer) are excluded — they exist in both
+    classes by design but flow through DeployConfig, not kwargs forwarding.
+    """
+    # Fields on both classes by design: orchestrator captures them to prevent
+    # uniform per-stage propagation; redistributed via DeployConfig.
+    orchestrator_captured = {"async_chunk", "tokenizer"}
 
     orch = orchestrator_field_names()
     engine = {f.name for f in fields(OmniEngineArgs)}
     overlap = orch & engine
-    unexpected = overlap - SHARED_FIELDS
+    unexpected = overlap - SHARED_FIELDS - orchestrator_captured
     assert not unexpected, (
-        f"Real OmniEngineArgs has ambiguous overlap with OrchestratorArgs: "
+        f"OmniEngineArgs has ambiguous overlap with OrchestratorArgs: "
         f"{sorted(unexpected)}. Update SHARED_FIELDS or remove duplication."
     )
 
 
-# ============================================================================
-# Invariant 2 — split_kwargs partitions correctly.
-# ============================================================================
-
-
-def test_split_orchestrator_only():
-    """Pure orchestrator fields go to OrchestratorArgs, not engine_kwargs."""
-    raw = {"stage_init_timeout": 500, "worker_backend": "ray"}
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert orch.stage_init_timeout == 500
-    assert orch.worker_backend == "ray"
-    assert "stage_init_timeout" not in engine
-    assert "worker_backend" not in engine
-
-
-def test_split_engine_only():
-    """Pure engine fields go to engine_kwargs, not OrchestratorArgs."""
-    raw = {"max_num_seqs": 128, "gpu_memory_utilization": 0.85}
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert engine["max_num_seqs"] == 128
-    assert engine["gpu_memory_utilization"] == 0.85
-    # These fields don't exist on OrchestratorArgs at all.
-
-
-def test_split_shared_fields_go_to_both():
-    """Fields in SHARED_FIELDS are copied to both buckets."""
-    raw = {"model": "Qwen/Qwen2.5-Omni-7B", "log_stats": True}
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert orch.log_stats is True
-    assert engine["model"] == "Qwen/Qwen2.5-Omni-7B"
-    assert engine["log_stats"] is True
-
-
-def test_split_drops_unclassified():
-    """Unclassified fields (uvicorn/server) are dropped silently."""
-    raw = {
-        "max_num_seqs": 64,  # engine
-        "host": "0.0.0.0",  # unclassified (server)
-        "port": 8091,  # unclassified (server)
-        "ssl_keyfile": "key.pem",  # unclassified (server)
-    }
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert engine == {"max_num_seqs": 64}
-    assert "host" not in engine
-    assert "port" not in engine
-    assert "ssl_keyfile" not in engine
-
-
-def test_split_mixed_real_world():
-    """End-to-end: raw CLI kwargs with all three classes present."""
-    raw = {
-        # orchestrator
-        "stage_init_timeout": 400,
-        "deploy_config": "/tmp/deploy.yaml",
-        "worker_backend": "multi_process",
-        "async_chunk": True,
-        # engine
-        "max_num_seqs": 32,
-        "gpu_memory_utilization": 0.8,
-        # shared
-        "model": "Qwen/Qwen3-Omni",
-        "log_stats": False,
-        # server / unclassified
-        "host": "0.0.0.0",
-        "port": 8091,
-        "api_key": "secret",
-        # None values
-        "ray_address": None,
-    }
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-
-    # Orchestrator side
-    assert orch.stage_init_timeout == 400
-    assert orch.deploy_config == "/tmp/deploy.yaml"
-    assert orch.worker_backend == "multi_process"
-    assert orch.async_chunk is True
-    assert orch.log_stats is False  # shared, read from raw
-    assert orch.ray_address is None  # default preserved
-
-    # Engine side
-    assert engine["max_num_seqs"] == 32
-    assert engine["gpu_memory_utilization"] == 0.8
-    assert engine["model"] == "Qwen/Qwen3-Omni"
-    assert engine["log_stats"] is False
-    assert "host" not in engine
-    assert "port" not in engine
-    assert "api_key" not in engine
-    # orchestrator-only keys never reach engine
-    assert "stage_init_timeout" not in engine
-    assert "deploy_config" not in engine
-    # async_chunk is in SHARED_FIELDS so it flows to both
-    assert "async_chunk" in engine
-
-
-# ============================================================================
-# Invariant 3 — user-typed unclassifiable flags warn (don't fail silently).
-# ============================================================================
-
-
-def test_user_typed_unclassified_warns(monkeypatch):
-    """If the user types a flag we can't route, warn — don't silently drop."""
-    import io
-
-    buf = io.StringIO()
-    vllm_logger = logging.getLogger("vllm")
-    orig_handlers = list(vllm_logger.handlers)
-    vllm_logger.handlers.clear()
-    handler = logging.StreamHandler(buf)
-    handler.setLevel(logging.WARNING)
-    vllm_logger.addHandler(handler)
-    try:
-        raw = {"bogus_flag": "value", "max_num_seqs": 64}
-        split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed={"bogus_flag"})
-        assert "bogus_flag" in buf.getvalue(), f"Expected 'bogus_flag' in log output, got: {buf.getvalue()!r}"
-    finally:
-        vllm_logger.handlers.clear()
-        for h in orig_handlers:
-            vllm_logger.addHandler(h)
-
-
-def test_unclassified_without_user_typed_silent():
-    """Without user_typed, unclassified keys drop silently (argparse defaults
-    for server flags shouldn't spam logs on every launch)."""
-    import io
-
-    vllm_logger = logging.getLogger("vllm")
-    orig_handlers = list(vllm_logger.handlers)
-    vllm_logger.handlers.clear()
-    buf = io.StringIO()
-    handler = logging.StreamHandler(buf)
-    handler.setLevel(logging.WARNING)
-    vllm_logger.addHandler(handler)
-    try:
-        raw = {"host": "0.0.0.0", "port": 8091, "max_num_seqs": 64}
-        split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed=None)
-        assert "host" not in buf.getvalue()
-    finally:
-        vllm_logger.handlers.clear()
-        for h in orig_handlers:
-            vllm_logger.addHandler(h)
-
-
-# ============================================================================
-# Invariant 4 — CLI flag classification completeness.
-# Catches new flags added without updating OrchestratorArgs or OmniEngineArgs.
-# ============================================================================
-
-
-def test_all_omni_cli_flags_classified():
-    """Every vllm-omni-added CLI flag must be classifiable.
-
-    Runs ``OmniServeCommand.subparser_init`` and checks that every new
-    argument (compared to vllm's base parser) is either:
-      - a field on OrchestratorArgs, OR
-      - a field on OmniEngineArgs, OR
-      - in SHARED_FIELDS
-    """
-    try:
-        from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-        from vllm_omni.entrypoints.cli.serve import OmniServeCommand
-    except Exception as exc:
-        pytest.skip(f"Cannot build parser in this environment: {exc}")
-
-    # Build the serve parser
-    root = FlexibleArgumentParser()
-    subparsers = root.add_subparsers()
-    cmd = OmniServeCommand()
-    try:
-        parser = cmd.subparser_init(subparsers)
-    except Exception as exc:
-        pytest.skip(f"subparser_init failed (dev env issue): {exc}")
-
-    all_dests = {a.dest for a in parser._actions if a.dest and a.dest not in {"help", "model_tag"}}
-
-    orch = orchestrator_field_names()
-    engine = {f.name for f in fields(OmniEngineArgs)}
-    server_derived = derive_server_dests_from_vllm_parser()
-
-    unclassified = all_dests - orch - engine - SHARED_FIELDS - server_derived
-    # Some argparse-internal dests (suppressed, private) may not match —
-    # filter those out.
-    unclassified = {d for d in unclassified if not d.startswith("_")}
-
-    assert not unclassified, (
-        f"These CLI flags are not classified as "
-        f"orchestrator/engine/shared/server: {sorted(unclassified)}. "
-        f"Add them to OrchestratorArgs (if consumed by orchestrator), "
-        f"OmniEngineArgs (if consumed by per-stage engine), or the known-server "
-        f"allowlist (if they're vllm frontend flags). "
-        f"If intentional (e.g. a new CLI-only flag that doesn't map to either "
-        f"dataclass), add it to a KNOWN_UNROUTED allowlist."
-    )
-
-
-# ============================================================================
-# argparse interop (Phase 3).
-# ============================================================================
-
-
-def test_orchestrator_args_from_argparse():
-    """Can build OrchestratorArgs from an argparse.Namespace."""
-    import argparse
-
-    ns = argparse.Namespace(
-        stage_init_timeout=500,
-        deploy_config="/tmp/x.yaml",
-        max_num_seqs=64,  # engine field — ignored
-        host="0.0.0.0",  # server field — ignored
-    )
-    orch = orchestrator_args_from_argparse(ns)
-    assert orch.stage_init_timeout == 500
-    assert orch.deploy_config == "/tmp/x.yaml"
-    assert orch.worker_backend == "multi_process"  # default
-
-
-def test_derive_server_dests_returns_frozenset():
-    """Server-dest derivation returns a frozenset (possibly empty)."""
-    result = derive_server_dests_from_vllm_parser()
-    assert isinstance(result, frozenset)
-
-
-# ============================================================================
-# internal_blacklist_keys — single source of truth for per-stage forwarding.
-# ============================================================================
-
-
 def test_internal_blacklist_keys_derived_from_orchestrator():
     """Blacklist is exactly OrchestratorArgs fields minus SHARED_FIELDS.
 
@@ -311,214 +49,3 @@ def test_internal_blacklist_keys_derived_from_orchestrator():
     """
     blacklist = internal_blacklist_keys()
     assert blacklist == orchestrator_field_names() - SHARED_FIELDS
-    # Spot-check expected entries
-    assert "stage_init_timeout" in blacklist
-    assert "deploy_config" in blacklist
-    assert "async_chunk" not in blacklist  # now in SHARED_FIELDS
-    # Shared fields must NOT appear — they flow to both orchestrator and engine
-    assert "model" not in blacklist
-    assert "log_stats" not in blacklist
-
-
-# ============================================================================
-# Boundary value analysis — edge cases around split_kwargs.
-# ============================================================================
-
-
-def test_split_empty_kwargs():
-    """Empty kwargs yields default OrchestratorArgs and empty engine dict."""
-    orch, engine = split_kwargs({}, engine_cls=_FakeEngineArgs)
-    assert orch.stage_init_timeout == 300  # dataclass default
-    assert orch.worker_backend == "multi_process"  # dataclass default
-    assert engine == {}
-
-
-def test_split_all_none_values_preserved_on_orchestrator():
-    """None values for orchestrator fields are kept (represents 'not set')."""
-    raw = {"ray_address": None, "deploy_config": None, "max_num_seqs": None}
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert orch.ray_address is None
-    assert orch.deploy_config is None
-    # Engine-side None still passes through; caller decides semantics downstream.
-    assert engine.get("max_num_seqs") is None
-
-
-def test_split_user_typed_with_empty_kwargs_no_warn(caplog):
-    """user_typed non-empty but kwargs empty — no warnings emitted."""
-    with caplog.at_level(logging.WARNING, logger="vllm_omni.engine.arg_utils"):
-        split_kwargs({}, engine_cls=_FakeEngineArgs, user_typed={"nothing"})
-    assert not caplog.records
-
-
-def test_ambiguous_field_strict_raises():
-    """strict=True raises ValueError on overlap outside SHARED_FIELDS."""
-
-    # deploy_config is on OrchestratorArgs; declaring it on the engine class
-    # too (without adding to SHARED_FIELDS) creates an ambiguous route.
-    @dataclass
-    class _AmbiguousEngine:
-        deploy_config: str | None = None
-
-    with pytest.raises(ValueError, match="both OrchestratorArgs and"):
-        split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=True)
-
-
-def test_ambiguous_field_non_strict_routes_to_orchestrator():
-    """strict=False logs ERROR but routes the ambiguous field to orchestrator."""
-
-    @dataclass
-    class _AmbiguousEngine:
-        deploy_config: str | None = None
-
-    import io
-
-    vllm_logger = logging.getLogger("vllm")
-    orig_handlers = list(vllm_logger.handlers)
-    vllm_logger.handlers.clear()
-    buf = io.StringIO()
-    handler = logging.StreamHandler(buf)
-    handler.setLevel(logging.ERROR)
-    vllm_logger.addHandler(handler)
-    try:
-        orch, engine = split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=False)
-        assert orch.deploy_config == "x"
-        assert "deploy_config" not in engine
-        assert "both OrchestratorArgs" in buf.getvalue()
-    finally:
-        vllm_logger.handlers.clear()
-        for h in orig_handlers:
-            vllm_logger.addHandler(h)
-
-
-# Sentinel-default precedence invariants (#3035)
-
-
-def _build_full_serve_parser():
-    from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-    try:
-        from vllm.entrypoints.openai.cli_args import make_arg_parser
-    except ImportError:
-        pytest.skip("vllm parser not importable")
-    return make_arg_parser(FlexibleArgumentParser())
-
-
-def test_nullify_stage_engine_defaults_resets_inherited_defaults():
-    import argparse
-
-    from vllm_omni.config.stage_config import deploy_override_field_names
-    from vllm_omni.engine.arg_utils import (
-        nullify_stage_engine_defaults,
-    )
-
-    parser = _build_full_serve_parser()
-    nullify_stage_engine_defaults(parser)
-
-    override_dests = deploy_override_field_names()
-    offenders = [
-        (a.dest, a.default)
-        for a in parser._actions
-        if a.dest not in ("help", "version")
-        and a.option_strings
-        and a.dest in override_dests
-        and a.default is not None
-        and a.default is not argparse.SUPPRESS
-    ]
-    assert not offenders, f"Stage flags with non-None defaults after nullify: {offenders}"
-
-
-def test_non_override_flags_keep_real_defaults_after_nullify():
-    import argparse
-
-    from vllm_omni.config.stage_config import deploy_override_field_names
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--batch-timeout", type=int, default=10, help="Batch timeout.")
-    parser.add_argument("--max-num-seqs", type=int, default=64, help="Max num seqs.")
-    nullify_stage_engine_defaults(parser)
-
-    assert "batch_timeout" not in deploy_override_field_names()
-
-    batch_timeout = next(a for a in parser._actions if a.dest == "batch_timeout")
-    max_num_seqs = next(a for a in parser._actions if a.dest == "max_num_seqs")
-    assert batch_timeout.default == 10
-    assert max_num_seqs.default is None
-
-
-def test_help_text_preserves_default_after_nullify():
-    # Real defaults must stay visible in --help even though parser stores None.
-    import argparse
-
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--max-num-seqs", type=int, default=42, help="Example knob.")
-    nullify_stage_engine_defaults(parser)
-
-    action = next(a for a in parser._actions if a.dest == "max_num_seqs")
-    assert action.default is None
-    assert "(default: 42)" in action.help
-
-
-_OMNIENGINEARGS_USER_INPUT_FIELDS = frozenset(
-    {
-        "model_arch",
-        "engine_output_type",
-        "hf_config_name",
-        "custom_process_next_stage_input_func",
-        "subtalker_sampling_params",
-        "omni_kv_config",
-        "quantization_config",
-        "worker_type",
-        "task_type",
-        "worker_cls",
-        "omni_master_address",
-        "omni_master_port",
-        "stage_configs_path",
-        "output_modalities",
-        "custom_pipeline_args",
-    }
-)
-
-
-def test_omniengineargs_user_input_fields_default_to_none():
-    try:
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-    except Exception as exc:
-        pytest.skip(f"OmniEngineArgs not importable: {exc}")
-
-    offenders = [
-        (f.name, f.default)
-        for f in fields(OmniEngineArgs)
-        if f.name in _OMNIENGINEARGS_USER_INPUT_FIELDS
-        and f.default is not dataclasses.MISSING
-        and f.default is not None
-    ]
-    assert not offenders, f"User-input fields with non-None defaults: {offenders}"
-
-
-def test_omniengineargs_create_tracks_explicit_fields():
-    try:
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-    except Exception as exc:
-        pytest.skip(f"OmniEngineArgs not importable: {exc}")
-
-    ea = OmniEngineArgs.create(model="x", gpu_memory_utilization=0.5)
-    assert ea._explicit_fields == frozenset({"model", "gpu_memory_utilization"})
-    assert ea.explicit_kwargs() == {"model": "x", "gpu_memory_utilization": 0.5}
-
-
-def test_omniengineargs_bare_constructor_has_no_explicit_tracking():
-    try:
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-    except Exception as exc:
-        pytest.skip(f"OmniEngineArgs not importable: {exc}")
-
-    ea = OmniEngineArgs(model="x")
-    assert not hasattr(ea, "_explicit_fields")
-    assert "model" in ea.explicit_kwargs()
-
-
-# dataclasses already imported via ``from dataclasses import dataclass, fields``
-import dataclasses  # noqa: E402  -- needed for MISSING sentinel above

tests/test_config_factory.py

@@ -29,7 +29,6 @@ from vllm_omni.config.stage_config import (
     _deep_merge_stage,
     _resolve_scheduler,
     build_stage_runtime_overrides,
-    deploy_override_field_names,
     load_deploy_config,
     merge_pipeline_deploy,
     register_pipeline,
@@ -39,6 +38,8 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_black
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class TestStageType:
     """Tests for StageType enum."""
@@ -156,23 +157,6 @@ class TestStageConfig:
         omega_config = config.to_omegaconf()
         assert omega_config.engine_args.max_num_seqs == 32
 
-    def test_to_omegaconf_omits_none_deploy_overrides_for_engine_args(self):
-        """None deploy overrides must fall through to EngineArgs defaults."""
-
-        config = StageConfig(
-            stage_id=0,
-            model_stage="thinker",
-            runtime_overrides={name: None for name in deploy_override_field_names()},
-        )
-
-        omega_config = config.to_omegaconf()
-        engine_args = dict(omega_config.engine_args)
-
-        assert "devices" not in engine_args
-        assert "max_batch_size" not in engine_args
-        for name in deploy_override_field_names() - {"devices"}:
-            assert name not in engine_args
-
     def test_to_omegaconf_diffusion_parallel_overrides_replace_nested_values(self):
         config = StageConfig(
             stage_id=1,
@@ -1018,64 +1002,7 @@ class TestPipelineRegistry:
 
 
 class TestDeployConfigLoading:
-    def test_deploy_override_fields_include_deploy_schema_fields(self):
-        expected_fields = {
-            "active_stream_window",
-            "async_chunk",
-            # StageDeployConfig: stage placement and runtime fields.
-            "devices",
-            # StageDeployConfig: vLLM EngineArgs fields.
-            "async_scheduling",
-            "compilation_config",
-            "config_format",
-            "custom_voice_dir",
-            "data_parallel_size",
-            "devices",
-            "disable_hybrid_kv_cache_manager",
-            "enable_flashinfer_autotune",
-            "enforce_eager",
-            "gpu_memory_utilization",
-            "load_format",
-            "max_model_len",
-            "max_num_batched_tokens",
-            "max_num_seqs",
-            "mm_processor_cache_gb",
-            "profiler_config",
-            "skip_mm_profiling",
-            "subtalker_sampling_params",
-            "tensor_parallel_size",
-            "tokenizer_mode",
-            # StageDeployConfig: diffusion parallel_config deploy override fields.
-            "cfg_parallel_size",
-            "enable_expert_parallel",
-            "hsdp_replicate_size",
-            "hsdp_shard_size",
-            "ring_degree",
-            "sequence_parallel_size",
-            "ulysses_degree",
-            "ulysses_mode",
-            "use_hsdp",
-            "vae_patch_parallel_size",
-            # DeployConfig: pipeline-wide engine settings.
-            "data_parallel_size",
-            "distributed_executor_backend",
-            "dtype",
-            "enable_chunked_prefill",
-            "enable_prefix_caching",
-            "pipeline_parallel_size",
-            "quantization",
-            "trust_remote_code",
-        }
-
-        actual_fields = deploy_override_field_names()
-        assert expected_fields == actual_fields, (
-            f"added={actual_fields - expected_fields}, removed={expected_fields - actual_fields}"
-        )
-
     def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path):
-        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml"
         custom_voice_dir = tmp_path / "voices"
         deploy_path.write_text(
@@ -1913,16 +1840,6 @@ class TestSentinelDefaultPrecedence:
         named = [p for p in sig.parameters.values() if p.kind != p.VAR_KEYWORD]
         assert "cli_explicit_keys" not in {p.name for p in named}
 
-    def test_cli_explicit_keys_kwarg_emits_deprecation(self):
-        with warnings.catch_warnings(record=True) as w:
-            warnings.simplefilter("always")
-            StageConfigFactory._create_from_registry(
-                "qwen3_omni_moe",
-                cli_overrides={},
-                cli_explicit_keys={"max_num_seqs"},
-            )
-            assert any(issubclass(x.category, DeprecationWarning) for x in w)
-
     def test_async_chunk_dispatches_processors(self):
         """A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end
         processors based on ``deploy.async_chunk``, without needing a

tests/utils/test_tracking_parser.py

@@ -0,0 +1,545 @@
+"""Tests for TrackingArgumentParser and related utilities."""
+
+import argparse
+import json
+from unittest.mock import patch
+
+import pytest
+import yaml
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.config.stage_config import (
+    _PIPELINE_REGISTRY,
+    DeployConfig,
+    PipelineConfig,
+    StageConfigFactory,
+    StageDeployConfig,
+    StagePipelineConfig,
+)
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
+
+### Fake pipeline/deploy config for integration tests
+
+_TEST_MODEL = "_test_tracking"
+
+_TEST_PIPELINE = PipelineConfig(
+    model_type=_TEST_MODEL,
+    stages=(
+        StagePipelineConfig(stage_id=0, model_stage="stage0"),
+        StagePipelineConfig(stage_id=1, model_stage="stage1"),
+        StagePipelineConfig(stage_id=2, model_stage="stage2"),
+    ),
+)
+
+_TEST_DEPLOY = DeployConfig(
+    async_chunk=False,
+    stages=[
+        StageDeployConfig(stage_id=0, gpu_memory_utilization=0.5, max_num_seqs=32),
+        StageDeployConfig(stage_id=1, gpu_memory_utilization=0.6, max_num_seqs=64),
+        StageDeployConfig(stage_id=2, gpu_memory_utilization=0.1, max_num_seqs=16),
+    ],
+)
+
+
+@pytest.fixture()
+def mock_stages(monkeypatch):
+    """Register a fake pipeline and mock deploy YAML loading."""
+    monkeypatch.setitem(_PIPELINE_REGISTRY._loaded, _TEST_MODEL, _TEST_PIPELINE)
+    monkeypatch.setattr(
+        "vllm_omni.config.stage_config.load_deploy_config",
+        lambda _path: _TEST_DEPLOY,
+    )
+    return __file__
+
+
+### Tests for TrackingNamespace
+def test_tracking_namespaces_cant_be_nested():
+    """Ensure tracking namespaces explode if we try to nest them."""
+    track_ns = TrackingNamespace(
+        unfiltered_ns=argparse.Namespace(foo="bar"),
+        explicit_keys=frozenset(),
+    )
+
+    with pytest.raises(ValueError):
+        TrackingNamespace(
+            unfiltered_ns=track_ns,
+            explicit_keys=frozenset(),
+        )
+
+
+def test_tracking_namespaces_init():
+    """Check simple initialization for tracking namespaces."""
+    unfiltered_ns = argparse.Namespace(foo="bar")
+    tracked_ns = TrackingNamespace(
+        unfiltered_ns=unfiltered_ns,
+        explicit_keys=frozenset({"foo"}),
+    )
+    assert tracked_ns.foo == "bar"
+    assert tracked_ns.explicit_keys == frozenset({"foo"})
+
+
+def test_tracking_filtering():
+    """Ensure tracking namespaces are filterable."""
+    unfiltered_ns = argparse.Namespace(foo="bar", baz="foobar")
+    tracked_ns = TrackingNamespace(
+        unfiltered_ns=unfiltered_ns,
+        explicit_keys=frozenset({"foo"}),
+    )
+    assert tracked_ns.foo == "bar"
+    assert tracked_ns.baz == "foobar"
+    assert tracked_ns.explicit_keys == frozenset({"foo"})
+    # baz gets dropped because it's not marked in explicit_keys
+    assert tracked_ns.get_explicit_kwargs_dict() == {"foo": "bar"}
+
+
+def test_setattr_writes_through_to_unfiltered_ns():
+    """Ensure mutating an attribute on TrackingNamespace forwards to
+    get_explicit_kwargs_dict() and vars()."""
+    unfiltered_ns = argparse.Namespace(model="original")
+    tracked_ns = TrackingNamespace(
+        unfiltered_ns=unfiltered_ns,
+        explicit_keys=frozenset({"model"}),
+    )
+    assert tracked_ns.model == "original"
+    assert tracked_ns.get_explicit_kwargs_dict() == {"model": "original"}
+    assert vars(tracked_ns) == {"model": "original"}
+
+    # Ensure if we update the namespace, it's forwarded correctly
+    tracked_ns.model = "updated"
+    assert tracked_ns.model == "updated"
+    assert tracked_ns.get_explicit_kwargs_dict() == {"model": "updated"}
+    assert vars(tracked_ns) == {"model": "updated"}
+
+
+### Tests for simple cases (no nested parsers or groups)
+def test_vars_does_not_expose_internals():
+    """Ensure vars on a TrackingArgumentParser is identical to running on the real one."""
+    tracking = TrackingArgumentParser()
+    flexible = FlexibleArgumentParser()
+    namespaces = []
+
+    # For both parser, register two args, but only pass one
+    for parser in [tracking, flexible]:
+        parser.add_argument("--foo", type=int, default=42)
+        parser.add_argument("--bar", type=str, default="x")
+        namespaces.append(parser.parse_args(["--foo", "100"]))
+
+    # Ensure that vars are the same result, since it isn't filtered
+    tracked_ns, flexible_ns = namespaces
+    tracked_kwargs = vars(tracked_ns)
+    real_kwargs = vars(flexible_ns)
+    assert tracked_kwargs == real_kwargs == {"foo": 100, "bar": "x"}
+    assert tracked_ns.explicit_keys == {"foo"}
+
+
+def test_default_not_detected():
+    """Ensure omitted defaults aren't in explicit keys and take defaults."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=42)
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo == 42
+
+
+@pytest.mark.parametrize("val", ["42", "100"])
+def test_explicit_value_equal_to_default(val):
+    """Ensure explicit keys correctly handles passed values."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=42)
+    ns = p.parse_args(["--foo", val])
+    assert ns.explicit_keys == {"foo"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo == int(val)
+
+
+def test_equals_syntax():
+    """Ensure equals syntax is handled correctly wrt explicit keys."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=42)
+    ns = p.parse_args(["--foo=100"])
+    assert ns.explicit_keys == {"foo"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo == 100
+
+
+def test_explicit_none_via_store_const():
+    """Ensure passing `None` to the namespace, it isn't filtered from explicit_keys."""
+    parser = TrackingArgumentParser()
+    parser.add_argument("--foo", action="store_const", const=None, default="Something else")
+    ns = parser.parse_args(["--foo"])
+    # User explicitly passed --foo, so it should be in the explicit keys even though it's None
+    assert ns.explicit_keys == {"foo"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo is None
+
+
+def test_multiple_args_mixed():
+    """Ensure that explicit keys are correct when some are passed and others aren't."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=1)
+    p.add_argument("--bar", type=str, default="x")
+    p.add_argument("--baz", type=float, default=0.5)
+    ns = p.parse_args(["--foo", "10", "--baz", "0.9"])
+    assert ns.explicit_keys == {"foo", "baz"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.bar == "x"
+
+
+def test_store_true_default():
+    """Ensure that store true is handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("--verbose", action="store_true")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose is False
+
+
+def test_store_true_explicit():
+    """Ensure that store true is handled correctly when passed."""
+    p = TrackingArgumentParser()
+    p.add_argument("--verbose", action="store_true")
+    ns = p.parse_args(["--verbose"])
+    assert ns.explicit_keys == {"verbose"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose is True
+
+
+def test_store_false_default():
+    """Ensure that store false is handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("--disable-x", action="store_false", dest="enable_x", default=True)
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.enable_x is True
+
+
+def test_store_false_explicit():
+    """Ensure that store false is handled correctly when passed."""
+    p = TrackingArgumentParser()
+    p.add_argument("--disable-x", action="store_false")
+    ns = p.parse_args(["--disable-x"])
+    assert ns.explicit_keys == {"disable_x"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.disable_x is False
+
+
+def test_dest_is_reflected_in_explicit_keys():
+    """Ensure that explicit keys use dest correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, dest="bar")
+    ns = p.parse_args(["--foo", "100"])
+    assert ns.explicit_keys == {"bar"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.bar == 100
+
+
+def test_boolean_optional_action_default():
+    """Check that boolean optional actions are handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("--flag", action=argparse.BooleanOptionalAction)
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flag is None
+
+
+def test_boolean_optional_action_positive():
+    """Check that boolean optional actions are handled correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--flag", action=argparse.BooleanOptionalAction, default=None)
+    ns = p.parse_args(["--flag"])
+    assert ns.explicit_keys == {"flag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flag is True
+
+
+def test_no_option_strings_are_handled():
+    """Ensure --no-<feature> sets <feature> in the explicit keys correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--flag", action=argparse.BooleanOptionalAction, default=None)
+    ns = p.parse_args(["--no-flag"])
+    assert ns.explicit_keys == {"flag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flag is False
+
+
+def test_json_type():
+    """Ensure json type is handled correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--cfg", type=json.loads, default="{}")
+    ns = p.parse_args(["--cfg", '{"a": 1}'])
+    assert "cfg" in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.cfg == {"a": 1}
+
+
+def test_choices():
+    """Ensure choices are handled correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--mode", choices=["fast", "slow"], default="fast")
+    ns = p.parse_args(["--mode", "slow"])
+    assert "mode" in ns.explicit_keys
+    assert ns is not None
+    assert ns.mode == "slow"
+
+
+def test_explicit_positional_arg():
+    """Ensure positional args are handled correctly when provided."""
+    p = TrackingArgumentParser()
+    p.add_argument("name", nargs="?", default=None)
+    ns = p.parse_args(["hello"])
+    assert "name" in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.name == "hello"
+
+
+def test_omitted_positional_arg():
+    """Ensure positional args are handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("name", nargs="?", default=None)
+    ns = p.parse_args([])
+    assert "name" not in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.name is None
+
+
+def test_explicit_nargs():
+    """Ensure that variable num args are handled correctly when provided."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", nargs="*", default=None)
+    ns = p.parse_args(["--items", "a", "b"])
+    assert "items" in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items == ["a", "b"]
+
+
+def test_omitted_nargs():
+    """Ensure that variable num args are handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", nargs="*", default=None)
+    ns = p.parse_args([])
+    assert "items" not in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items is None
+
+
+def test_parse_known_args_tracking():
+    """Ensure parse_known_args is also trackable"""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=42)
+    ns, remaining = p.parse_known_args(["--foo", "10", "--unknown", "val"])
+    assert ns.explicit_keys == {"foo"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo == 10
+    assert remaining == ["--unknown", "val"]
+
+
+### Tests for group handling
+def test_group_arg_default():
+    """Ensure that that groups with defaults are handled correctly."""
+    p = TrackingArgumentParser()
+    g = p.add_argument_group("TestGroup")
+    g.add_argument("--bar", type=str, default="baz")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.bar == "baz"
+
+
+def test_group_arg_explicit():
+    """Ensure that that groups are handled correctly."""
+    p = TrackingArgumentParser()
+    g = p.add_argument_group("TestGroup")
+    g.add_argument("--bar", type=str, default="baz")
+    ns = p.parse_args(["--bar", "qux"])
+    assert ns.explicit_keys == {"bar"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.bar == "qux"
+
+
+def test_multiple_groups():
+    """Test multiple groups behave correctly."""
+    p = TrackingArgumentParser()
+    g1 = p.add_argument_group("Group1")
+    g2 = p.add_argument_group("Group2")
+    g1.add_argument("--a", type=int, default=1)
+    g2.add_argument("--b", type=int, default=2)
+    ns = p.parse_args(["--b", "20"])
+    assert ns.explicit_keys == {"b"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.a == 1
+    assert ns.b == 20
+
+
+def test_omitted_mutually_exclusive_group():
+    """Ensure that that mutually exclusive groups with defaults are handled correctly."""
+    p = TrackingArgumentParser()
+    grp = p.add_mutually_exclusive_group()
+    grp.add_argument("--json", action="store_true")
+    grp.add_argument("--text", action="store_true")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+
+
+def test_mutually_exclusive_group():
+    """Ensure that that mutually exclusive groups are handled correctly."""
+    p = TrackingArgumentParser()
+    grp = p.add_mutually_exclusive_group()
+    grp.add_argument("--json", action="store_true")
+    grp.add_argument("--text", action="store_true")
+    ns = p.parse_args(["--json"])
+    assert ns.explicit_keys == {"json"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.json is True
+    assert ns.text is False
+
+
+### Tests for subparser handling
+def test_subparser_explicit_detection():
+    """Ensure that subparsers are handled correctly."""
+    p = TrackingArgumentParser()
+    sub = p.add_subparsers()
+    child = sub.add_parser("foo")
+    child.add_argument("--bar", type=str)
+    ns = p.parse_args(["foo", "--bar", "baz"])
+    assert isinstance(child, TrackingArgumentParser)
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.explicit_keys == {"bar"}
+    assert ns.bar == "baz"
+
+
+def test_subparser_group_args():
+    """Ensure that subparsers with groups are handled correctly."""
+    p = TrackingArgumentParser()
+    sub = p.add_subparsers()
+    child = sub.add_parser("foo")
+    g = child.add_argument_group("Config")
+    g.add_argument("--port", type=int, default=8000)
+    g.add_argument("--host", type=str, default="localhost")
+    ns = p.parse_args(["foo", "--port", "9000"])
+    assert ns.explicit_keys == {"port"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.host == "localhost"
+
+
+### Tests for specific behaviors against FlexibleArgumentParser
+def test_config_file_args_detected(tmp_path):
+    """Ensure config is handled correctly for vLLM's FlexibleArgumentParser."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int)
+    p.add_argument("--bar", type=int)
+    cfg = tmp_path / "test.yaml"
+    cfg.write_text(yaml.dump({"foo": 100}))
+    ns = p.parse_args(["--config", str(cfg)])
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.explicit_keys == {"foo"}
+    assert ns.foo == 100
+
+
+def test_cli_overrides_config(tmp_path):
+    """Ensure tracking parser handles config vs cli overrides correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int)
+    cfg = tmp_path / "test.yaml"
+    cfg.write_text(yaml.dump({"foo": 100}))
+    ns = p.parse_args(["--config", str(cfg), "--foo", "200"])
+    assert isinstance(ns, TrackingNamespace)
+    assert "foo" in ns.explicit_keys
+    assert ns.foo == 200
+
+
+### Integration tests for arg resolution through StageConfigFactory
+def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages):
+    """Explicitly passed CLI values reach runtime_overrides on all stages."""
+    p = TrackingArgumentParser()
+    p.add_argument("--max-num-seqs", type=int, default=64)
+    ns = p.parse_args(["--max-num-seqs", "999"])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+    for stage in stages:
+        assert stage.runtime_overrides.get("max_num_seqs") == 999
+
+
+def test_omitted_default_not_in_runtime_overrides(mock_stages):
+    """Omitted defaults are overridden by deploy config values"""
+    p = TrackingArgumentParser()
+    p.add_argument("--max-num-seqs", type=int, default=64)
+    ns = p.parse_args([])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+    for stage in stages:
+        assert stage.runtime_overrides == {}
+
+
+def test_config_file_args_reach_runtime_overrides(mock_stages):
+    """Args from --config YAML must be treated as explicitly passed and
+    flow through to runtime_overrides."""
+    p = TrackingArgumentParser()
+    p.add_argument("--max-num-seqs", type=int, default=64)
+    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+    with patch(
+        "vllm.utils.argparse_utils.FlexibleArgumentParser.load_config_file", return_value=["--max-num-seqs", "999"]
+    ):
+        ns = p.parse_args(["--config", "fake.yaml"])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+    for stage in stages:
+        assert stage.runtime_overrides.get("max_num_seqs") == 999
+        assert "gpu_memory_utilization" not in stage.runtime_overrides
+
+
+def test_per_stage_override_routes_correctly(mock_stages):
+    """Ensure stage_<N>_<key> only affects the targeted stage."""
+    p = TrackingArgumentParser()
+    p.add_argument("--stage-0-gpu-memory-utilization", type=float)
+    ns = p.parse_args(["--stage-0-gpu-memory-utilization", "0.42"])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+    assert stages[0].runtime_overrides == {"gpu_memory_utilization": 0.42}
+    assert stages[1].runtime_overrides == {}
+    assert stages[2].runtime_overrides == {}
+
+
+def test_explicit_args_omitted_from_yaml(mock_stages):
+    """Ensure only passed args end up in runtime overrides (regardless of whether
+    they are defined in the yaml config)."""
+    p = TrackingArgumentParser()
+    p.add_argument("--enforce-eager", action="store_true")
+    p.add_argument("--max-num-seqs", type=int, default=64)
+    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+    # NOTE: enforce eager is not set in the mock config yaml.
+    ns = p.parse_args(["--enforce-eager"])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+
+    for stage in stages:
+        assert stage.runtime_overrides == {"enforce_eager": True}