ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/runner/click_api.py

@@ -0,0 +1,717 @@
+import os
+import sys
+
+_py_ver = sys.version_info[:2]
+
+if _py_ver >= (3, 8):
+    from metaflow._vendor.typeguard import TypeCheckError, check_type
+elif _py_ver >= (3, 7):
+    from metaflow._vendor.v3_7.typeguard import TypeCheckError, check_type
+else:
+    raise RuntimeError(
+        """
+        The Metaflow Programmatic API is not supported for versions of Python less than 3.7
+    """
+    )
+
+import datetime
+import functools
+import importlib
+import inspect
+import itertools
+import uuid
+import json
+from collections import OrderedDict
+from typing import Any, Callable, Dict, List, Optional, Type
+from typing import OrderedDict as TOrderedDict
+from typing import Tuple as TTuple
+from typing import Union
+
+from metaflow import FlowSpec, Parameter
+from metaflow._vendor import click
+from metaflow._vendor.click.types import (
+    BoolParamType,
+    Choice,
+    DateTime,
+    File,
+    FloatParamType,
+    IntParamType,
+    Path,
+    StringParamType,
+    Tuple,
+    UUIDParameterType,
+)
+from metaflow.decorators import add_decorator_options
+from metaflow.exception import MetaflowException
+from metaflow.flowspec import FlowStateItems
+from metaflow.includefile import FilePathClass
+from metaflow.metaflow_config import CLICK_API_PROCESS_CONFIG
+from metaflow.parameters import JSONTypeClass, flow_context
+from metaflow.user_configs.config_options import (
+    ConfigValue,
+    ConvertDictOrStr,
+    ConvertPath,
+    LocalFileInput,
+    MultipleTuple,
+    config_options_with_config_input,
+)
+from metaflow.user_decorators.user_flow_decorator import FlowMutator
+
+# Define a recursive type alias for JSON
+JSON = Union[Dict[str, "JSON"], List["JSON"], str, int, float, bool, None]
+
+click_to_python_types = {
+    StringParamType: str,
+    IntParamType: int,
+    FloatParamType: float,
+    BoolParamType: bool,
+    UUIDParameterType: uuid.UUID,
+    Path: str,
+    DateTime: datetime.datetime,
+    Tuple: tuple,
+    Choice: str,
+    File: str,
+    JSONTypeClass: JSON,
+    FilePathClass: str,
+    LocalFileInput: str,
+    MultipleTuple: TTuple[str, Union[JSON, ConfigValue]],
+}
+
+
+def _method_sanity_check(
+    possible_arg_params: TOrderedDict[str, click.Argument],
+    possible_opt_params: TOrderedDict[str, click.Option],
+    annotations: TOrderedDict[str, Any],
+    defaults: TOrderedDict[str, Any],
+    **kwargs
+) -> Dict[str, Any]:
+    method_params = {"args": {}, "options": {}, "defaults": defaults}
+
+    possible_params = OrderedDict()
+    possible_params.update(possible_arg_params)
+    possible_params.update(possible_opt_params)
+
+    # supplied kwargs
+    for supplied_k, supplied_v in kwargs.items():
+        if supplied_k not in possible_params:
+            raise ValueError(
+                "Unknown argument: '%s', possible args are: %s"
+                % (supplied_k, ", ".join(possible_params.keys()))
+            )
+
+        try:
+            check_type(supplied_v, annotations[supplied_k])
+        except TypeCheckError:
+            raise TypeError(
+                "Invalid type for '%s' (%s), expected: '%s', default is '%s' but found '%s'"
+                % (
+                    supplied_k,
+                    type(supplied_k),
+                    annotations[supplied_k],
+                    defaults[supplied_k],
+                    str(supplied_v),
+                )
+            )
+
+        # Clean up values to make them into what click expects
+        if annotations[supplied_k] == JSON:
+            # JSON should be a string (json dumps)
+            supplied_v = json.dumps(supplied_v)
+        elif supplied_k == "config_value":
+            # Special handling of config value because we need to go look in the tuple
+            new_list = []
+            for cfg_name, cfg_value in supplied_v:
+                if isinstance(cfg_value, ConfigValue):
+                    # ConfigValue should be JSONified and converted to a string
+                    new_list.append((cfg_name, json.dumps(cfg_value.to_dict())))
+                elif isinstance(cfg_value, dict):
+                    # ConfigValue passed as a dictionary
+                    new_list.append((cfg_name, json.dumps(cfg_value)))
+                else:
+                    raise TypeError(
+                        "Invalid type for a config-value, expected a ConfigValue or "
+                        "dict but got '%s'" % type(cfg_value)
+                    )
+            supplied_v = new_list
+
+        if supplied_k in possible_arg_params:
+            cli_name = possible_arg_params[supplied_k].opts[0].strip("-")
+            method_params["args"][cli_name] = supplied_v
+        elif supplied_k in possible_opt_params:
+            if possible_opt_params[supplied_k].is_bool_flag:
+                # it is a boolean flag..
+                if supplied_v == True:
+                    cli_name = possible_opt_params[supplied_k].opts[0].strip("-")
+                elif supplied_v == False:
+                    if possible_opt_params[supplied_k].secondary_opts:
+                        cli_name = (
+                            possible_opt_params[supplied_k].secondary_opts[0].strip("-")
+                        )
+                    else:
+                        continue
+                supplied_v = "flag"
+            else:
+                cli_name = possible_opt_params[supplied_k].opts[0].strip("-")
+            method_params["options"][cli_name] = supplied_v
+
+    # possible kwargs
+    for _, possible_v in possible_params.items():
+        cli_name = possible_v.opts[0].strip("-")
+        if (
+            (cli_name not in method_params["args"])
+            and (cli_name not in method_params["options"])
+        ) and possible_v.required:
+            raise ValueError("Missing argument: %s is required." % cli_name)
+
+    return method_params
+
+
+def _cleanup_flow_parameters(cmd_obj: Union[click.Command, click.Group]):
+    if hasattr(cmd_obj, "original_params"):
+        cmd_obj.params = list(cmd_obj.original_params)
+
+    if isinstance(cmd_obj, click.Group):
+        for sub_cmd_name in cmd_obj.list_commands(None):
+            sub_cmd = cmd_obj.get_command(None, sub_cmd_name)
+            if sub_cmd:
+                _cleanup_flow_parameters(sub_cmd)
+
+
+def _lazy_load_command(
+    cli_collection: click.Group,
+    flow_parameters: Union[str, List[Parameter]],
+    _self,
+    name: str,
+):
+    # Context is not used in get_command so we can pass None. Since we pin click,
+    # this won't change from under us.
+
+    if isinstance(flow_parameters, str):
+        # Resolve flow_parameters -- for start, this is a function which we
+        # need to call to figure out the actual parameters (may be changed by configs)
+        flow_parameters = getattr(_self, flow_parameters)()
+    cmd_obj = cli_collection.get_command(None, name)
+    if cmd_obj:
+        _cleanup_flow_parameters(cmd_obj)
+        if isinstance(cmd_obj, click.Group):
+            # TODO: possibly check for fake groups with cmd_obj.name in ["cli", "main"]
+            result = functools.partial(extract_group(cmd_obj, flow_parameters), _self)
+        elif isinstance(cmd_obj, click.Command):
+            result = functools.partial(extract_command(cmd_obj, flow_parameters), _self)
+        else:
+            raise RuntimeError(
+                "Cannot handle %s of type %s" % (cmd_obj.name, type(cmd_obj))
+            )
+        setattr(_self, name, result)
+        return result
+    else:
+        raise AttributeError()
+
+
+def get_annotation(param: click.Parameter) -> TTuple[Type, bool]:
+    py_type = click_to_python_types[type(param.type)]
+    if param.nargs == -1:
+        # This is the equivalent of *args effectively
+        # so the type annotation should be the type of the
+        # elements in the list
+        return py_type, True
+    if not param.required:
+        if param.multiple or param.nargs > 1:
+            return Optional[Union[List[py_type], TTuple[py_type]]], False
+        else:
+            return Optional[py_type], False
+    else:
+        if param.multiple or param.nargs > 1:
+            return Union[List[py_type], TTuple[py_type]], False
+        else:
+            return py_type, False
+
+
+def get_inspect_param_obj(p: Union[click.Argument, click.Option], kind: str):
+    annotation, is_vararg = get_annotation(p)
+    return (
+        inspect.Parameter(
+            name="args" if is_vararg else p.name,
+            kind=inspect.Parameter.VAR_POSITIONAL if is_vararg else kind,
+            default=inspect.Parameter.empty if is_vararg else p.default,
+            annotation=annotation,
+        ),
+        (
+            Optional[Union[TTuple[annotation], List[annotation]]]
+            if is_vararg
+            else annotation
+        ),
+    )
+
+
+# Cache to store already loaded modules
+loaded_modules = {}
+
+
+def extract_flow_class_from_file(flow_file: str) -> FlowSpec:
+    if not os.path.exists(flow_file):
+        raise FileNotFoundError("Flow file not present at '%s'" % flow_file)
+
+    flow_dir = os.path.dirname(os.path.abspath(flow_file))
+    path_was_added = False
+
+    # Only add to path if it's not already there
+    if flow_dir not in sys.path:
+        sys.path.insert(0, flow_dir)
+        path_was_added = True
+
+    try:
+        # Get module name from the file path
+        module_name = os.path.splitext(os.path.basename(flow_file))[0]
+
+        # Check if the module has already been loaded
+        if flow_file in loaded_modules:
+            module = loaded_modules[flow_file]
+        else:
+            # Load the module if it's not already loaded
+            spec = importlib.util.spec_from_file_location(module_name, flow_file)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            # Cache the loaded module
+            loaded_modules[flow_file] = module
+
+        classes = inspect.getmembers(
+            module, lambda x: inspect.isclass(x) or isinstance(x, FlowMutator)
+        )
+        flow_cls = None
+
+        for _, kls in classes:
+            if isinstance(kls, FlowMutator):
+                kls = kls._flow_cls
+            if (
+                kls is not FlowSpec
+                and kls.__module__ == module_name
+                and issubclass(kls, FlowSpec)
+            ):
+                if flow_cls is not None:
+                    raise MetaflowException(
+                        "Multiple FlowSpec classes found in %s" % flow_file
+                    )
+                flow_cls = kls
+
+        if flow_cls is None:
+            raise MetaflowException("No FlowSpec class found in %s" % flow_file)
+        return flow_cls
+    finally:
+        # Only remove from path if we added it
+        if path_was_added:
+            try:
+                sys.path.remove(flow_dir)
+            except ValueError:
+                # User's code might have removed it already
+                pass
+
+
+class MetaflowAPI(object):
+    def __init__(self, parent=None, flow_cls=None, config_input=None, **kwargs):
+        self._parent = parent
+        self._chain = [{self._API_NAME: kwargs}]
+        self._flow_cls = flow_cls
+        self._config_input = config_input
+        self._cached_computed_parameters = None
+
+    @property
+    def parent(self):
+        if self._parent:
+            return self._parent
+        return None
+
+    @property
+    def chain(self):
+        return self._chain
+
+    @property
+    def name(self):
+        return self._API_NAME
+
+    @classmethod
+    def from_cli(cls, flow_file: str, cli_collection: Callable) -> Callable:
+        flow_cls = extract_flow_class_from_file(flow_file)
+
+        with flow_context(flow_cls) as _:
+            cli_collection, config_input = config_options_with_config_input(
+                cli_collection
+            )
+            cli_collection = add_decorator_options(cli_collection)
+
+        def getattr_wrapper(_self, name):
+            # Functools.partial do not automatically bind self (no __get__)
+            with flow_context(flow_cls) as _:
+                # We also wrap this in the proper flow context because since commands
+                # are loaded lazily, we need the proper flow context to compute things
+                # like parameters. If we do not do this, the outer flow's context will
+                # be used.
+                return _self._internal_getattr(_self, name)
+
+        class_dict = {
+            "__module__": "metaflow",
+            "_API_NAME": flow_file,
+            "_internal_getattr": staticmethod(
+                functools.partial(
+                    _lazy_load_command, cli_collection, "_compute_flow_parameters"
+                )
+            ),
+            "__getattr__": getattr_wrapper,
+        }
+
+        to_return = type(flow_file, (MetaflowAPI,), class_dict)
+        to_return.__name__ = flow_file
+
+        (
+            params_sigs,
+            possible_arg_params,
+            possible_opt_params,
+            annotations,
+            defaults,
+        ) = extract_all_params(cli_collection)
+
+        def _method(_self, *args, **kwargs):
+            method_params = _method_sanity_check(
+                possible_arg_params,
+                possible_opt_params,
+                annotations,
+                defaults,
+                **kwargs,
+            )
+            return to_return(
+                parent=None,
+                flow_cls=flow_cls,
+                config_input=config_input,
+                **method_params,
+            )
+
+        m = _method
+        m.__name__ = cli_collection.name
+        m.__doc__ = getattr(cli_collection, "help", None)
+        m.__signature__ = inspect.signature(_method).replace(
+            parameters=params_sigs.values()
+        )
+        m.__annotations__ = annotations
+        m.__defaults__ = tuple(defaults.values())
+
+        return m
+
+    def execute(self) -> List[str]:
+        parents = []
+        current = self
+        while current.parent:
+            parents.append(current.parent)
+            current = current.parent
+
+        parents.reverse()
+
+        final_chain = list(itertools.chain.from_iterable([p.chain for p in parents]))
+        final_chain.extend(self.chain)
+
+        components = []
+        for each_cmd in final_chain:
+            for cmd, params in each_cmd.items():
+                components.append(cmd)
+                args = params.pop("args", {})
+                options = params.pop("options", {})
+
+                for _, v in args.items():
+                    if v is None:
+                        continue
+                    if isinstance(v, (list, tuple)):
+                        for i in v:
+                            components.append(i)
+                    else:
+                        components.append(v)
+                for k, v in options.items():
+                    if v is None:
+                        continue
+                    if isinstance(v, list):
+                        for i in v:
+                            if isinstance(i, tuple):
+                                components.append("--%s" % k)
+                                components.extend(map(str, i))
+                            else:
+                                components.append("--%s" % k)
+                                components.append(str(i))
+                    elif v is None:
+                        continue  # Skip None values -- they are defaults and converting
+                        # them to string will not be what the user wants
+                    else:
+                        components.append("--%s" % k)
+                        if v != "flag":
+                            components.append(str(v))
+
+        return components
+
+    def _compute_flow_parameters(self):
+        if (
+            self._flow_cls is None
+            or self._config_input is None
+            or self._parent is not None
+        ):
+            raise RuntimeError(
+                "Computing flow-level parameters for a non start API. "
+                "Please report to the Metaflow team."
+            )
+
+        if self._cached_computed_parameters is not None:
+            return self._cached_computed_parameters
+        self._cached_computed_parameters = []
+
+        config_options = None
+        if CLICK_API_PROCESS_CONFIG:
+            with flow_context(self._flow_cls) as _:
+                # We are going to resolve the configs first and then get the parameters.
+                # Note that configs may update/add parameters so the order is important
+                # Since part of the processing of configs happens by click, we need to
+                # "fake" it.
+
+                # Extract any config options as well as datastore and quiet options
+                method_params = self._chain[0][self._API_NAME]
+                opts = method_params["options"]
+                defaults = method_params["defaults"]
+
+                ds = opts.get("datastore", defaults["datastore"])
+                quiet = opts.get("quiet", defaults["quiet"])
+
+                # Order to find config or config_value:
+                # 1. Passed directly to the Click API
+                # 2. If not found, check if passed through an environment variable
+                # 3. If not found, use the default value
+                is_default = False
+                config_file = opts.get("config")
+                if config_file is None:
+                    # Check if it was set through an environment variable -- we
+                    # don't have click process them here so we need to "fake" it.
+                    env_config_file = os.environ.get("METAFLOW_FLOW_CONFIG")
+                    if env_config_file:
+                        # Convert dict items to list of tuples
+                        config_file = list(json.loads(env_config_file).items())
+                        is_default = False
+                    else:
+                        is_default = True
+                        config_file = defaults.get("config")
+
+                if config_file:
+                    config_file = dict(
+                        map(
+                            lambda x: (
+                                x[0],
+                                ConvertPath.convert_value(x[1], is_default),
+                            ),
+                            config_file,
+                        )
+                    )
+
+                is_default = False
+                config_value = opts.get("config-value")
+                if config_value is None:
+                    env_config_value = os.environ.get("METAFLOW_FLOW_CONFIG_VALUE")
+                    if env_config_value:
+                        # Parse environment variable using MultipleTuple logic
+                        loaded = json.loads(env_config_value)
+                        # Convert dict items to list of tuples with JSON-serialized values
+                        config_value = [
+                            (k, json.dumps(v) if not isinstance(v, str) else v)
+                            for k, v in loaded.items()
+                        ]
+                        is_default = False
+                    else:
+                        is_default = True
+                        config_value = defaults.get("config_value")
+
+                if config_value:
+                    config_value = dict(
+                        map(
+                            lambda x: (
+                                x[0],
+                                ConvertDictOrStr.convert_value(x[1], is_default),
+                            ),
+                            config_value,
+                        )
+                    )
+
+                if (config_file is None) ^ (config_value is None):
+                    # If we have one, we should have the other
+                    raise MetaflowException(
+                        "Options were not properly set -- this is an internal error."
+                    )
+
+                if config_file:
+                    # Process both configurations; the second one will return all the merged
+                    # configuration options properly processed.
+                    self._config_input.process_configs(
+                        self._flow_cls.__name__, "config", config_file, quiet, ds
+                    )
+                    config_options = self._config_input.process_configs(
+                        self._flow_cls.__name__, "config_value", config_value, quiet, ds
+                    )
+
+        # At this point, we are like in start() in cli.py -- we obtained the
+        # properly processed config_options which we can now use to process
+        # the config decorators (including StepMutator/FlowMutator)
+        # Note that if CLICK_API_PROCESS_CONFIG is False, we still do this because
+        # it will init all parameters (config_options will be None)
+        # We ignore any errors if we don't check the configs in the click API.
+
+        # Init all values in the flow mutators and then process them
+        for decorator in self._flow_cls._flow_state[FlowStateItems.FLOW_MUTATORS]:
+            decorator.external_init()
+
+        new_cls = self._flow_cls._process_config_decorators(
+            config_options, process_configs=CLICK_API_PROCESS_CONFIG
+        )
+        if new_cls:
+            self._flow_cls = new_cls
+
+        for _, param in self._flow_cls._get_parameters():
+            if param.IS_CONFIG_PARAMETER:
+                continue
+            self._cached_computed_parameters.append(param)
+        return self._cached_computed_parameters
+
+
+def extract_all_params(cmd_obj: Union[click.Command, click.Group]):
+    arg_params_sigs = OrderedDict()
+    opt_params_sigs = OrderedDict()
+    params_sigs = OrderedDict()
+
+    arg_parameters = OrderedDict()
+    opt_parameters = OrderedDict()
+    annotations = OrderedDict()
+    defaults = OrderedDict()
+
+    for each_param in cmd_obj.params:
+        if isinstance(each_param, click.Argument):
+            (
+                arg_params_sigs[each_param.name],
+                annotations[each_param.name],
+            ) = get_inspect_param_obj(each_param, inspect.Parameter.POSITIONAL_ONLY)
+            arg_parameters[each_param.name] = each_param
+        elif isinstance(each_param, click.Option):
+            (
+                opt_params_sigs[each_param.name],
+                annotations[each_param.name],
+            ) = get_inspect_param_obj(each_param, inspect.Parameter.KEYWORD_ONLY)
+            opt_parameters[each_param.name] = each_param
+
+        defaults[each_param.name] = each_param.default
+
+    # first, fill in positional arguments
+    for name, each_arg_param in arg_params_sigs.items():
+        params_sigs[name] = each_arg_param
+    # then, fill in keyword arguments
+    for name, each_opt_param in opt_params_sigs.items():
+        params_sigs[name] = each_opt_param
+
+    return params_sigs, arg_parameters, opt_parameters, annotations, defaults
+
+
+def extract_group(cmd_obj: click.Group, flow_parameters: List[Parameter]) -> Callable:
+    class_dict = {"__module__": "metaflow", "_API_NAME": cmd_obj.name}
+    for _, sub_cmd_obj in cmd_obj.commands.items():
+        if isinstance(sub_cmd_obj, click.Group):
+            # recursion
+            class_dict[sub_cmd_obj.name] = extract_group(sub_cmd_obj, flow_parameters)
+        elif isinstance(sub_cmd_obj, click.Command):
+            class_dict[sub_cmd_obj.name] = extract_command(sub_cmd_obj, flow_parameters)
+        else:
+            raise RuntimeError(
+                "Cannot handle %s of type %s" % (sub_cmd_obj.name, type(sub_cmd_obj))
+            )
+
+    resulting_class = type(cmd_obj.name, (MetaflowAPI,), class_dict)
+    resulting_class.__name__ = cmd_obj.name
+
+    (
+        params_sigs,
+        possible_arg_params,
+        possible_opt_params,
+        annotations,
+        defaults,
+    ) = extract_all_params(cmd_obj)
+
+    def _method(_self, *args, **kwargs):
+        method_params = _method_sanity_check(
+            possible_arg_params, possible_opt_params, annotations, defaults, **kwargs
+        )
+        return resulting_class(parent=_self, flow_cls=None, **method_params)
+
+    m = _method
+    m.__name__ = cmd_obj.name
+    m.__doc__ = getattr(cmd_obj, "help", None)
+    m.__signature__ = inspect.signature(_method).replace(
+        parameters=params_sigs.values()
+    )
+    m.__annotations__ = annotations
+    m.__defaults__ = tuple(defaults.values())
+
+    return m
+
+
+def extract_command(
+    cmd_obj: click.Command, flow_parameters: List[Parameter]
+) -> Callable:
+    if getattr(cmd_obj, "has_flow_params", False):
+        for p in flow_parameters[::-1]:
+            cmd_obj.params.insert(0, click.Option(("--" + p.name,), **p.kwargs))
+
+    (
+        params_sigs,
+        possible_arg_params,
+        possible_opt_params,
+        annotations,
+        defaults,
+    ) = extract_all_params(cmd_obj)
+
+    def _method(_self, *args, **kwargs):
+        method_params = _method_sanity_check(
+            possible_arg_params, possible_opt_params, annotations, defaults, **kwargs
+        )
+        _self._chain.append({cmd_obj.name: method_params})
+        return _self.execute()
+
+    m = _method
+    m.__name__ = cmd_obj.name
+    m.__doc__ = getattr(cmd_obj, "help", None)
+    m.__signature__ = inspect.signature(_method).replace(
+        parameters=params_sigs.values()
+    )
+    m.__annotations__ = annotations
+    m.__defaults__ = tuple(defaults.values())
+
+    return m
+
+
+if __name__ == "__main__":
+    from metaflow.cli import start
+
+    api = MetaflowAPI.from_cli("../try.py", start)
+
+    command = api(metadata="local").run(
+        tags=["abc", "def"],
+        decospecs=["kubernetes"],
+        max_workers=5,
+        alpha=3,
+        myfile="path/to/file",
+    )
+    print(" ".join(command))
+
+    command = (
+        api(metadata="local")
+        .kubernetes()
+        .step(
+            step_name="process",
+            code_package_metadata="some_version",
+            code_package_sha="some_sha",
+            code_package_url="some_url",
+        )
+    )
+    print(" ".join(command))
+
+    command = api().tag().add(tags=["abc", "def"])
+    print(" ".join(command))
+
+    command = getattr(api(decospecs=["retry"]), "argo-workflows")().create()
+    print(" ".join(command))