Move factory/builder methods from datastruct specific "specs.api" to "specs.builders". (#499)

Summary:
Pull Request resolved: #499

Extract functional module out out `specs.api` to new `specs.builders` module.

- Diff/PR should be 100% refactoring without any functional impact.
- `parse_app_handle` could not be moved due to external usage via `specs.api` reference and will be handled separately

#368

Reviewed By: d4l3k

Differential Revision: D36617402

fbshipit-source-id: 563aa15aa354244d09a8eb32da7f5275e7c9ade0

Author: kurman

torchx/cli/cmd_cancel.py

@@ -10,7 +10,7 @@
 
 from torchx.cli.cmd_base import SubCommand
 from torchx.runner import get_runner
-from torchx.specs import api
+from torchx.specs.api import parse_app_handle
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -25,6 +25,6 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
-        _, session_name, _ = api.parse_app_handle(app_handle)
+        _, session_name, _ = parse_app_handle(app_handle)
         runner = get_runner(name=session_name)
         runner.cancel(app_handle)

torchx/cli/cmd_describe.py

@@ -13,7 +13,7 @@
 
 from torchx.cli.cmd_base import SubCommand
 from torchx.runner import get_runner
-from torchx.specs import api
+from torchx.specs.api import parse_app_handle
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -28,7 +28,7 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
-        scheduler, session_name, app_id = api.parse_app_handle(app_handle)
+        scheduler, session_name, app_id = parse_app_handle(app_handle)
         runner = get_runner(name=session_name)
         app = runner.describe(app_handle)
 

torchx/cli/cmd_log.py

@@ -20,7 +20,8 @@
 from torchx.cli.colors import ENDC, GREEN
 from torchx.runner import get_runner, Runner
 from torchx.schedulers.api import Stream
-from torchx.specs.api import is_started, make_app_handle
+from torchx.specs.api import is_started
+from torchx.specs.builders import make_app_handle
 
 logger: logging.Logger = logging.getLogger(__name__)
 

torchx/cli/cmd_status.py

@@ -17,7 +17,7 @@
 from torchx.cli.cmd_base import SubCommand
 from torchx.runner import get_runner
 from torchx.specs import api
-from torchx.specs.api import NONE
+from torchx.specs.api import NONE, parse_app_handle
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -159,7 +159,7 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
-        scheduler, session_name, app_id = api.parse_app_handle(app_handle)
+        scheduler, session_name, app_id = parse_app_handle(app_handle)
         runner = get_runner(name=session_name)
         app_status = runner.status(app_handle)
         filter_roles = parse_list_arg(args.roles)

torchx/components/component_test_base.py

@@ -18,7 +18,7 @@
 import unittest
 from types import ModuleType
 
-from torchx.specs.api import _create_args_parser
+from torchx.specs.builders import _create_args_parser
 from torchx.specs.finder import get_component
 
 

torchx/runner/api.py

@@ -23,8 +23,8 @@
     AppHandle,
     AppStatus,
     CfgVal,
-    from_function,
     make_app_handle,
+    materialize_appdef,
     parse_app_handle,
     runopts,
     UnknownAppException,
@@ -161,7 +161,7 @@ def dryrun_component(
         component, but just returns what "would" have run.
         """
         component_def = get_component(component)
-        app = from_function(
+        app = materialize_appdef(
             component_def.fn,
             component_args,
             self._component_defaults.get(component, None),

torchx/specs/__init__.py

@@ -26,18 +26,15 @@
     BindMount,
     CfgVal,
     DeviceMount,
-    from_function,
     get_type_name,
     InvalidRunConfigException,
     is_terminal,
     macros,
-    make_app_handle,
     MalformedAppHandleException,
     MISSING,
     NONE,
     NULL_RESOURCE,
     parse_app_handle,
-    parse_mounts,
     ReplicaState,
     ReplicaStatus,
     Resource,
@@ -50,7 +47,7 @@
     UnknownSchedulerException,
     VolumeMount,
 )
-
+from .builders import make_app_handle, materialize_appdef, parse_mounts  # noqa
 
 GiB: int = 1024
 

torchx/specs/api.py

@@ -5,9 +5,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import argparse
 import copy
-import inspect
 import json
 from dataclasses import asdict, dataclass, field
 from enum import Enum
@@ -28,15 +26,7 @@
 )
 
 import yaml
-from torchx.specs.file_linter import get_fn_docstring, TorchXArgumentHelpFormatter
-from torchx.util.types import (
-    decode_from_string,
-    decode_optional,
-    get_argparse_param_type,
-    is_bool,
-    is_primitive,
-    to_dict,
-)
+from torchx.util.types import to_dict
 
 
 # ========================================
@@ -832,10 +822,6 @@ def __init__(self, app_handle: "AppHandle") -> None:
         )
 
 
-def make_app_handle(scheduler_backend: str, session_name: str, app_id: str) -> str:
-    return f"{scheduler_backend}://{session_name}/{app_id}"
-
-
 def parse_app_handle(app_handle: AppHandle) -> Tuple[str, str, str]:
     """
     parses the app handle into ```(scheduler_backend, session_name, and app_id)```
@@ -851,215 +837,3 @@ def parse_app_handle(app_handle: AppHandle) -> Tuple[str, str, str]:
         raise MalformedAppHandleException(app_handle)
     gd = match.groupdict()
     return gd["scheduler_backend"], gd["session_name"], gd["app_id"]
-
-
-def _create_args_parser(
-    cmpnt_fn: Callable[..., AppDef], cmpnt_defaults: Optional[Dict[str, str]] = None
-) -> argparse.ArgumentParser:
-    parameters = inspect.signature(cmpnt_fn).parameters
-    function_desc, args_desc = get_fn_docstring(cmpnt_fn)
-    script_parser = argparse.ArgumentParser(
-        prog=f"torchx run <run args...> {cmpnt_fn.__name__} ",
-        description=function_desc,
-        formatter_class=TorchXArgumentHelpFormatter,
-        # enables components to have "h" as a parameter
-        # otherwise argparse by default adds -h/--help as the help argument
-        # we still add --help but reserve "-"h" to be used as a component argument
-        add_help=False,
-    )
-    # add help manually since we disabled auto help to allow "h" in component arg
-    script_parser.add_argument(
-        "--help",
-        action="help",
-        default=argparse.SUPPRESS,
-        help="show this help message and exit",
-    )
-
-    class _reminder_action(argparse.Action):
-        def __call__(
-            self,
-            parser: argparse.ArgumentParser,
-            namespace: argparse.Namespace,
-            values: Any,
-            option_string: Optional[str] = None,
-        ) -> None:
-            setattr(
-                namespace,
-                self.dest,
-                (self.default or "").split() if len(values) == 0 else values,
-            )
-
-    for param_name, parameter in parameters.items():
-        param_desc = args_desc[parameter.name]
-        args: Dict[str, Any] = {
-            "help": param_desc,
-            "type": get_argparse_param_type(parameter),
-        }
-        # set defaults specified in the component function declaration
-        if parameter.default != inspect.Parameter.empty:
-            if is_bool(type(parameter.default)):
-                args["default"] = str(parameter.default)
-            else:
-                args["default"] = parameter.default
-
-        # set defaults supplied directly to this method (overwrites the declared defaults)
-        # the defaults are given as str (as option values passed from CLI) since
-        # these are typically read from .torchxconfig
-        if cmpnt_defaults and param_name in cmpnt_defaults:
-            args["default"] = cmpnt_defaults[param_name]
-
-        if parameter.kind == inspect._ParameterKind.VAR_POSITIONAL:
-            args["nargs"] = argparse.REMAINDER
-            args["action"] = _reminder_action
-            script_parser.add_argument(param_name, **args)
-        else:
-            arg_names = [f"--{param_name}"]
-            if len(param_name) == 1:
-                arg_names = [f"-{param_name}"] + arg_names
-            if "default" not in args:
-                args["required"] = True
-            script_parser.add_argument(*arg_names, **args)
-    return script_parser
-
-
-# FIXME this function invokes the component fn to materialize the AppDef
-#      rename to something more appropriate like - materialize_appdef or eval_component_fn
-#      seee: https://github.com/pytorch/torchx/issues/368
-def from_function(
-    cmpnt_fn: Callable[..., AppDef],
-    cmpnt_args: List[str],
-    cmpnt_defaults: Optional[Dict[str, str]] = None,
-) -> AppDef:
-    """
-    Creates an application by running user defined ``app_fn``.
-
-    ``app_fn`` has the following restrictions:
-        * Name must be ``app_fn``
-        * All arguments should be annotated
-        * Supported argument types:
-            - primitive: int, str, float
-            - Dict[primitive, primitive]
-            - List[primitive]
-            - Optional[Dict[primitive, primitive]]
-            - Optional[List[primitive]]
-        * ``app_fn`` can define a vararg (*arg) at the end
-        * There should be a docstring for the function that defines
-            All arguments in a google-style format
-        * There can be default values for the function arguments.
-        * The return object must be ``AppDef``
-
-    Args:
-        cmpnt_fn: Component function
-        cmpnt_args: Function args
-        cmpnt_defaults: Additional default values for parameters of ``app_fn``
-                          (overrides the defaults set on the fn declaration)
-    Returns:
-        An application spec
-    """
-
-    script_parser = _create_args_parser(cmpnt_fn, cmpnt_defaults)
-    parsed_args = script_parser.parse_args(cmpnt_args)
-
-    function_args = []
-    var_arg = []
-    kwargs = {}
-
-    parameters = inspect.signature(cmpnt_fn).parameters
-    for param_name, parameter in parameters.items():
-        arg_value = getattr(parsed_args, param_name)
-        parameter_type = parameter.annotation
-        parameter_type = decode_optional(parameter_type)
-        if is_bool(parameter_type):
-            arg_value = arg_value.lower() == "true"
-        elif not is_primitive(parameter_type):
-            arg_value = decode_from_string(arg_value, parameter_type)
-        if parameter.kind == inspect.Parameter.VAR_POSITIONAL:
-            var_arg = arg_value
-        elif parameter.kind == inspect.Parameter.KEYWORD_ONLY:
-            kwargs[param_name] = arg_value
-        elif parameter.kind == inspect.Parameter.VAR_KEYWORD:
-            raise TypeError("**kwargs are not supported for component definitions")
-        else:
-            function_args.append(arg_value)
-    if len(var_arg) > 0 and var_arg[0] == "--":
-        var_arg = var_arg[1:]
-
-    return cmpnt_fn(*function_args, *var_arg, **kwargs)
-
-
-_MOUNT_OPT_MAP: Mapping[str, str] = {
-    "type": "type",
-    "destination": "dst",
-    "dst": "dst",
-    "target": "dst",
-    "read_only": "readonly",
-    "readonly": "readonly",
-    "source": "src",
-    "src": "src",
-    "perm": "perm",
-}
-
-
-def parse_mounts(opts: List[str]) -> List[Union[BindMount, VolumeMount, DeviceMount]]:
-    """
-    parse_mounts parses a list of options into typed mounts following a similar
-    format to Dockers bind mount.
-
-    Multiple mounts can be specified in the same list. ``type`` must be
-    specified first in each.
-
-    Ex:
-        type=bind,src=/host,dst=/container,readonly,[type=bind,src=...,dst=...]
-
-    Supported types:
-        BindMount: type=bind,src=<host path>,dst=<container path>[,readonly]
-        VolumeMount: type=volume,src=<name/id>,dst=<container path>[,readonly]
-        DeviceMount: type=device,src=/dev/<dev>[,dst=<container path>][,perm=rwm]
-    """
-    mount_opts = []
-    cur = {}
-    for opt in opts:
-        key, _, val = opt.partition("=")
-        if key not in _MOUNT_OPT_MAP:
-            raise KeyError(
-                f"unknown mount option {key}, must be one of {list(_MOUNT_OPT_MAP.keys())}"
-            )
-        key = _MOUNT_OPT_MAP[key]
-        if key == "type":
-            cur = {}
-            mount_opts.append(cur)
-        elif len(mount_opts) == 0:
-            raise KeyError("type must be specified first")
-        cur[key] = val
-
-    mounts = []
-    for opts in mount_opts:
-        typ = opts.get("type")
-        if typ == MountType.BIND:
-            mounts.append(
-                BindMount(
-                    src_path=opts["src"],
-                    dst_path=opts["dst"],
-                    read_only="readonly" in opts,
-                )
-            )
-        elif typ == MountType.VOLUME:
-            mounts.append(
-                VolumeMount(
-                    src=opts["src"], dst_path=opts["dst"], read_only="readonly" in opts
-                )
-            )
-        elif typ == MountType.DEVICE:
-            src = opts["src"]
-            dst = opts.get("dst", src)
-            perm = opts.get("perm", "rwm")
-            for c in perm:
-                if c not in "rwm":
-                    raise ValueError(
-                        f"{c} is not a valid permission flags must one of r,w,m"
-                    )
-            mounts.append(DeviceMount(src_path=src, dst_path=dst, permissions=perm))
-        else:
-            valid = list(str(item.value) for item in MountType)
-            raise ValueError(f"invalid mount type {repr(typ)}, must be one of {valid}")
-    return mounts