diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py index 437fbd50a..d2e7e4034 100644 --- a/tests/utils/test_tracking_parser.py +++ b/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() diff --git a/vllm_omni/utils/tracking_parser.py b/vllm_omni/utils/tracking_parser.py index c6dc5fee7..64849a5aa 100644 --- a/vllm_omni/utils/tracking_parser.py +++ b/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