[Bugfix] Support append/extend/count actions in TrackingArgumentParser (#4639)

Signed-off-by: Junhee Yoo <junhee.yoo@navercorp.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

tests/utils/test_tracking_parser.py

@@ -16,7 +16,12 @@ from vllm_omni.config.stage_config import (
     StageDeployConfig,
     StagePipelineConfig,
 )
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
+from vllm_omni.utils.tracking_parser import (
+    UNSET,
+    TrackingArgumentParser,
+    TrackingNamespace,
+    build_shadow_kwargs,
+)
 
 ### Fake pipeline/deploy config for integration tests
 
@@ -326,6 +331,132 @@ def test_omitted_nargs():
     assert ns.items is None
 
 
+### Tests for in-place mutating actions (append / extend / count).
+# These actions mutate their default in place, which would crash on the bare
+# ``UNSET`` sentinel. The shadow parser remaps them to a non-mutating store-style
+# action (see build_shadow_kwargs in tracking_parser), so the real namespace
+# still accumulates while explicit-arg tracking keeps working.
+def test_append_action_omitted():
+    """Omitted append args parse without error and aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--tag", action="append")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.tag is None
+
+
+def test_append_action_explicit():
+    """Repeated append args are collected and marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--tag", action="append")
+    ns = p.parse_args(["--tag", "a", "--tag", "b"])
+    assert ns.explicit_keys == {"tag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.tag == ["a", "b"]
+
+
+def test_append_const_action_omitted():
+    """Omitted append_const args aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--dbg", action="append_const", const="x", dest="flags")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flags is None
+
+
+def test_append_const_action_explicit():
+    """Repeated append_const args accumulate the const and are explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--dbg", action="append_const", const="x", dest="flags")
+    ns = p.parse_args(["--dbg", "--dbg"])
+    assert ns.explicit_keys == {"flags"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flags == ["x", "x"]
+
+
+def test_extend_action_omitted():
+    """Omitted extend args parse without error and aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", action="extend", nargs="+")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items is None
+
+
+def test_extend_action_explicit():
+    """Repeated extend args are flattened into one list and marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", action="extend", nargs="+")
+    ns = p.parse_args(["--items", "a", "b", "--items", "c"])
+    assert ns.explicit_keys == {"items"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items == ["a", "b", "c"]
+
+
+def test_count_action_omitted():
+    """Omitted count args keep their default and aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("-v", "--verbose", action="count", default=0)
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose == 0
+
+
+def test_count_action_explicit():
+    """Repeated count flags increment and are marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("-v", "--verbose", action="count", default=0)
+    ns = p.parse_args(["-vv"])
+    assert ns.explicit_keys == {"verbose"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose == 2
+
+
+def test_group_append_action_explicit():
+    """append args added via a group go through the same shadow-default path."""
+    p = TrackingArgumentParser()
+    g = p.add_argument_group("TestGroup")
+    g.add_argument("--tag", action="append")
+    ns = p.parse_args(["--tag", "a", "--tag", "b"])
+    assert ns.explicit_keys == {"tag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.tag == ["a", "b"]
+
+
+@pytest.mark.parametrize(
+    ("action", "expected"),
+    [
+        ("append", "store"),
+        ("extend", "store"),
+        ("append_const", "store_const"),
+        ("count", "store_const"),
+    ],
+)
+def test_build_shadow_kwargs_remaps_mutating_actions(action, expected):
+    """Mutating actions are remapped to a store-style action with an UNSET default."""
+    shadow = build_shadow_kwargs({"action": action, "const": "c"})
+    assert shadow["action"] == expected
+    assert shadow["default"] is UNSET
+
+
+def test_build_shadow_kwargs_count_gets_marker_const():
+    """count carries no const, so the remapped store_const gets a non-UNSET marker."""
+    shadow = build_shadow_kwargs({"action": "count"})
+    assert shadow["action"] == "store_const"
+    assert shadow["const"] is not UNSET
+
+
+def test_build_shadow_kwargs_leaves_other_actions_untouched():
+    """Non-mutating actions keep their action and only get the UNSET default."""
+    shadow = build_shadow_kwargs({"action": "store_true"})
+    assert shadow["action"] == "store_true"
+    assert shadow["default"] is UNSET
+
+
 def test_parse_known_args_tracking():
     """Ensure parse_known_args is also trackable"""
     p = TrackingArgumentParser()

vllm_omni/utils/tracking_parser.py

@@ -9,6 +9,29 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
 UNSET = object()
 
 
+def build_shadow_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Build kwargs for the shadow argument with an ``UNSET`` default.
+
+    Actions that mutate their default in place (append/extend/append_const/
+    count) would crash on the bare ``UNSET`` sentinel, so they are remapped to
+    an equivalent store-style action; the shadow value only needs to flip away
+    from ``UNSET`` when the arg is passed explicitly.
+    """
+    shadow_kwargs = {**kwargs, "default": UNSET}
+    action = kwargs.get("action")
+
+    if action in ("append", "extend"):
+        shadow_kwargs["action"] = "store"
+
+    elif action in ("append_const", "count"):
+        shadow_kwargs["action"] = "store_const"
+
+        if action == "count":
+            shadow_kwargs["const"] = True
+
+    return shadow_kwargs
+
+
 class TrackingNamespace(argparse.Namespace):
     """Proxy that wraps an argparse namespace with explicit keys, which
     can be filtered down to a dict containing only explicitly passed values.
@@ -58,7 +81,7 @@ class TrackingGroup:
     def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action:
         """Add an argument to the real group and to the shadow group."""
         action = self._real.add_argument(*args, **kwargs)
-        default_kwargs = {**kwargs, "default": UNSET}
+        default_kwargs = build_shadow_kwargs(kwargs)
         self._shadow.add_argument(*args, **default_kwargs)
         return action
 
@@ -115,7 +138,7 @@ class TrackingArgumentParser(FlexibleArgumentParser):
     def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action:
         """Add an arg to the parser & the shadow, where the latter has UNSET for the default."""
         action = super().add_argument(*args, **kwargs)
-        shadow_kwargs = {**kwargs, "default": UNSET}
+        shadow_kwargs = build_shadow_kwargs(kwargs)
         self._shadow.add_argument(*args, **shadow_kwargs)
         return action