ob-metaflow 2.12.36.3__py2.py3-none-any.whl → 2.13.0.1__py2.py3-none-any.whl

metaflow/flowspec.py

@@ -4,15 +4,18 @@ import sys
 import traceback
 import reprlib
 
+from enum import Enum
 from itertools import islice
 from types import FunctionType, MethodType
-from typing import Any, Callable, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Generator, List, Optional, Tuple
 
 from . import cmd_with_io, parameters
+from .debug import debug
 from .parameters import DelayedEvaluationParameter, Parameter
 from .exception import (
     MetaflowException,
     MissingInMergeArtifactsException,
+    MetaflowInternalError,
     UnhandledInMergeArtifactsException,
 )
 
@@ -20,6 +23,13 @@ from .extension_support import extension_info
 
 from .graph import FlowGraph
 from .unbounded_foreach import UnboundedForeachInput
+from .user_configs.config_decorators import (
+    ConfigValue,
+    CustomFlowDecorator,
+    CustomStepDecorator,
+    MutableFlow,
+    MutableStep,
+)
 from .util import to_pod
 from .metaflow_config import INCLUDE_FOREACH_STACK, MAXIMUM_FOREACH_VALUE_CHARS
 
@@ -65,15 +75,36 @@ class ParallelUBF(UnboundedForeachInput):
         return item or 0  # item is None for the control task, but it is also split 0
 
 
+class _FlowState(Enum):
+    CONFIGS = 1
+    CONFIG_DECORATORS = 2
+    CACHED_PARAMETERS = 3
+
+
 class FlowSpecMeta(type):
-    def __new__(cls, name, bases, dct):
-        f = super().__new__(cls, name, bases, dct)
-        # This makes sure to give _flow_decorators to each
-        # child class (and not share it with the FlowSpec base
-        # class). This is important to not make a "global"
-        # _flow_decorators
-        f._flow_decorators = {}
-        return f
+    def __init__(cls, name, bases, attrs):
+        super().__init__(name, bases, attrs)
+        if name == "FlowSpec":
+            return
+        # We store some state in the flow class itself. This is primarily used to
+        # attach global state to a flow. It is *not* an actual global because of
+        # Runner/NBRunner. This is also created here in the meta class to avoid it being
+        # shared between different children classes.
+
+        # We should move _flow_decorators into this structure as well but keeping it
+        # out to limit the changes for now.
+        cls._flow_decorators = {}
+
+        # Keys are _FlowState enum values
+        cls._flow_state = {}
+
+        cls._init_attrs()
+
+    def _init_attrs(cls):
+        # Graph and steps are specific to the class -- store here so we can access
+        # in class method _process_config_decorators
+        cls._graph = FlowGraph(cls)
+        cls._steps = [getattr(cls, node.name) for node in cls._graph]
 
 
 class FlowSpec(metaclass=FlowSpecMeta):
@@ -96,6 +127,7 @@ class FlowSpec(metaclass=FlowSpecMeta):
         "_cached_input",
         "_graph",
         "_flow_decorators",
+        "_flow_state",
         "_steps",
         "index",
         "input",
@@ -122,9 +154,6 @@ class FlowSpec(metaclass=FlowSpecMeta):
         self._transition = None
         self._cached_input = {}
 
-        self._graph = FlowGraph(self.__class__)
-        self._steps = [getattr(self, node.name) for node in self._graph]
-
         if use_cli:
             with parameters.flow_context(self.__class__) as _:
                 from . import cli
@@ -148,16 +177,12 @@ class FlowSpec(metaclass=FlowSpecMeta):
             fname = fname[:-1]
         return os.path.basename(fname)
 
-    def _set_constants(self, graph, kwargs):
-        from metaflow.decorators import (
-            flow_decorators,
-        )  # To prevent circular dependency
-
-        # Persist values for parameters and other constants (class level variables)
-        # only once. This method is called before persist_constants is called to
-        # persist all values set using setattr
+    @classmethod
+    def _check_parameters(cls, config_parameters=False):
         seen = set()
-        for var, param in self._get_parameters():
+        for _, param in cls._get_parameters():
+            if param.IS_CONFIG_PARAMETER != config_parameters:
+                continue
             norm = param.name.lower()
             if norm in seen:
                 raise MetaflowException(
@@ -166,17 +191,136 @@ class FlowSpec(metaclass=FlowSpecMeta):
                     "case-insensitive." % param.name
                 )
             seen.add(norm)
-        seen.clear()
+
+    @classmethod
+    def _process_config_decorators(cls, config_options, ignore_errors=False):
+
+        # Fast path for no user configurations
+        if not cls._flow_state.get(_FlowState.CONFIG_DECORATORS):
+            # Process parameters to allow them to also use config values easily
+            for var, param in cls._get_parameters():
+                if param.IS_CONFIG_PARAMETER:
+                    continue
+                param.init(ignore_errors)
+            return None
+
+        debug.userconf_exec("Processing mutating step/flow decorators")
+        # We need to convert all the user configurations from DelayedEvaluationParameters
+        # to actual values so they can be used as is in the config decorators.
+
+        # We then reset them to be proper configs so they can be re-evaluated in
+        # _set_constants
+        to_reset_configs = []
+        cls._check_parameters(config_parameters=True)
+        for var, param in cls._get_parameters():
+            if not param.IS_CONFIG_PARAMETER:
+                continue
+            # Note that a config with no default and not required will be None
+            val = config_options.get(param.name.replace("-", "_").lower())
+            if isinstance(val, DelayedEvaluationParameter):
+                val = val()
+            # We store the value as well so that in _set_constants, we don't try
+            # to recompute (no guarantee that it is stable)
+            param._store_value(val)
+            to_reset_configs.append((var, param))
+            debug.userconf_exec("Setting config %s to %s" % (var, str(val)))
+            setattr(cls, var, val)
+
+        # Run all the decorators. Step decorators are directly in the step and
+        # we will run those first and *then* we run all the flow level decorators
+        for step in cls._steps:
+            for deco in step.config_decorators:
+                if isinstance(deco, CustomStepDecorator):
+                    debug.userconf_exec(
+                        "Evaluating step level decorator %s for %s"
+                        % (deco.__class__.__name__, step.name)
+                    )
+                    deco.evaluate(MutableStep(cls, step))
+                else:
+                    raise MetaflowInternalError(
+                        "A non CustomFlowDecorator found in step custom decorators"
+                    )
+            if step.config_decorators:
+                # We remove all mention of the custom step decorator
+                setattr(cls, step.name, step)
+
+        mutable_flow = MutableFlow(cls)
+        for deco in cls._flow_state[_FlowState.CONFIG_DECORATORS]:
+            if isinstance(deco, CustomFlowDecorator):
+                # Sanity check to make sure we are applying the decorator to the right
+                # class
+                if not deco._flow_cls == cls and not issubclass(cls, deco._flow_cls):
+                    raise MetaflowInternalError(
+                        "CustomFlowDecorator registered on the wrong flow -- "
+                        "expected %s but got %s"
+                        % (deco._flow_cls.__name__, cls.__name__)
+                    )
+                debug.userconf_exec(
+                    "Evaluating flow level decorator %s" % deco.__class__.__name__
+                )
+                deco.evaluate(mutable_flow)
+                # We reset cached_parameters on the very off chance that the user added
+                # more configurations based on the configuration
+                if _FlowState.CACHED_PARAMETERS in cls._flow_state:
+                    del cls._flow_state[_FlowState.CACHED_PARAMETERS]
+            else:
+                raise MetaflowInternalError(
+                    "A non CustomFlowDecorator found in flow custom decorators"
+                )
+
+        # Process parameters to allow them to also use config values easily
+        for var, param in cls._get_parameters():
+            if param.IS_CONFIG_PARAMETER:
+                continue
+            param.init()
+        # Reset all configs that were already present in the class.
+        # TODO: This means that users can't override configs directly. Not sure if this
+        # is a pattern we want to support
+        for var, param in to_reset_configs:
+            setattr(cls, var, param)
+
+        # Reset cached parameters again since we added back the config parameters
+        if _FlowState.CACHED_PARAMETERS in cls._flow_state:
+            del cls._flow_state[_FlowState.CACHED_PARAMETERS]
+
+        # Set the current flow class we are in (the one we just created)
+        parameters.replace_flow_context(cls)
+
+        # Re-calculate class level attributes after modifying the class
+        cls._init_attrs()
+        return cls
+
+    def _set_constants(self, graph, kwargs, config_options):
+        from metaflow.decorators import (
+            flow_decorators,
+        )  # To prevent circular dependency
+
+        # Persist values for parameters and other constants (class level variables)
+        # only once. This method is called before persist_constants is called to
+        # persist all values set using setattr
+        self._check_parameters(config_parameters=False)
+
+        seen = set()
         self._success = True
 
         parameters_info = []
         for var, param in self._get_parameters():
             seen.add(var)
-            val = kwargs[param.name.replace("-", "_").lower()]
+            if param.IS_CONFIG_PARAMETER:
+                # Use computed value if already evaluated, else get from config_options
+                val = param._computed_value or config_options.get(
+                    param.name.replace("-", "_").lower()
+                )
+            else:
+                val = kwargs[param.name.replace("-", "_").lower()]
             # Support for delayed evaluation of parameters.
             if isinstance(val, DelayedEvaluationParameter):
                 val = val()
             val = val.split(param.separator) if val and param.separator else val
+            if isinstance(val, ConfigValue):
+                # We store config values as dict so they are accessible with older
+                # metaflow clients. It also makes it easier to access.
+                val = val.to_dict()
             setattr(self, var, val)
             parameters_info.append({"name": var, "type": param.__class__.__name__})
 
@@ -218,6 +362,12 @@ class FlowSpec(metaclass=FlowSpecMeta):
 
     @classmethod
     def _get_parameters(cls):
+        cached = cls._flow_state.get(_FlowState.CACHED_PARAMETERS)
+        if cached is not None:
+            for var in cached:
+                yield var, getattr(cls, var)
+            return
+        build_list = []
         for var in dir(cls):
             if var[0] == "_" or var in cls._NON_PARAMETERS:
                 continue
@@ -226,7 +376,9 @@ class FlowSpec(metaclass=FlowSpecMeta):
             except:
                 continue
             if isinstance(val, Parameter):
+                build_list.append(var)
                 yield var, val
+        cls._flow_state[_FlowState.CACHED_PARAMETERS] = build_list
 
     def _set_datastore(self, datastore):
         self._datastore = datastore

metaflow/graph.py

@@ -45,9 +45,12 @@ def deindent_docstring(doc):
 
 
 class DAGNode(object):
-    def __init__(self, func_ast, decos, doc):
+    def __init__(self, func_ast, decos, doc, source_file, lineno):
         self.name = func_ast.name
-        self.func_lineno = func_ast.lineno
+        self.source_file = source_file
+        # lineno is the start line of decorators in source_file
+        # func_ast.lineno is lines from decorators start to def of function
+        self.func_lineno = lineno + func_ast.lineno - 1
         self.decorators = decos
         self.doc = deindent_docstring(doc)
         self.parallel_step = any(getattr(deco, "IS_PARALLEL", False) for deco in decos)
@@ -62,7 +65,7 @@ class DAGNode(object):
         self.foreach_param = None
         self.num_parallel = 0
         self.parallel_foreach = False
-        self._parse(func_ast)
+        self._parse(func_ast, lineno)
 
         # these attributes are populated by _traverse_graph
         self.in_funcs = set()
@@ -74,7 +77,7 @@ class DAGNode(object):
     def _expr_str(self, expr):
         return "%s.%s" % (expr.value.id, expr.attr)
 
-    def _parse(self, func_ast):
+    def _parse(self, func_ast, lineno):
         self.num_args = len(func_ast.args.args)
         tail = func_ast.body[-1]
 
@@ -94,7 +97,7 @@ class DAGNode(object):
 
             self.has_tail_next = True
             self.invalid_tail_next = True
-            self.tail_next_lineno = tail.lineno
+            self.tail_next_lineno = lineno + tail.lineno - 1
             self.out_funcs = [e.attr for e in tail.value.args]
 
             keywords = dict(
@@ -131,7 +134,7 @@ class DAGNode(object):
             return
 
     def __str__(self):
-        return """*[{0.name} {0.type} (line {0.func_lineno})]*
+        return """*[{0.name} {0.type} ({0.source_file} line {0.func_lineno})]*
     in_funcs={in_funcs}
     out_funcs={out_funcs}
     split_parents={parents}
@@ -156,18 +159,6 @@ class DAGNode(object):
         )
 
 
-class StepVisitor(ast.NodeVisitor):
-    def __init__(self, nodes, flow):
-        self.nodes = nodes
-        self.flow = flow
-        super(StepVisitor, self).__init__()
-
-    def visit_FunctionDef(self, node):
-        func = getattr(self.flow, node.name)
-        if hasattr(func, "is_step"):
-            self.nodes[node.name] = DAGNode(node, func.decorators, func.__doc__)
-
-
 class FlowGraph(object):
     def __init__(self, flow):
         self.name = flow.__name__
@@ -179,13 +170,20 @@ class FlowGraph(object):
         self._postprocess()
 
     def _create_nodes(self, flow):
-        module = __import__(flow.__module__)
-        tree = ast.parse(inspect.getsource(module)).body
-        root = [n for n in tree if isinstance(n, ast.ClassDef) and n.name == self.name][
-            0
-        ]
         nodes = {}
-        StepVisitor(nodes, flow).visit(root)
+        for element in dir(flow):
+            func = getattr(flow, element)
+            if callable(func) and hasattr(func, "is_step"):
+                source_file = inspect.getsourcefile(func)
+                source_lines, lineno = inspect.getsourcelines(func)
+                # This also works for code (strips out leading whitspace based on
+                # first line)
+                source_code = deindent_docstring("".join(source_lines))
+                function_ast = ast.parse(source_code).body[0]
+                node = DAGNode(
+                    function_ast, func.decorators, func.__doc__, source_file, lineno
+                )
+                nodes[element] = node
         return nodes
 
     def _postprocess(self):
@@ -201,6 +199,10 @@ class FlowGraph(object):
 
     def _traverse_graph(self):
         def traverse(node, seen, split_parents):
+            try:
+                self.sorted_nodes.remove(node.name)
+            except ValueError:
+                pass
             self.sorted_nodes.append(node.name)
             if node.type in ("split", "foreach"):
                 node.split_parents = split_parents
@@ -240,9 +242,7 @@ class FlowGraph(object):
         return iter(self.nodes.values())
 
     def __str__(self):
-        return "\n".join(
-            str(n) for _, n in sorted((n.func_lineno, n) for n in self.nodes.values())
-        )
+        return "\n".join(str(self[n]) for n in self.sorted_nodes)
 
     def output_dot(self):
         def edge_specs():
@@ -286,6 +286,7 @@ class FlowGraph(object):
                 "name": name,
                 "type": node_to_type(node),
                 "line": node.func_lineno,
+                "source_file": node.source_file,
                 "doc": node.doc,
                 "decorators": [
                     {

metaflow/includefile.py

@@ -20,6 +20,7 @@ from .parameters import (
 )
 
 from .plugins import DATACLIENTS
+from .user_configs.config_parameters import ConfigValue
 from .util import get_username
 
 import functools
@@ -136,6 +137,7 @@ class FilePathClass(click.ParamType):
                 parameter_name=param.name,
                 logger=ctx.obj.echo,
                 ds_type=ctx.obj.datastore_impl.TYPE,
+                configs=None,
             )
 
         if len(value) > 0 and (value.startswith("{") or value.startswith('"{')):
@@ -243,29 +245,63 @@ class IncludeFile(Parameter):
     default : Union[str, Callable[ParameterContext, str]]
         Default path to a local file. A function
         implies that the parameter corresponds to a *deploy-time parameter*.
-    is_text : bool, default True
+    is_text : bool, optional, default None
         Convert the file contents to a string using the provided `encoding`.
-        If False, the artifact is stored in `bytes`.
-    encoding : str, optional, default 'utf-8'
-        Use this encoding to decode the file contexts if `is_text=True`.
-    required : bool, default False
+        If False, the artifact is stored in `bytes`. A value of None is equivalent to
+        True.
+    encoding : str, optional, default None
+        Use this encoding to decode the file contexts if `is_text=True`. A value of None
+        is equivalent to "utf-8".
+    required : bool, optional, default None
         Require that the user specified a value for the parameter.
-        `required=True` implies that the `default` is not used.
+        `required=True` implies that the `default` is not used. A value of None is
+        equivalent to False
     help : str, optional
         Help text to show in `run --help`.
     show_default : bool, default True
-        If True, show the default value in the help text.
+        If True, show the default value in the help text. A value of None is equivalent
+        to True.
     """
 
     def __init__(
         self,
         name: str,
-        required: bool = False,
-        is_text: bool = True,
-        encoding: str = "utf-8",
+        required: Optional[bool] = None,
+        is_text: Optional[bool] = None,
+        encoding: Optional[str] = None,
         help: Optional[str] = None,
         **kwargs: Dict[str, str]
     ):
+        self._includefile_overrides = {}
+        if is_text is not None:
+            self._includefile_overrides["is_text"] = is_text
+        if encoding is not None:
+            self._includefile_overrides["encoding"] = encoding
+        # NOTA: Right now, there is an issue where these can't be overridden by config
+        # in all circumstances. Ignoring for now.
+        super(IncludeFile, self).__init__(
+            name,
+            required=required,
+            help=help,
+            type=FilePathClass(
+                self._includefile_overrides.get("is_text", True),
+                self._includefile_overrides.get("encoding", "utf-8"),
+            ),
+            **kwargs,
+        )
+
+    def init(self, ignore_errors=False):
+        super(IncludeFile, self).init(ignore_errors)
+
+        # This will use the values set explicitly in the args if present, else will
+        # use and remove from kwargs else will use True/utf-8
+        is_text = self._includefile_overrides.get(
+            "is_text", self.kwargs.pop("is_text", True)
+        )
+        encoding = self._includefile_overrides.get(
+            "encoding", self.kwargs.pop("encoding", "utf-8")
+        )
+
         # If a default is specified, it needs to be uploaded when the flow is deployed
         # (for example when doing a `step-functions create`) so we make the default
         # be a DeployTimeField. This means that it will be evaluated in two cases:
@@ -275,7 +311,7 @@ class IncludeFile(Parameter):
         # In the first case, we will need to fully upload the file whereas in the
         # second case, we can just return the string as the FilePath.convert method
         # will take care of evaluating things.
-        v = kwargs.get("default")
+        v = self.kwargs.get("default")
         if v is not None:
             # If the default is a callable, we have two DeployTimeField:
             #  - the callable nature of the default will require us to "call" the default
@@ -288,23 +324,15 @@ class IncludeFile(Parameter):
             # (call the default)
             if callable(v) and not isinstance(v, DeployTimeField):
                 # If default is a callable, make it a DeployTimeField (the inner one)
-                v = DeployTimeField(name, str, "default", v, return_str=True)
-            kwargs["default"] = DeployTimeField(
-                name,
+                v = DeployTimeField(self.name, str, "default", v, return_str=True)
+            self.kwargs["default"] = DeployTimeField(
+                self.name,
                 str,
                 "default",
                 IncludeFile._eval_default(is_text, encoding, v),
                 print_representation=v,
             )
 
-        super(IncludeFile, self).__init__(
-            name,
-            required=required,
-            help=help,
-            type=FilePathClass(is_text, encoding),
-            **kwargs,
-        )
-
     def load_parameter(self, v):
         if v is None:
             return v