torchx-nightly 2024.1.6__py3-none-any.whl → 2025.12.24__py3-none-any.whl

torchx/specs/builders.py

@@ -4,27 +4,46 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import argparse
 import inspect
-from typing import Any, Callable, Dict, List, Mapping, Optional, Union
+import os
+from argparse import Namespace
+from typing import Any, Callable, Dict, List, Mapping, NamedTuple, Optional, Union
 
 from torchx.specs.api import BindMount, MountType, VolumeMount
 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,
-)
+from torchx.util.types import decode, decode_optional, get_argparse_param_type, is_bool
 
 from .api import AppDef, DeviceMount
 
 
+class ComponentArgs(NamedTuple):
+    """Parsed component function arguments"""
+
+    positional_args: dict[str, Any]
+    var_args: list[str]
+    kwargs: dict[str, Any]
+
+
 def _create_args_parser(
-    cmpnt_fn: Callable[..., AppDef], cmpnt_defaults: Optional[Dict[str, str]] = None
+    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_defaults: Optional[Dict[str, str]] = None,
+    config: Optional[Dict[str, Any]] = None,
 ) -> argparse.ArgumentParser:
     parameters = inspect.signature(cmpnt_fn).parameters
+    return _create_args_parser_from_parameters(
+        cmpnt_fn, parameters, cmpnt_defaults, config
+    )
+
+
+def _create_args_parser_from_parameters(
+    cmpnt_fn: Callable[..., AppDef],
+    parameters: Mapping[str, inspect.Parameter],
+    cmpnt_defaults: Optional[Dict[str, str]] = None,
+    config: Optional[Dict[str, Any]] = None,
+) -> argparse.ArgumentParser:
     function_desc, args_desc = get_fn_docstring(cmpnt_fn)
     script_parser = argparse.ArgumentParser(
         prog=f"torchx run <run args...> {cmpnt_fn.__name__} ",
@@ -85,15 +104,144 @@ def _create_args_parser(
             if len(param_name) == 1:
                 arg_names = [f"-{param_name}"] + arg_names
             if "default" not in args:
-                args["required"] = True
+                if (config and param_name not in config) or not config:
+                    args["required"] = True
+
             script_parser.add_argument(*arg_names, **args)
     return script_parser
 
 
+def _merge_config_values_with_args(
+    parsed_args: argparse.Namespace, config: Dict[str, Any]
+) -> None:
+    for key, val in config.items():
+        if key in parsed_args:
+            setattr(parsed_args, key, val)
+
+
+def parse_args(
+    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_args: List[str],
+    cmpnt_defaults: Optional[Dict[str, Any]] = None,
+    config: Optional[Dict[str, Any]] = None,
+) -> Namespace:
+    """
+    Parse passed arguments, defaults, and config values into a namespace for
+    a component function.
+
+    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)
+    config: Optional dict containing additional configuration for the component from a passed config file
+
+    Returns:
+    A Namespace object with the args, defaults, and config values incorporated.
+    """
+
+    script_parser = _create_args_parser(cmpnt_fn, cmpnt_defaults, config)
+    parsed_args = script_parser.parse_args(cmpnt_args)
+    if config:
+        _merge_config_values_with_args(parsed_args, config)
+
+    return parsed_args
+
+
+def component_args_from_str(
+    cmpnt_fn: Callable[..., AppDef],
+    cmpnt_args: list[str],
+    cmpnt_args_defaults: Optional[Dict[str, Any]] = None,
+    config: Optional[Dict[str, Any]] = None,
+) -> ComponentArgs:
+    """
+    Parses and decodes command-line arguments for a component function.
+
+    This function takes a component function and its arguments, parses them using argparse,
+    and decodes the arguments into their expected types based on the function's signature.
+    It separates positional arguments, variable positional arguments (*args), and keyword-only arguments.
+
+    Args:
+        cmpnt_fn: The component function whose arguments are to be parsed and decoded.
+        cmpnt_args: List of command-line arguments to be parsed. Supports both space separated and '=' separated arguments.
+        cmpnt_args_defaults: Optional dictionary of default values for the component function's parameters.
+        config: Optional dictionary containing additional configuration values.
+
+    Returns:
+        ComponentArgs representing the input args to a component function containing:
+            - positional_args: Dictionary of positional and positional-or-keyword arguments.
+            - var_args: List of variable positional arguments (*args).
+            - kwargs: Dictionary of keyword-only arguments.
+
+        Usage:
+
+        .. doctest::
+            from torchx.specs.api import AppDef
+            from torchx.specs.builders import component_args_from_str
+
+            def example_component_fn(foo: str, *args: str, bar: str = "asdf") -> AppDef:
+                return AppDef(name="example")
+
+            # Supports space separated arguments
+            args = ["--foo", "fooval", "--bar", "barval", "arg1", "arg2"]
+            parsed_args = component_args_from_str(example_component_fn, args)
+
+            assert parsed_args.positional_args == {"foo": "fooval"}
+            assert parsed_args.var_args == ["arg1", "arg2"]
+            assert parsed_args.kwargs == {"bar": "barval"}
+
+            # Supports '=' separated arguments
+            args = ["--foo=fooval", "--bar=barval", "arg1", "arg2"]
+            parsed_args = component_args_from_str(example_component_fn, args)
+
+            assert parsed_args.positional_args == {"foo": "fooval"}
+            assert parsed_args.var_args == ["arg1", "arg2"]
+            assert parsed_args.kwargs == {"bar": "barval"}
+
+
+    """
+    parsed_args: Namespace = parse_args(
+        cmpnt_fn, cmpnt_args, cmpnt_args_defaults, config
+    )
+
+    positional_args = {}
+    var_args = []
+    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 (
+            parameter_type != arg_value.__class__
+            and parameter.kind != inspect.Parameter.VAR_POSITIONAL
+        ):
+            arg_value = decode(arg_value, parameter_type)
+        if parameter.kind == inspect.Parameter.VAR_POSITIONAL:
+            var_args = arg_value
+        elif parameter.kind == inspect.Parameter.KEYWORD_ONLY:
+            kwargs[param_name] = arg_value
+        elif parameter.kind == inspect.Parameter.VAR_KEYWORD:
+            raise TypeError(
+                f"component fn param `{param_name}` is a '**kwargs' which is not supported; consider changing the "
+                f"type to a dict or explicitly declare the params"
+            )
+        else:
+            # POSITIONAL or POSITIONAL_OR_KEYWORD
+            positional_args[param_name] = arg_value
+
+    if len(var_args) > 0 and var_args[0] == "--":
+        var_args = var_args[1:]
+
+    return ComponentArgs(positional_args, var_args, kwargs)
+
+
 def materialize_appdef(
     cmpnt_fn: Callable[..., AppDef],
     cmpnt_args: List[str],
-    cmpnt_defaults: Optional[Dict[str, str]] = None,
+    cmpnt_defaults: Optional[Dict[str, Any]] = None,
+    config: Optional[Dict[str, Any]] = None,
 ) -> AppDef:
     """
     Creates an application by running user defined ``app_fn``.
@@ -118,38 +266,25 @@ def materialize_appdef(
         cmpnt_args: Function args
         cmpnt_defaults: Additional default values for parameters of ``app_fn``
                           (overrides the defaults set on the fn declaration)
+        config: Optional dict containing additional configuration for the component from a passed config file
     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 = {}
+    component_args: ComponentArgs = component_args_from_str(
+        cmpnt_fn, cmpnt_args, cmpnt_defaults, config
+    )
+    positional_arg_values = list(component_args.positional_args.values())
+    appdef = cmpnt_fn(
+        *positional_arg_values, *component_args.var_args, **component_args.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 and 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:]
+    if not isinstance(appdef, AppDef):
+        raise TypeError(
+            f"Expected a component that returns `AppDef`, but got `{type(appdef)}`"
+        )
 
-    return cmpnt_fn(*function_args, *var_arg, **kwargs)
+    return appdef
 
 
 def make_app_handle(scheduler_backend: str, session_name: str, app_id: str) -> str:
@@ -205,9 +340,12 @@ def parse_mounts(opts: List[str]) -> List[Union[BindMount, VolumeMount, DeviceMo
     for opts in mount_opts:
         typ = opts.get("type")
         if typ == MountType.BIND:
+            src_path = opts["src"]
+            if src_path.startswith("~"):
+                src_path = os.path.expanduser(src_path)
             mounts.append(
                 BindMount(
-                    src_path=opts["src"],
+                    src_path=src_path,
                     dst_path=opts["dst"],
                     read_only="readonly" in opts,
                 )

torchx/specs/file_linter.py

@@ -5,12 +5,15 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 import abc
 import argparse
 import ast
 import inspect
+import sys
 from dataclasses import dataclass
-from typing import Callable, cast, Dict, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from docstring_parser import parse
 from torchx.util.io import read_conf_file
@@ -29,7 +32,11 @@ def _get_default_arguments_descriptions(fn: Callable[..., object]) -> Dict[str,
     return args_decs
 
 
-class TorchXArgumentHelpFormatter(argparse.HelpFormatter):
+class TorchXArgumentHelpFormatter(
+    argparse.RawDescriptionHelpFormatter,
+    argparse.ArgumentDefaultsHelpFormatter,
+    argparse.MetavarTypeHelpFormatter,
+):
     """Help message formatter which adds default values and required to argument help.
 
     If the argument is required, the class appends `(required)` at the end of the help message.
@@ -68,7 +75,7 @@ def get_fn_docstring(fn: Callable[..., object]) -> Tuple[str, Dict[str, str]]:
             if the description
     """
     default_fn_desc = f"""{fn.__name__} TIP: improve this help string by adding a docstring
-to your component (see: https://pytorch.org/torchx/latest/component_best_practices.html)"""
+to your component (see: https://meta-pytorch.org/torchx/latest/component_best_practices.html)"""
     args_description = _get_default_arguments_descriptions(fn)
     func_description = inspect.getdoc(fn)
     if not func_description:
@@ -79,7 +86,7 @@ to your component (see: https://pytorch.org/torchx/latest/component_best_practic
             args_description[param.arg_name] = param.description
     short_func_description = docstring.short_description or default_fn_desc
     if docstring.long_description:
-        short_func_description += " ..."
+        short_func_description += "\n" + docstring.long_description
     return (short_func_description or default_fn_desc, args_description)
 
 
@@ -92,7 +99,7 @@ class LinterMessage:
     severity: str = "error"
 
 
-class TorchxFunctionValidator(abc.ABC):
+class ComponentFunctionValidator(abc.ABC):
     @abc.abstractmethod
     def validate(self, app_specs_func_def: ast.FunctionDef) -> List[LinterMessage]:
         """
@@ -110,7 +117,55 @@ class TorchxFunctionValidator(abc.ABC):
         )
 
 
-class TorchxFunctionArgsValidator(TorchxFunctionValidator):
+def OK() -> list[LinterMessage]:
+    return []  # empty linter error means validation passes
+
+
+def is_primitive(arg: ast.expr) -> bool:
+    # whether the arg is a primitive type (e.g. int, float, str, bool)
+    return isinstance(arg, ast.Name)
+
+
+def get_generic_type(arg: ast.expr) -> ast.expr:
+    # returns the slice expr of a subscripted type
+    # `arg` must be an instance of ast.Subscript (caller checks)
+    # in this validator's context, this is the generic type of a container type
+    # e.g. for Optional[str] returns the expr for str
+
+    assert isinstance(arg, ast.Subscript)  # e.g. arg = C[T]
+
+    if isinstance(arg.slice, ast.Index):  # python>=3.10
+        return arg.slice.value
+    else:  # python-3.9
+        return arg.slice
+
+
+def get_optional_type(arg: ast.expr) -> Optional[ast.expr]:
+    """
+    Returns the type parameter ``T`` of ``Optional[T]`` or ``None`` if `arg``
+    is not an ``Optional``. Handles both:
+        1. ``typing.Optional[T]`` (python<3.10)
+        2.  ``T | None`` or ``None | T`` (python>=3.10 - PEP 604)
+    """
+    # case 1: 'a: Optional[T]'
+    if isinstance(arg, ast.Subscript) and arg.value.id == "Optional":
+        return get_generic_type(arg)
+
+    # case 2: 'a: T | None' or 'a: None | T'
+    if sys.version_info >= (3, 10):  # PEP 604 introduced in python-3.10
+        if isinstance(arg, ast.BinOp) and isinstance(arg.op, ast.BitOr):
+            if isinstance(arg.right, ast.Constant) and arg.right.value is None:
+                return arg.left
+            if isinstance(arg.left, ast.Constant) and arg.left.value is None:
+                return arg.right
+
+    # case 3: is not optional
+    return None
+
+
+class ArgTypeValidator(ComponentFunctionValidator):
+    """Validates component function's argument types."""
+
     def validate(self, app_specs_func_def: ast.FunctionDef) -> List[LinterMessage]:
         linter_errors = []
         for arg_def in app_specs_func_def.args.args:
@@ -127,53 +182,73 @@ class TorchxFunctionArgsValidator(TorchxFunctionValidator):
         return linter_errors
 
     def _validate_arg_def(
-        self, function_name: str, arg_def: ast.arg
+        self, function_name: str, arg: ast.arg
     ) -> List[LinterMessage]:
-        if not arg_def.annotation:
-            return [
-                self._gen_linter_message(
-                    f"Arg {arg_def.arg} missing type annotation", arg_def.lineno
-                )
-            ]
-        if isinstance(arg_def.annotation, ast.Name):
+        arg_type = arg.annotation  # type hint
+
+        def ok() -> list[LinterMessage]:
+            # return value when validation passes (e.g. no linter errors)
             return []
-        complex_type_def = cast(ast.Subscript, none_throws(arg_def.annotation))
-        if complex_type_def.value.id == "Optional":
-            # ast module in python3.9 does not have ast.Index wrapper
-            if isinstance(complex_type_def.slice, ast.Index):
-                complex_type_def = complex_type_def.slice.value
-            else:
-                complex_type_def = complex_type_def.slice
-            # Check if type is Optional[primitive_type]
-            if isinstance(complex_type_def, ast.Name):
-                return []
-        # Check if type is Union[Dict,List]
-        type_name = complex_type_def.value.id
-        if type_name != "Dict" and type_name != "List":
-            desc = (
-                f"`{function_name}` allows only Dict, List as complex types."
-                f"Argument `{arg_def.arg}` has: {type_name}"
-            )
-            return [self._gen_linter_message(desc, arg_def.lineno)]
-        linter_errors = []
-        # ast module in python3.9 does not have objects wrapped in ast.Index
-        if isinstance(complex_type_def.slice, ast.Index):
-            sub_type = complex_type_def.slice.value
+
+        def err(reason: str) -> list[LinterMessage]:
+            msg = f"{reason} for argument {ast.unparse(arg)!r} in function {function_name!r}"
+            return [self._gen_linter_message(msg, arg.lineno)]
+
+        if not arg_type:
+            return err("Missing type annotation")
+
+        # Case 1: optional
+        if T := get_optional_type(arg_type):
+            # NOTE: optional types can be primitives or any of the allowed container types
+            #   so check if arg is an optional, and if so, run the rest of the validation with the unpacked type
+            arg_type = T
+
+        # Case 2: int, float, str, bool
+        if is_primitive(arg_type):
+            return ok()
+        # Case 3: Containers (Dict, List, Tuple)
+        elif isinstance(arg_type, ast.Subscript):
+            container_type = arg_type.value.id
+
+            if container_type in ["Dict", "dict"]:
+                KV = get_generic_type(arg_type)
+
+                assert isinstance(KV, ast.Tuple)  # dict[K,V] has ast.Tuple slice
+
+                K, V = KV.elts
+                if not is_primitive(K):
+                    return err(f"Non-primitive key type {ast.unparse(K)!r}")
+                if not is_primitive(V):
+                    return err(f"Non-primitive value type {ast.unparse(V)!r}")
+                return ok()
+            elif container_type in ["List", "list"]:
+                T = get_generic_type(arg_type)
+                if is_primitive(T):
+                    return ok()
+                else:
+                    return err(f"Non-primitive element type {ast.unparse(T)!r}")
+            elif container_type in ["Tuple", "tuple"]:
+                E_N = get_generic_type(arg_type)
+                assert isinstance(E_N, ast.Tuple)  # tuple[...] has ast.Tuple slice
+
+                for e in E_N.elts:
+                    if not is_primitive(e):
+                        return err(f"Non-primitive element type '{ast.unparse(e)!r}'")
+
+                return ok()
+
+            return err(f"Unsupported container type {container_type!r}")
         else:
-            sub_type = complex_type_def.slice
-        if type_name == "Dict":
-            sub_type_tuple = cast(ast.Tuple, sub_type)
-            for el in sub_type_tuple.elts:
-                if not isinstance(el, ast.Name):
-                    desc = "Dict can only have primitive types"
-                    linter_errors.append(self._gen_linter_message(desc, arg_def.lineno))
-        elif not isinstance(sub_type, ast.Name):
-            desc = "List can only have primitive types"
-            linter_errors.append(self._gen_linter_message(desc, arg_def.lineno))
-        return linter_errors
+            return err(f"Unsupported argument type {ast.unparse(arg_type)!r}")
+
 
+class ReturnTypeValidator(ComponentFunctionValidator):
+    """Validates that component functions always return AppDef type"""
+
+    def __init__(self, supported_return_type: str) -> None:
+        super().__init__()
+        self._supported_return_type = supported_return_type
 
-class TorchxReturnValidator(TorchxFunctionValidator):
     def _get_return_annotation(
         self, app_specs_func_def: ast.FunctionDef
     ) -> Optional[str]:
@@ -197,7 +272,7 @@ class TorchxReturnValidator(TorchxFunctionValidator):
             * AppDef
             * specs.AppDef
         """
-        supported_return_annotation = "AppDef"
+        supported_return_annotation = self._supported_return_type
         return_annotation = self._get_return_annotation(app_specs_func_def)
         linter_errors = []
         if not return_annotation:
@@ -220,7 +295,7 @@ class TorchxReturnValidator(TorchxFunctionValidator):
         return linter_errors
 
 
-class TorchFunctionVisitor(ast.NodeVisitor):
+class ComponentFnVisitor(ast.NodeVisitor):
     """
     Visitor that finds the component_function and runs registered validators on it.
     Current registered validators:
@@ -238,11 +313,18 @@ class TorchFunctionVisitor(ast.NodeVisitor):
 
     """
 
-    def __init__(self, component_function_name: str) -> None:
-        self.validators = [
-            TorchxFunctionArgsValidator(),
-            TorchxReturnValidator(),
-        ]
+    def __init__(
+        self,
+        component_function_name: str,
+        validators: Optional[List[ComponentFunctionValidator]],
+    ) -> None:
+        if validators is None:
+            self.validators: List[ComponentFunctionValidator] = [
+                ArgTypeValidator(),
+                ReturnTypeValidator("AppDef"),
+            ]
+        else:
+            self.validators = validators
         self.linter_errors: List[LinterMessage] = []
         self.component_function_name = component_function_name
         self.visited_function = False
@@ -258,7 +340,11 @@ class TorchFunctionVisitor(ast.NodeVisitor):
             self.linter_errors += validator.validate(node)
 
 
-def validate(path: str, component_function: str) -> List[LinterMessage]:
+def validate(
+    path: str,
+    component_function: str,
+    validators: Optional[List[ComponentFunctionValidator]] = None,
+) -> List[LinterMessage]:
     """
     Validates the function to make sure it complies the component standard.
 
@@ -287,7 +373,7 @@ def validate(path: str, component_function: str) -> List[LinterMessage]:
             severity="error",
         )
         return [linter_message]
-    visitor = TorchFunctionVisitor(component_function)
+    visitor = ComponentFnVisitor(component_function, validators)
     visitor.visit(module)
     linter_errors = visitor.linter_errors
     if not visitor.visited_function: