metaflow 2.12.36__py2.py3-none-any.whl → 2.12.38__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,12 @@ from .extension_support import extension_info
 
 from .graph import FlowGraph
 from .unbounded_foreach import UnboundedForeachInput
+from .user_configs.config_decorators import (
+    CustomFlowDecorator,
+    CustomStepDecorator,
+    MutableFlow,
+    MutableStep,
+)
 from .util import to_pod
 from .metaflow_config import INCLUDE_FOREACH_STACK, MAXIMUM_FOREACH_VALUE_CHARS
 
@@ -65,14 +74,27 @@ 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
+        # 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.
         f._flow_decorators = {}
+
+        # Keys are _FlowState enum values
+        f._flow_state = {}
+
         return f
 
 
@@ -96,6 +118,7 @@ class FlowSpec(metaclass=FlowSpecMeta):
         "_cached_input",
         "_graph",
         "_flow_decorators",
+        "_flow_state",
         "_steps",
         "index",
         "input",
@@ -148,16 +171,11 @@ 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
+    def _check_parameters(self, config_parameters=False):
         seen = set()
-        for var, param in self._get_parameters():
+        for _, param in self._get_parameters():
+            if param.IS_CONFIG_PARAMETER != config_parameters:
+                continue
             norm = param.name.lower()
             if norm in seen:
                 raise MetaflowException(
@@ -166,13 +184,127 @@ class FlowSpec(metaclass=FlowSpecMeta):
                     "case-insensitive." % param.name
                 )
             seen.add(norm)
-        seen.clear()
+
+    def _process_config_decorators(self, config_options):
+        current_cls = self.__class__
+
+        # Fast path for no user configurations
+        if not self._flow_state.get(_FlowState.CONFIG_DECORATORS):
+            # Process parameters to allow them to also use config values easily
+            for var, param in self._get_parameters():
+                if param.IS_CONFIG_PARAMETER:
+                    continue
+                param.init()
+            return self
+
+        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 = []
+        self._check_parameters(config_parameters=True)
+        for var, param in self._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(current_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 self._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(current_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(current_cls, step.name, step)
+
+        mutable_flow = MutableFlow(current_cls)
+        for deco in self._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 == current_cls and not issubclass(
+                    current_cls, deco._flow_cls
+                ):
+                    raise MetaflowInternalError(
+                        "CustomFlowDecorator registered on the wrong flow -- "
+                        "expected %s but got %s"
+                        % (deco._flow_cls.__name__, current_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 current_cls._flow_state:
+                    del current_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 self._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(current_cls, var, param)
+
+        # Reset cached parameters again since we added back the config parameters
+        if _FlowState.CACHED_PARAMETERS in current_cls._flow_state:
+            del current_cls._flow_state[_FlowState.CACHED_PARAMETERS]
+
+        # Set the current flow class we are in (the one we just created)
+        parameters.replace_flow_context(current_cls)
+        return current_cls(use_cli=False)
+
+    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()
@@ -218,6 +350,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 +364,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/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):
+        super(IncludeFile, self).init()
+
+        # 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

metaflow/metaflow_config.py

@@ -440,7 +440,7 @@ ESCAPE_HATCH_WARNING = from_conf("ESCAPE_HATCH_WARNING", True)
 ###
 # Debug configuration
 ###
-DEBUG_OPTIONS = ["subcommand", "sidecar", "s3client", "tracing", "stubgen"]
+DEBUG_OPTIONS = ["subcommand", "sidecar", "s3client", "tracing", "stubgen", "userconf"]
 
 for typ in DEBUG_OPTIONS:
     vars()["DEBUG_%s" % typ.upper()] = from_conf("DEBUG_%s" % typ.upper(), False)

metaflow/package.py

@@ -6,6 +6,7 @@ import time
 import json
 from io import BytesIO
 
+from .user_configs.config_parameters import CONFIG_FILE, dump_config_values
 from .extension_support import EXT_PKG, package_mfext_all
 from .metaflow_config import DEFAULT_PACKAGE_SUFFIXES
 from .exception import MetaflowException
@@ -151,11 +152,23 @@ class MetaflowPackage(object):
             for path_tuple in self._walk(flowdir, suffixes=self.suffixes):
                 yield path_tuple
 
+    def _add_configs(self, tar):
+        buf = BytesIO()
+        buf.write(json.dumps(dump_config_values(self._flow)).encode("utf-8"))
+        self._add_file(tar, os.path.basename(CONFIG_FILE), buf)
+
     def _add_info(self, tar):
-        info = tarfile.TarInfo(os.path.basename(INFO_FILE))
-        env = self.environment.get_environment_info(include_ext_info=True)
         buf = BytesIO()
-        buf.write(json.dumps(env).encode("utf-8"))
+        buf.write(
+            json.dumps(
+                self.environment.get_environment_info(include_ext_info=True)
+            ).encode("utf-8")
+        )
+        self._add_file(tar, os.path.basename(INFO_FILE), buf)
+
+    @staticmethod
+    def _add_file(tar, filename, buf):
+        info = tarfile.TarInfo(filename)
         buf.seek(0)
         info.size = len(buf.getvalue())
         # Setting this default to Dec 3, 2019
@@ -175,6 +188,7 @@ class MetaflowPackage(object):
             fileobj=buf, mode="w:gz", compresslevel=3, dereference=True
         ) as tar:
             self._add_info(tar)
+            self._add_configs(tar)
             for path, arcname in self.path_tuples():
                 tar.add(path, arcname=arcname, recursive=False, filter=no_mtime)
 

metaflow/parameters.py

@@ -3,7 +3,7 @@ import json
 from contextlib import contextmanager
 from threading import local
 
-from typing import Any, Callable, Dict, NamedTuple, Optional, Type, Union
+from typing import Any, Callable, Dict, NamedTuple, Optional, TYPE_CHECKING, Type, Union
 
 from metaflow._vendor import click
 
@@ -14,6 +14,9 @@ from .exception import (
     MetaflowException,
 )
 
+if TYPE_CHECKING:
+    from .user_configs.config_parameters import ConfigValue
+
 try:
     # Python2
     strtype = basestring
@@ -32,6 +35,7 @@ ParameterContext = NamedTuple(
         ("parameter_name", str),
         ("logger", Callable[..., None]),
         ("ds_type", str),
+        ("configs", Optional["ConfigValue"]),
     ],
 )
 
@@ -72,6 +76,16 @@ def flow_context(flow_cls):
 context_proto = None
 
 
+def replace_flow_context(flow_cls):
+    """
+    Replace the current flow context with a new flow class. This is used
+    when we change the current flow class after having run user configuration functions
+    """
+    current_flow.flow_cls_stack = current_flow.flow_cls_stack[1:]
+    current_flow.flow_cls_stack.insert(0, flow_cls)
+    current_flow.flow_cls = current_flow.flow_cls_stack[0]
+
+
 class JSONTypeClass(click.ParamType):
     name = "JSON"
 
@@ -210,12 +224,18 @@ class DeployTimeField(object):
 def deploy_time_eval(value):
     if isinstance(value, DeployTimeField):
         return value(deploy_time=True)
+    elif isinstance(value, DelayedEvaluationParameter):
+        return value(return_str=True)
     else:
         return value
 
 
 # this is called by cli.main
-def set_parameter_context(flow_name, echo, datastore):
+def set_parameter_context(flow_name, echo, datastore, configs):
+    from .user_configs.config_parameters import (
+        ConfigValue,
+    )  # Prevent circular dependency
+
     global context_proto
     context_proto = ParameterContext(
         flow_name=flow_name,
@@ -223,6 +243,7 @@ def set_parameter_context(flow_name, echo, datastore):
         parameter_name=None,
         logger=echo,
         ds_type=datastore.TYPE,
+        configs=ConfigValue(dict(configs)),
     )
 
 
@@ -279,7 +300,11 @@ class Parameter(object):
     ----------
     name : str
         User-visible parameter name.
-    default : str or float or int or bool or `JSONType` or a function.
+    default : Union[str, float, int, bool, Dict[str, Any],
+                Callable[
+                    [ParameterContext], Union[str, float, int, bool, Dict[str, Any]]
+                ],
+            ], optional, default None
         Default value for the parameter. Use a special `JSONType` class to
         indicate that the value must be a valid JSON object. A function
         implies that the parameter corresponds to a *deploy-time parameter*.
@@ -288,15 +313,19 @@ class Parameter(object):
         If `default` is not specified, define the parameter type. Specify
         one of `str`, `float`, `int`, `bool`, or `JSONType`. If None, defaults
         to the type of `default` or `str` if none specified.
-    help : str, optional
+    help : str, optional, default None
         Help text to show in `run --help`.
-    required : bool, default False
-        Require that the user specified a value for the parameter.
-        `required=True` implies that the `default` is not used.
-    show_default : bool, default True
-        If True, show the default value in the help text.
+    required : bool, optional, default None
+        Require that the user specified a value for the parameter. Note that if
+        a default is provide, the required flag is ignored.
+        A value of None is equivalent to False.
+    show_default : bool, optional, default None
+        If True, show the default value in the help text. A value of None is equivalent
+        to True.
     """
 
+    IS_CONFIG_PARAMETER = False
+
     def __init__(
         self,
         name: str,
@@ -307,31 +336,53 @@ class Parameter(object):
                 int,
                 bool,
                 Dict[str, Any],
-                Callable[[], Union[str, float, int, bool, Dict[str, Any]]],
+                Callable[
+                    [ParameterContext], Union[str, float, int, bool, Dict[str, Any]]
+                ],
             ]
         ] = None,
         type: Optional[
             Union[Type[str], Type[float], Type[int], Type[bool], JSONTypeClass]
         ] = None,
         help: Optional[str] = None,
-        required: bool = False,
-        show_default: bool = True,
+        required: Optional[bool] = None,
+        show_default: Optional[bool] = None,
         **kwargs: Dict[str, Any]
     ):
         self.name = name
         self.kwargs = kwargs
-        for k, v in {
+        self._override_kwargs = {
             "default": default,
             "type": type,
             "help": help,
             "required": required,
             "show_default": show_default,
-        }.items():
-            if v is not None:
-                self.kwargs[k] = v
+        }
+
+    def init(self):
+        # Prevent circular import
+        from .user_configs.config_parameters import (
+            resolve_delayed_evaluator,
+            unpack_delayed_evaluator,
+        )
+
+        # Resolve any value from configurations
+        self.kwargs = unpack_delayed_evaluator(self.kwargs)
+        self.kwargs = resolve_delayed_evaluator(self.kwargs)
+
+        # This was the behavior before configs: values specified in args would override
+        # stuff in kwargs which is what we implement here as well
+        for key, value in self._override_kwargs.items():
+            if value is not None:
+                self.kwargs[key] = value
+        # Set two default values if no-one specified them
+        self.kwargs.setdefault("required", False)
+        self.kwargs.setdefault("show_default", True)
+
+        # Continue processing kwargs free of any configuration values :)
 
         # TODO: check that the type is one of the supported types
-        param_type = self.kwargs["type"] = self._get_type(kwargs)
+        param_type = self.kwargs["type"] = self._get_type(self.kwargs)
 
         reserved_params = [
             "params",
@@ -356,23 +407,27 @@ class Parameter(object):
             raise MetaflowException(
                 "Parameter name '%s' is a reserved "
                 "word. Please use a different "
-                "name for your parameter." % (name)
+                "name for your parameter." % (self.name)
             )
 
         # make sure the user is not trying to pass a function in one of the
         # fields that don't support function-values yet
         for field in ("show_default", "separator", "required"):
-            if callable(kwargs.get(field)):
+            if callable(self.kwargs.get(field)):
                 raise MetaflowException(
                     "Parameter *%s*: Field '%s' cannot "
-                    "have a function as its value" % (name, field)
+                    "have a function as its value" % (self.name, field)
                 )
 
         # default can be defined as a function
         default_field = self.kwargs.get("default")
         if callable(default_field) and not isinstance(default_field, DeployTimeField):
             self.kwargs["default"] = DeployTimeField(
-                name, param_type, "default", self.kwargs["default"], return_str=True
+                self.name,
+                param_type,
+                "default",
+                self.kwargs["default"],
+                return_str=True,
             )
 
         # note that separator doesn't work with DeployTimeFields unless you
@@ -381,7 +436,7 @@ class Parameter(object):
         if self.separator and not self.is_string_type:
             raise MetaflowException(
                 "Parameter *%s*: Separator is only allowed "
-                "for string parameters." % name
+                "for string parameters." % self.name
             )
 
     def __repr__(self):
@@ -437,7 +492,9 @@ def add_custom_parameters(deploy_mode=False):
         flow_cls = getattr(current_flow, "flow_cls", None)
         if flow_cls is None:
             return cmd
-        parameters = [p for _, p in flow_cls._get_parameters()]
+        parameters = [
+            p for _, p in flow_cls._get_parameters() if not p.IS_CONFIG_PARAMETER
+        ]
         for arg in parameters[::-1]:
             kwargs = arg.option_kwargs(deploy_mode)
             cmd.params.insert(0, click.Option(("--" + arg.name,), **kwargs))

metaflow/plugins/__init__.py

@@ -164,6 +164,10 @@ def get_plugin_cli():
     return resolve_plugins("cli")
 
 
+def get_plugin_cli_path():
+    return resolve_plugins("cli", path_only=True)
+
+
 STEP_DECORATORS = resolve_plugins("step_decorator")
 FLOW_DECORATORS = resolve_plugins("flow_decorator")
 ENVIRONMENTS = resolve_plugins("environment")

metaflow/plugins/airflow/airflow_cli.py

@@ -283,6 +283,7 @@ def make_flow(
 ):
     # Attach @kubernetes.
     decorators._attach_decorators(obj.flow, [KubernetesDecorator.name])
+    decorators._init(obj.flow)
 
     decorators._init_step_decorators(
         obj.flow, obj.graph, obj.environment, obj.flow_datastore, obj.logger