metaflow 2.11.16__py2.py3-none-any.whl → 2.12.0__py2.py3-none-any.whl

metaflow/cli.py

@@ -3,39 +3,16 @@ import sys
 import traceback
 from datetime import datetime
 from functools import wraps
-import metaflow.tracing as tracing
 
+import metaflow.tracing as tracing
 from metaflow._vendor import click
 
-from . import lint
-from . import plugins
-from . import parameters
-from . import decorators
-from . import metaflow_version
-from . import namespace
-from .metaflow_current import current
+from . import decorators, lint, metaflow_version, namespace, parameters, plugins
 from .cli_args import cli_args
-from .tagging_util import validate_tags
-from .util import (
-    resolve_identity,
-    decompress_list,
-    write_latest_run_id,
-    get_latest_run_id,
-)
-from .task import MetaflowTask
+from .client.core import get_metadata
+from .datastore import FlowDataStore, TaskDataStore, TaskDataStoreSet
 from .exception import CommandException, MetaflowException
 from .graph import FlowGraph
-from .datastore import FlowDataStore, TaskDataStoreSet, TaskDataStore
-
-from .runtime import NativeRuntime
-from .package import MetaflowPackage
-from .plugins import (
-    DATASTORES,
-    ENVIRONMENTS,
-    LOGGING_SIDECARS,
-    METADATA_PROVIDERS,
-    MONITOR_SIDECARS,
-)
 from .metaflow_config import (
     DEFAULT_DATASTORE,
     DEFAULT_ENVIRONMENT,
@@ -44,12 +21,29 @@ from .metaflow_config import (
     DEFAULT_MONITOR,
     DEFAULT_PACKAGE_SUFFIXES,
 )
+from .metaflow_current import current
 from .metaflow_environment import MetaflowEnvironment
+from .mflog import LOG_SOURCES, mflog
+from .package import MetaflowPackage
+from .plugins import (
+    DATASTORES,
+    ENVIRONMENTS,
+    LOGGING_SIDECARS,
+    METADATA_PROVIDERS,
+    MONITOR_SIDECARS,
+)
 from .pylint_wrapper import PyLint
-from .R import use_r, metaflow_r_version
-from .mflog import mflog, LOG_SOURCES
+from .R import metaflow_r_version, use_r
+from .runtime import NativeRuntime
+from .tagging_util import validate_tags
+from .task import MetaflowTask
 from .unbounded_foreach import UBF_CONTROL, UBF_TASK
-
+from .util import (
+    decompress_list,
+    get_latest_run_id,
+    resolve_identity,
+    write_latest_run_id,
+)
 
 ERASE_TO_EOL = "\033[K"
 HIGHLIGHT = "red"
@@ -557,6 +551,13 @@ def common_run_options(func):
         type=str,
         help="Write the ID of this run to the file specified.",
     )
+    @click.option(
+        "--runner-attribute-file",
+        default=None,
+        show_default=True,
+        type=str,
+        help="Write the metadata and pathspec of this run to the file specified. Used internally for Metaflow's Runner API.",
+    )
     @wraps(func)
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
@@ -615,6 +616,7 @@ def resume(
     decospecs=None,
     run_id_file=None,
     resume_identifier=None,
+    runner_attribute_file=None,
 ):
     before_run(obj, tags, decospecs + obj.environment.decospecs())
 
@@ -670,9 +672,14 @@ def resume(
         max_log_size=max_log_size * 1024 * 1024,
         resume_identifier=resume_identifier,
     )
-    write_run_id(run_id_file, runtime.run_id)
+    write_file(run_id_file, runtime.run_id)
     runtime.print_workflow_info()
+
     runtime.persist_constants()
+    write_file(
+        runner_attribute_file,
+        "%s:%s" % (get_metadata(), "/".join((obj.flow.name, runtime.run_id))),
+    )
     if clone_only:
         runtime.clone_original_run()
     else:
@@ -703,6 +710,7 @@ def run(
     max_log_size=None,
     decospecs=None,
     run_id_file=None,
+    runner_attribute_file=None,
     user_namespace=None,
     **kwargs
 ):
@@ -726,18 +734,22 @@ def run(
         max_log_size=max_log_size * 1024 * 1024,
     )
     write_latest_run_id(obj, runtime.run_id)
-    write_run_id(run_id_file, runtime.run_id)
+    write_file(run_id_file, runtime.run_id)
 
     obj.flow._set_constants(obj.graph, kwargs)
     runtime.print_workflow_info()
     runtime.persist_constants()
+    write_file(
+        runner_attribute_file,
+        "%s:%s" % (get_metadata(), "/".join((obj.flow.name, runtime.run_id))),
+    )
     runtime.execute()
 
 
-def write_run_id(run_id_file, run_id):
-    if run_id_file is not None:
-        with open(run_id_file, "w") as f:
-            f.write(str(run_id))
+def write_file(file_path, content):
+    if file_path is not None:
+        with open(file_path, "w") as f:
+            f.write(str(content))
 
 
 def before_run(obj, tags, decospecs):

metaflow/cmd/develop/stubs.py

@@ -23,6 +23,8 @@ def _check_stubs_supported():
         if _py_ver >= (3, 4):
             if _py_ver >= (3, 8):
                 from importlib import metadata
+            elif _py_ver >= (3, 7):
+                from metaflow._vendor import importlib_metadata as metadata
             elif _py_ver >= (3, 6):
                 from metaflow._vendor.v3_6 import importlib_metadata as metadata
             else:

metaflow/extension_support/__init__.py

@@ -262,6 +262,8 @@ if _py_ver >= (3, 4):
 
     if _py_ver >= (3, 8):
         from importlib import metadata
+    elif _py_ver >= (3, 7):
+        from metaflow._vendor import importlib_metadata as metadata
     elif _py_ver >= (3, 6):
         from metaflow._vendor.v3_6 import importlib_metadata as metadata
     else:

metaflow/parameters.py

@@ -388,6 +388,7 @@ def add_custom_parameters(deploy_mode=False):
     # deploy_mode determines whether deploy-time functions should or should
     # not be evaluated for this command
     def wrapper(cmd):
+        cmd.has_flow_params = True
         # Iterate over parameters in reverse order so cmd.params lists options
         # in the order they are defined in the FlowSpec subclass
         for arg in parameters[::-1]:

metaflow/plugins/aws/batch/batch_decorator.py

@@ -88,15 +88,15 @@ class BatchDecorator(StepDecorator):
         Alias for inferentia. Use only one of the two.
     efa : int, default 0
         Number of elastic fabric adapter network devices to attach to container
-    ephemeral_storage: int, default None
-        The total amount, in GiB, of ephemeral storage to set for the task (21-200)
+    ephemeral_storage : int, default None
+        The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
         This is only relevant for Fargate compute environments
     log_driver: str, optional, default None
         The log driver to use for the Amazon ECS container.
     log_options: List[str], optional, default None
         List of strings containing options for the chosen log driver. The configurable values
         depend on the `log driver` chosen. Validation of these options is not supported yet.
-        Example usage: ["awslogs-group:aws/batch/job"]
+        Example: [`awslogs-group:aws/batch/job`]
     """
 
     name = "batch"

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -275,11 +275,6 @@ class KubernetesJob(object):
         # (unique UID) per Metaflow task attempt.
         client = self._client.get()
 
-        # tmpfs variables
-        use_tmpfs = self._kwargs["use_tmpfs"]
-        tmpfs_size = self._kwargs["tmpfs_size"]
-        tmpfs_enabled = use_tmpfs or (tmpfs_size and not use_tmpfs)
-
         self._job = client.V1Job(
             api_version="batch/v1",
             kind="Job",

metaflow/runner/click_api.py

@@ -0,0 +1,406 @@
+import sys
+
+if sys.version_info < (3, 7):
+    raise RuntimeError(
+        """
+        The Metaflow Programmatic API is not supported for versions of Python less than 3.7
+    """
+    )
+
+import datetime
+import importlib
+import inspect
+import itertools
+import uuid
+from collections import OrderedDict
+from typing import Any, Callable, Dict, List, Optional
+from typing import OrderedDict as TOrderedDict
+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._vendor.typeguard import TypeCheckError, check_type
+from metaflow.cli import start
+from metaflow.includefile import FilePathClass
+from metaflow.parameters import JSONTypeClass
+
+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: str,
+    FilePathClass: str,
+}
+
+
+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": {}}
+
+    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', expected: '%s', default is '%s'"
+                % (supplied_k, annotations[supplied_k], defaults[supplied_k])
+            )
+
+        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 get_annotation(param: Union[click.Argument, click.Option]):
+    py_type = click_to_python_types[type(param.type)]
+    if not param.required:
+        if param.multiple or param.nargs == -1:
+            return Optional[List[py_type]]
+        else:
+            return Optional[py_type]
+    else:
+        if param.multiple or param.nargs == -1:
+            return List[py_type]
+        else:
+            return py_type
+
+
+def get_inspect_param_obj(p: Union[click.Argument, click.Option], kind: str):
+    return inspect.Parameter(
+        name=p.name,
+        kind=kind,
+        default=p.default,
+        annotation=get_annotation(p),
+    )
+
+
+# Cache to store already loaded modules
+loaded_modules = {}
+
+
+def extract_flowspec_params(flow_file: str) -> List[Parameter]:
+    # 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", 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, inspect.isclass)
+
+    parameters = []
+    for _, kls in classes:
+        if kls != FlowSpec and issubclass(kls, FlowSpec):
+            for _, obj in inspect.getmembers(kls):
+                if isinstance(obj, Parameter):
+                    parameters.append(obj)
+
+    return parameters
+
+
+class MetaflowAPI(object):
+    def __init__(self, parent=None, **kwargs):
+        self._parent = parent
+        self._chain = [{self._API_NAME: kwargs}]
+
+    @property
+    def parent(self):
+        if self._parent:
+            return self._parent
+        return None
+
+    @property
+    def chain(self):
+        return self._chain
+
+    @classmethod
+    def from_cli(cls, flow_file: str, cli_collection: Callable) -> Callable:
+        flow_parameters = extract_flowspec_params(flow_file)
+        class_dict = {"__module__": "metaflow", "_API_NAME": flow_file}
+        command_groups = cli_collection.sources
+        for each_group in command_groups:
+            for _, cmd_obj in each_group.commands.items():
+                if isinstance(cmd_obj, click.Group):
+                    # TODO: possibly check for fake groups with cmd_obj.name in ["cli", "main"]
+                    class_dict[cmd_obj.name] = extract_group(cmd_obj, flow_parameters)
+                elif isinstance(cmd_obj, click.Command):
+                    class_dict[cmd_obj.name] = extract_command(cmd_obj, flow_parameters)
+                else:
+                    raise RuntimeError(
+                        "Cannot handle %s of type %s" % (cmd_obj.name, type(cmd_obj))
+                    )
+
+        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, **kwargs):
+            method_params = _method_sanity_check(
+                possible_arg_params,
+                possible_opt_params,
+                annotations,
+                defaults,
+                **kwargs,
+            )
+            return to_return(parent=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 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 isinstance(v, list):
+                        for i in v:
+                            components.append(i)
+                    else:
+                        components.append(v)
+                for k, v in options.items():
+                    if isinstance(v, list):
+                        for i in v:
+                            components.append("--%s" % k)
+                            components.append(str(i))
+                    else:
+                        components.append("--%s" % k)
+                        if v != "flag":
+                            components.append(str(v))
+
+        return components
+
+
+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] = 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] = get_inspect_param_obj(
+                each_param, inspect.Parameter.KEYWORD_ONLY
+            )
+            opt_parameters[each_param.name] = each_param
+
+        annotations[each_param.name] = get_annotation(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, **kwargs):
+        method_params = _method_sanity_check(
+            possible_arg_params, possible_opt_params, annotations, defaults, **kwargs
+        )
+        return resulting_class(parent=_self, **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, **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__":
+    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_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))