Make torchx error out when the same arg is passed in twice

Differential Revision: D55216738

Pull Request resolved: #854

Author: andywag

torchx/cli/argparse_util.py

@@ -6,19 +6,25 @@
 
 # pyre-strict
 
+import logging
+import sys
 from argparse import Action, ArgumentParser, Namespace
-from typing import Any, Dict, Optional, Sequence, Text
+from typing import Any, Dict, List, Optional, Sequence, Set, Text
 
 from torchx.runner import config
 
+logger: logging.Logger = logging.getLogger(__name__)
 
-class _torchxconfig(Action):
+
+class torchxconfig(Action):
     """
     Custom argparse action that loads default torchx CLI options
     from .torchxconfig file.
 
     """
 
+    called_args: Set[str] = set()
+
     # since this action is used for each argparse argument
     # load the config section for the subcmd once
     _subcmd_configs: Dict[str, Dict[str, str]] = {}
@@ -66,13 +72,18 @@ def __call__(
         values: Any,  # pyre-ignore[2] declared as Any in superclass Action
         option_string: Optional[str] = None,
     ) -> None:
+        if option_string is not None:
+            if option_string in self.called_args:
+                logger.error(f"{option_string} is specified more than once")
+                sys.exit(1)
+            self.called_args.add(option_string)
         setattr(namespace, self.dest, values)
 
 
 # argparse takes the action as a Type[Action] so we can't have custom constructors
 # hence for each subcommand we need to subclass the base _torchxconfig Action
 # this is also how store_true and store_false builtin actions are implemented in argparse
-class torchxconfig_run(_torchxconfig):
+class torchxconfig_run(torchxconfig):
     """
     Custom action that gets the default argument from .torchxconfig.
     """
@@ -94,3 +105,25 @@ def __init__(
             option_strings=option_strings,
             **kwargs,
         )
+
+
+class ArgOnceAction(Action):
+    """
+    Custom argparse action only allows argument to be specified once
+    """
+
+    called_args: Set[str] = set()
+
+    def __call__(
+        self,
+        parser: ArgumentParser,
+        namespace: Namespace,
+        values: List[str],
+        option_string: Optional[str] = None,
+    ) -> None:
+        if option_string is not None:
+            if option_string in self.called_args:
+                logger.error(f"{option_string} is specified more than once")
+                sys.exit(1)
+            self.called_args.add(option_string)
+        setattr(namespace, self.dest, values)

torchx/cli/cmd_run.py

@@ -17,7 +17,7 @@
 from typing import Dict, List, Optional, Tuple
 
 import torchx.specs as specs
-from torchx.cli.argparse_util import torchxconfig_run
+from torchx.cli.argparse_util import ArgOnceAction, torchxconfig_run
 from torchx.cli.cmd_base import SubCommand
 from torchx.cli.cmd_log import get_logs
 from torchx.runner import config, get_runner, Runner
@@ -133,6 +133,7 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
             "-cfg",
             "--scheduler_args",
             type=str,
+            action=ArgOnceAction,
             help="Arguments to pass to the scheduler (Ex:`cluster=foo,user=bar`)."
             " For a list of scheduler run options run: `torchx runopts`",
         )
@@ -165,6 +166,7 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
         subparser.add_argument(
             "--parent_run_id",
             type=str,
+            action=ArgOnceAction,
             help="optional parent run ID that this run belongs to."
             " It can be used to group runs for experiment tracking purposes",
         )

torchx/cli/test/argparse_util_test.py

@@ -19,7 +19,8 @@
 class ArgparseUtilTest(TestWithTmpDir):
     def setUp(self) -> None:
         super().setUp()
-        argparse_util._torchxconfig._subcmd_configs.clear()
+        argparse_util.torchxconfig._subcmd_configs.clear()
+        argparse_util.torchxconfig.called_args = set()
 
     def test_torchxconfig_action(self) -> None:
         with mock.patch(DEFAULT_CONFIG_DIRS, [str(self.tmpdir)]):

torchx/cli/test/cmd_run_test.py

@@ -20,6 +20,8 @@
 from typing import Generator
 from unittest.mock import MagicMock, patch
 
+from torchx.cli.argparse_util import ArgOnceAction, torchxconfig
+
 from torchx.cli.cmd_run import _parse_component_name_and_args, CmdBuiltins, CmdRun
 from torchx.schedulers.local_scheduler import SignalException
 
@@ -45,6 +47,28 @@ def setUp(self) -> None:
 
     def tearDown(self) -> None:
         shutil.rmtree(self.tmpdir, ignore_errors=True)
+        ArgOnceAction.called_args = set()
+        torchxconfig.called_args = set()
+
+    def test_run_with_multiple_scheduler_args(self) -> None:
+
+        args = ["--scheduler_args", "first_args", "--scheduler_args", "second_args"]
+        with self.assertRaises(SystemExit) as cm:
+            self.parser.parse_args(args)
+        self.assertEqual(cm.exception.code, 1)
+
+    def test_run_with_multiple_schedule_args(self) -> None:
+
+        args = [
+            "--scheduler",
+            "local_conda",
+            "--scheduler",
+            "local_cwd",
+        ]
+
+        with self.assertRaises(SystemExit) as cm:
+            self.parser.parse_args(args)
+        self.assertEqual(cm.exception.code, 1)
 
     def test_run_with_user_conf_abs_path(self) -> None:
         args = self.parser.parse_args(