(torchx/cli) allow --workspace to be configured through .torchxconfig (#397)

Summary:
Pull Request resolved: #397

For `torchx run`, allows `--workspace` to be configured through `.torchxconfig`. Adds a custom `argparse.Action` that can be used to read and default other CLI options through `.torchxconfig`. Only hooked it up to `workspace` for now since there is some cleanup to be done to hook it up for `component`, `component args` (we get this from a different section not `[cli:run]`))

Reviewed By: divchenko

Differential Revision: D34317215

fbshipit-source-id: 878f9d75002ab352f4d08d15802cfa405995f544

Author: Kiuk Chung

torchx/cli/argparse_util.py

@@ -0,0 +1,99 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from argparse import Action, ArgumentParser, Namespace
+from pathlib import Path
+from typing import Any, Dict, Optional, Sequence, Text
+
+from torchx.runner import config
+
+
+CONFIG_DIRS = [str(Path.home()), str(Path.cwd())]
+
+
+class _torchxconfig(Action):
+    """
+    Custom argparse action that loads default torchx CLI options
+    from .torchxconfig file.
+
+    """
+
+    # since this action is used for each argparse argument
+    # load the config section for the subcmd once
+    _subcmd_configs: Dict[str, Dict[str, str]] = {}
+
+    def __init__(
+        self,
+        subcmd: str,
+        dest: str,
+        option_strings: Sequence[Text],
+        required: bool = False,
+        # pyre-ignore[2] declared as Any in superclass Action
+        default: Any = None,
+        **kwargs: Any,
+    ) -> None:
+        cfg = self._subcmd_configs.setdefault(
+            subcmd,
+            config.get_configs(
+                prefix="cli",
+                name=subcmd,
+                dirs=CONFIG_DIRS,
+            ),
+        )
+
+        # if found in .torchxconfig make it the default for this argument
+        # otherwise use the default defined from add_argument(...)
+        default = cfg.get(dest, default)
+
+        # ``required`` means that it NEEDS to be present  in the CLI args
+        # if we found it in .torchxconfig then we don't "require" it to be
+        # in the CLI args so set it to False
+        if default:
+            required = False
+
+        super().__init__(
+            dest=dest,
+            default=default,
+            option_strings=option_strings,
+            required=required,
+            **kwargs,
+        )
+
+    def __call__(
+        self,
+        parser: ArgumentParser,
+        namespace: Namespace,
+        values: Any,  # pyre-ignore[2] declared as Any in superclass Action
+        option_string: Optional[str] = None,
+    ) -> None:
+        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):
+    """
+    Custom action that gets the default argument from .torchxconfig.
+    """
+
+    def __init__(
+        self,
+        dest: str,
+        option_strings: Sequence[Text],
+        required: bool = False,
+        # pyre-ignore[2] declared as Any in superclass Action
+        default: Any = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            "run",
+            dest=dest,
+            default=default,
+            required=required,
+            option_strings=option_strings,
+            **kwargs,
+        )

torchx/cli/cmd_run.py

@@ -16,6 +16,7 @@
 
 import torchx.specs as specs
 from pyre_extensions import none_throws
+from torchx.cli.argparse_util import CONFIG_DIRS, torchxconfig_run
 from torchx.cli.cmd_base import SubCommand
 from torchx.cli.cmd_log import get_logs
 from torchx.runner import Runner, config, get_runner
@@ -36,7 +37,6 @@
     "missing component name, either provide it from the CLI or in .torchxconfig"
 )
 
-CONFIG_DIRS = [str(Path.home()), str(Path.cwd())]
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -185,6 +185,7 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
             "--workspace",
             "--buck-target",
             default=f"file://{Path.cwd()}",
+            action=torchxconfig_run,
             help="local workspace to build/patch (buck-target of main binary if using buck)",
         )
         subparser.add_argument(

torchx/cli/conf_helpers.py

@@ -0,0 +1,146 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import shutil
+import tempfile
+import unittest
+from argparse import ArgumentParser
+from pathlib import Path
+from unittest import mock
+
+from torchx.cli import argparse_util
+from torchx.cli.argparse_util import torchxconfig_run
+
+
+CONFIG_DIRS = "torchx.cli.argparse_util.CONFIG_DIRS"
+
+
+class ArgparseUtilTest(unittest.TestCase):
+    def _write(self, filename: str, content: str) -> Path:
+        f = Path(self.test_dir) / filename
+        f.parent.mkdir(parents=True, exist_ok=True)
+        with open(f, "w") as fp:
+            fp.write(content)
+        return f
+
+    def setUp(self) -> None:
+        self.test_dir = tempfile.mkdtemp(prefix="torchx_argparse_util_test")
+        argparse_util._torchxconfig._subcmd_configs.clear()
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.test_dir)
+
+    def test_torchxconfig_action(self) -> None:
+        with mock.patch(CONFIG_DIRS, [self.test_dir]):
+            self._write(
+                ".torchxconfig",
+                """
+[cli:run]
+workspace = baz
+                """,
+            )
+
+            parser = ArgumentParser()
+
+            subparsers = parser.add_subparsers()
+            run_parser = subparsers.add_parser("run")
+
+            run_parser.add_argument(
+                "--workspace",
+                default="foo",
+                type=str,
+                action=torchxconfig_run,
+            )
+
+            # arguments specified in CLI should take outmost precedence
+            args = parser.parse_args(["run", "--workspace", "bar"])
+            self.assertEqual("bar", args.workspace)
+
+            # if not specified in CLI, then grab it from .torchxconfig
+            args = parser.parse_args(["run"])
+            self.assertEqual("baz", args.workspace)
+
+    def test_torchxconfig_action_argparse_default(self) -> None:
+        with mock.patch(CONFIG_DIRS, [self.test_dir]):
+            self._write(
+                ".torchxconfig",
+                """
+[cli:run]
+                """,
+            )
+
+            parser = ArgumentParser()
+
+            subparsers = parser.add_subparsers()
+            run_parser = subparsers.add_parser("run")
+
+            run_parser.add_argument(
+                "--workspace",
+                default="foo",
+                type=str,
+                action=torchxconfig_run,
+            )
+
+            # if not found in .torchxconfig should use argparse default
+            args = parser.parse_args(["run"])
+            self.assertEqual("foo", args.workspace)
+
+    def test_torchxconfig_action_required(self) -> None:
+        with mock.patch(CONFIG_DIRS, [self.test_dir]):
+            self._write(
+                ".torchxconfig",
+                """
+[cli:run]
+workspace = bazz
+                """,
+            )
+
+            parser = ArgumentParser()
+
+            subparsers = parser.add_subparsers()
+            run_parser = subparsers.add_parser("run")
+
+            run_parser.add_argument(
+                "--workspace",
+                required=True,
+                type=str,
+                action=torchxconfig_run,
+            )
+
+            # arguments specified in CLI should take outmost precedence
+            args = parser.parse_args(["run", "--workspace", "bar"])
+            self.assertEqual("bar", args.workspace)
+
+            # if not specified in CLI, then grab it from .torchxconfig
+            args = parser.parse_args(["run"])
+            self.assertEqual("bazz", args.workspace)
+
+    def test_torchxconfig_action_aliases(self) -> None:
+        # for aliases, the config file needs to declare the original arg
+        with mock.patch(CONFIG_DIRS, [self.test_dir]):
+            self._write(
+                ".torchxconfig",
+                """
+[cli:run]
+workspace = baz
+                """,
+            )
+
+            parser = ArgumentParser()
+
+            subparsers = parser.add_subparsers()
+            run_parser = subparsers.add_parser("run")
+
+            run_parser.add_argument(
+                "--workspace",
+                "--buck-target",
+                type=str,
+                required=True,
+                action=torchxconfig_run,
+            )
+
+            args = parser.parse_args(["run"])
+            self.assertEqual("baz", args.workspace)

torchx/cli/test/argparse_util_test.py

@@ -326,6 +326,31 @@ def strip_prefix(section_name: str) -> Optional[str]:
     return sections
 
 
+def get_configs(
+    prefix: str,
+    name: str,
+    dirs: Optional[List[str]],
+) -> Dict[str, str]:
+    """
+    Gets all the config values in the section ``["{prefix}:{name}"]``.
+    Or an empty map if the section does not exist.
+
+    Example:
+
+    ::
+
+     # for config file:
+     # [foo:bar]
+     # baz = 1
+
+     get_configs(prefix="foo", name="bar") # returns {"baz": "1"}
+     get_config(prefix="foo", name="barr") # returns {}
+
+    """
+    sections = load_sections(prefix, dirs)
+    return sections.get(name, {})
+
+
 def get_config(
     prefix: str,
     name: str,
@@ -350,9 +375,7 @@ def get_config(
      get_config(prefix="fooo", name="bar", key="baz") == None
 
     """
-    sections = load_sections(prefix, dirs)
-    section = sections.get(name, {})
-    return section.get(key, None)
+    return get_configs(prefix, name, dirs).get(key, None)
 
 
 def find_configs(dirs: Optional[Iterable[str]] = None) -> List[str]: