metaflow 2.15.21__py2.py3-none-any.whl → 2.16.0__py2.py3-none-any.whl

metaflow/extension_support/__init__.py

@@ -12,7 +12,7 @@ from importlib.abc import MetaPathFinder, Loader
 from itertools import chain
 from pathlib import Path
 
-from metaflow.info_file import read_info_file
+from metaflow.meta_files import read_info_file
 
 
 #
@@ -214,6 +214,10 @@ def package_mfext_all():
             yield path_tuple
 
 
+def package_mfext_all_descriptions():
+    return _all_packages
+
+
 def load_globals(module, dst_globals, extra_indent=False):
     if extra_indent:
         extra_indent = "    "
@@ -808,13 +812,16 @@ def _get_extension_packages(ignore_info_file=False, restrict_to_directories=None
                                 "    Extends '%s' with config '%s'"
                                 % (_extension_points[idx], config_module)
                             )
-            mf_pkg_list.append(package_name)
-            mf_ext_packages[package_name] = {
-                "root_paths": [package_path],
-                "meta_module": meta_module,
-                "files": files_to_include,
-                "version": "_local_",
-            }
+            if files_to_include:
+                mf_pkg_list.append(package_name)
+                mf_ext_packages[package_name] = {
+                    "root_paths": [package_path],
+                    "meta_module": meta_module,
+                    "files": files_to_include,
+                    "version": "_local_",
+                }
+            else:
+                _ext_debug("Skipping package as no files found (empty dir?)")
 
     # Sanity check that we only have one package per configuration file.
     # This prevents multiple packages from providing the same named configuration

metaflow/extension_support/_empty_file.py

@@ -1,2 +1,2 @@
-# This file serves as a __init__.py for metaflow_extensions when it is packaged
-# and needs to remain empty.
+# This file serves as a __init__.py for metaflow_extensions or metaflow
+# packages when they are packaged and needs to remain empty.

metaflow/flowspec.py

@@ -7,7 +7,7 @@ import reprlib
 from enum import Enum
 from itertools import islice
 from types import FunctionType, MethodType
-from typing import TYPE_CHECKING, Any, Callable, Generator, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple
 
 from . import cmd_with_io, parameters
 from .debug import debug
@@ -23,13 +23,14 @@ 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 .user_configs.config_parameters import ConfigValue
+
+from .user_decorators.mutable_flow import MutableFlow
+from .user_decorators.mutable_step import MutableStep
+from .user_decorators.user_flow_decorator import FlowMutator
+from .user_decorators.user_step_decorator import StepMutator
+
+
 from .util import to_pod
 from .metaflow_config import INCLUDE_FOREACH_STACK, MAXIMUM_FOREACH_VALUE_CHARS
 
@@ -79,6 +80,10 @@ class _FlowState(Enum):
     CONFIGS = 1
     CONFIG_DECORATORS = 2
     CACHED_PARAMETERS = 3
+    SET_CONFIG_PARAMETERS = (
+        4  # These are Parameters that now have a ConfigValue (converted)
+    )
+    # but we need to remember them.
 
 
 class FlowSpecMeta(type):
@@ -136,6 +141,7 @@ class FlowSpecMeta(type):
                     cls._flow_state.setdefault(_FlowState.CONFIG_DECORATORS, []).extend(
                         base_configs
                     )
+
         cls._init_graph()
 
     def _init_graph(cls):
@@ -232,8 +238,8 @@ class FlowSpec(metaclass=FlowSpecMeta):
 
     @classmethod
     def _process_config_decorators(cls, config_options, process_configs=True):
-
         if cls._configs_processed:
+            debug.userconf_exec("Mutating step/flow decorators already processed")
             return None
         cls._configs_processed = True
 
@@ -244,18 +250,20 @@ class FlowSpec(metaclass=FlowSpecMeta):
         ):
             # Process parameters to allow them to also use config values easily
             for var, param in cls._get_parameters():
-                if param.IS_CONFIG_PARAMETER:
+                if isinstance(param, ConfigValue) or param.IS_CONFIG_PARAMETER:
                     continue
                 param.init(not process_configs)
             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.
+        # to actual values so they can be used as is in the mutators.
 
-        # We then reset them to be proper configs so they can be re-evaluated in
-        # _set_constants
-        to_reset_configs = []
+        # We, however, need to make sure _get_parameters still works properly so
+        # we store what was a config and has been set to a specific value.
+        # This is safe to do for now because all other uses of _get_parameters typically
+        # do not rely on the variable itself but just the parameter.
+        to_save_configs = []
         cls._check_parameters(config_parameters=True)
         for var, param in cls._get_parameters():
             if not param.IS_CONFIG_PARAMETER:
@@ -267,71 +275,72 @@ class FlowSpec(metaclass=FlowSpecMeta):
             # 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))
+            to_save_configs.append((var, param))
             debug.userconf_exec("Setting config %s to %s" % (var, str(val)))
             setattr(cls, var, val)
 
-        # Reset cached parameters since we have replaced configs already with ConfigValue
-        # so they are not parameters anymore to be re-evaluated when we do _get_parameters
-        if _FlowState.CACHED_PARAMETERS in cls._flow_state:
-            del cls._flow_state[_FlowState.CACHED_PARAMETERS]
+        cls._flow_state[_FlowState.SET_CONFIG_PARAMETERS] = to_save_configs
+        # Run all the decorators. We first run the flow-level decorators
+        # and then the step level ones to maintain a consistent order with how
+        # other decorators are run.
 
-        # 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.get(_FlowState.CONFIG_DECORATORS, []):
-            if isinstance(deco, CustomFlowDecorator):
+            if isinstance(deco, FlowMutator):
+                inserted_by_value = [deco.decorator_name] + (deco.inserted_by or [])
+                mutable_flow = MutableFlow(
+                    cls,
+                    pre_mutate=True,
+                    statically_defined=deco.statically_defined,
+                    inserted_by=inserted_by_value,
+                )
                 # 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 -- "
+                        "FlowMutator 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)
+                deco.pre_mutate(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"
+                    "A non FlowMutator found in flow custom decorators"
                 )
 
+        for step in cls._steps:
+            for deco in step.config_decorators:
+                if isinstance(deco, StepMutator):
+                    inserted_by_value = [deco.decorator_name] + (deco.inserted_by or [])
+                    debug.userconf_exec(
+                        "Evaluating step level decorator %s for %s"
+                        % (deco.__class__.__name__, step.name)
+                    )
+                    deco.pre_mutate(
+                        MutableStep(
+                            cls,
+                            step,
+                            pre_mutate=True,
+                            statically_defined=deco.statically_defined,
+                            inserted_by=inserted_by_value,
+                        )
+                    )
+                else:
+                    raise MetaflowInternalError(
+                        "A non StepMutator found in step 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)
@@ -402,9 +411,19 @@ class FlowSpec(metaclass=FlowSpecMeta):
                     "name": deco.name,
                     "attributes": to_pod(deco.attributes),
                     "statically_defined": deco.statically_defined,
+                    "inserted_by": deco.inserted_by,
                 }
                 for deco in flow_decorators(self)
                 if not deco.name.startswith("_")
+            ]
+            + [
+                {
+                    "name": deco.__class__.__name__,
+                    "attributes": {},
+                    "statically_defined": deco.statically_defined,
+                    "inserted_by": deco.inserted_by,
+                }
+                for deco in self._flow_state.get(_FlowState.CONFIG_DECORATORS, [])
             ],
             "extensions": extension_info(),
         }
@@ -413,11 +432,19 @@ class FlowSpec(metaclass=FlowSpecMeta):
     @classmethod
     def _get_parameters(cls):
         cached = cls._flow_state.get(_FlowState.CACHED_PARAMETERS)
+        returned = set()
         if cached is not None:
+            for set_config in cls._flow_state.get(_FlowState.SET_CONFIG_PARAMETERS, []):
+                returned.add(set_config[0])
+                yield set_config[0], set_config[1]
             for var in cached:
-                yield var, getattr(cls, var)
+                if var not in returned:
+                    yield var, getattr(cls, var)
             return
         build_list = []
+        for set_config in cls._flow_state.get(_FlowState.SET_CONFIG_PARAMETERS, []):
+            returned.add(set_config[0])
+            yield set_config[0], set_config[1]
         for var in dir(cls):
             if var[0] == "_" or var in cls._NON_PARAMETERS:
                 continue
@@ -425,7 +452,7 @@ class FlowSpec(metaclass=FlowSpecMeta):
                 val = getattr(cls, var)
             except:
                 continue
-            if isinstance(val, Parameter):
+            if isinstance(val, Parameter) and var not in returned:
                 build_list.append(var)
                 yield var, val
         cls._flow_state[_FlowState.CACHED_PARAMETERS] = build_list

metaflow/graph.py

@@ -2,6 +2,8 @@ import inspect
 import ast
 import re
 
+from itertools import chain
+
 
 from .util import to_pod
 
@@ -45,13 +47,17 @@ def deindent_docstring(doc):
 
 
 class DAGNode(object):
-    def __init__(self, func_ast, decos, doc, source_file, lineno):
+    def __init__(
+        self, func_ast, decos, wrappers, config_decorators, doc, source_file, lineno
+    ):
         self.name = func_ast.name
         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.wrappers = wrappers
+        self.config_decorators = config_decorators
         self.doc = deindent_docstring(doc)
         self.parallel_step = any(getattr(deco, "IS_PARALLEL", False) for deco in decos)
 
@@ -181,7 +187,13 @@ class FlowGraph(object):
                 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
+                    function_ast,
+                    func.decorators,
+                    func.wrappers,
+                    func.config_decorators,
+                    func.__doc__,
+                    source_file,
+                    lineno,
                 )
                 nodes[element] = node
         return nodes
@@ -293,9 +305,19 @@ class FlowGraph(object):
                         "name": deco.name,
                         "attributes": to_pod(deco.attributes),
                         "statically_defined": deco.statically_defined,
+                        "inserted_by": deco.inserted_by,
                     }
                     for deco in node.decorators
                     if not deco.name.startswith("_")
+                ]
+                + [
+                    {
+                        "name": deco.decorator_name,
+                        "attributes": {"_args": deco._args, **deco._kwargs},
+                        "statically_defined": deco.statically_defined,
+                        "inserted_by": deco.inserted_by,
+                    }
+                    for deco in chain(node.wrappers, node.config_decorators)
                 ],
                 "next": node.out_funcs,
             }

metaflow/meta_files.py

@@ -0,0 +1,13 @@
+_UNINITIALIZED = object()
+_info_file_content = _UNINITIALIZED
+
+
+def read_info_file():
+    # Prevent circular import
+    from .packaging_sys import MetaflowCodeContent
+
+    global _info_file_content
+
+    if id(_info_file_content) == id(_UNINITIALIZED):
+        _info_file_content = MetaflowCodeContent.get_info()
+    return _info_file_content

metaflow/metadata_provider/metadata.py

@@ -674,11 +674,17 @@ class MetadataProvider(object):
         if code_sha:
             code_url = os.environ.get("METAFLOW_CODE_URL")
             code_ds = os.environ.get("METAFLOW_CODE_DS")
+            code_metadata = os.environ.get("METAFLOW_CODE_METADATA")
             metadata.append(
                 MetaDatum(
                     field="code-package",
                     value=json.dumps(
-                        {"ds_type": code_ds, "sha": code_sha, "location": code_url}
+                        {
+                            "ds_type": code_ds,
+                            "sha": code_sha,
+                            "location": code_url,
+                            "metadata": code_metadata,
+                        }
                     ),
                     type="code-package",
                     tags=["attempt_id:{0}".format(attempt)],

metaflow/metaflow_config.py

@@ -449,6 +449,10 @@ CONDA_USE_FAST_INIT = from_conf("CONDA_USE_FAST_INIT", False)
 # Print out warning if escape hatch is not used for the target packages
 ESCAPE_HATCH_WARNING = from_conf("ESCAPE_HATCH_WARNING", True)
 
+###
+# Features
+###
+FEAT_ALWAYS_UPLOAD_CODE_PACKAGE = from_conf("FEAT_ALWAYS_UPLOAD_CODE_PACKAGE", False)
 ###
 # Debug configuration
 ###
@@ -460,6 +464,7 @@ DEBUG_OPTIONS = [
     "stubgen",
     "userconf",
     "conda",
+    "package",
 ]
 
 for typ in DEBUG_OPTIONS:

metaflow/metaflow_environment.py

@@ -1,3 +1,4 @@
+import json
 import os
 import platform
 import sys
@@ -8,6 +9,8 @@ from . import metaflow_git
 from metaflow.exception import MetaflowException
 from metaflow.extension_support import dump_module_info
 from metaflow.mflog import BASH_MFLOG, BASH_FLUSH_LOGS
+from metaflow.package import MetaflowPackage
+
 from . import R
 
 
@@ -49,8 +52,36 @@ class MetaflowEnvironment(object):
 
     def add_to_package(self):
         """
-        A list of tuples (file, arcname) to add to the job package.
-        `arcname` is an alternative name for the file in the job package.
+        Called to add custom files needed for this environment. This hook will be
+        called in the `MetaflowPackage` class where metaflow compiles the code package
+        tarball. This hook can return one of two things (the first is for backwards
+        compatibility -- move to the second):
+          - a generator yielding a tuple of `(file_path, arcname)` to add files to
+            the code package. `file_path` is the path to the file on the local filesystem
+            and `arcname` is the path relative to the packaged code.
+          - a generator yielding a tuple of `(content, arcname, type)` where:
+            - type is one of
+            ContentType.{USER_CONTENT, CODE_CONTENT, MODULE_CONTENT, OTHER_CONTENT}
+            - for USER_CONTENT:
+              - the file will be included relative to the directory containing the
+                user's flow file.
+              - content: path to the file to include
+              - arcname: path relative to the directory containing the user's flow file
+            - for CODE_CONTENT:
+              - the file will be included relative to the code directory in the package.
+                This will be the directory containing `metaflow`.
+              - content: path to the file to include
+              - arcname: path relative to the code directory in the package
+            - for MODULE_CONTENT:
+              - the module will be added to the code package as a python module. It will
+                be accessible as usual (import <module_name>)
+              - content: name of the module
+              - arcname: None (ignored)
+            - for OTHER_CONTENT:
+              - the file will be included relative to any other configuration/metadata
+                files for the flow
+              - content: path to the file to include
+              - arcname: path relative to the config directory in the package
         """
         return []
 
@@ -157,29 +188,55 @@ class MetaflowEnvironment(object):
         # skip pip installs if we know that packages might already be available
         return "if [ -z $METAFLOW_SKIP_INSTALL_DEPENDENCIES ]; then {}; fi".format(cmd)
 
-    def get_package_commands(self, code_package_url, datastore_type):
-        cmds = [
-            BASH_MFLOG,
-            BASH_FLUSH_LOGS,
-            "mflog 'Setting up task environment.'",
-            self._get_install_dependencies_cmd(datastore_type),
-            "mkdir metaflow",
-            "cd metaflow",
-            "mkdir .metaflow",  # mute local datastore creation log
-            "i=0; while [ $i -le 5 ]; do "
-            "mflog 'Downloading code package...'; "
-            + self._get_download_code_package_cmd(code_package_url, datastore_type)
-            + " && mflog 'Code package downloaded.' && break; "
-            "sleep 10; i=$((i+1)); "
-            "done",
-            "if [ $i -gt 5 ]; then "
-            "mflog 'Failed to download code package from %s "
-            "after 6 tries. Exiting...' && exit 1; "
-            "fi" % code_package_url,
-            "TAR_OPTIONS='--warning=no-timestamp' tar xf job.tar",
-            "mflog 'Task is starting.'",
-            "flush_mflogs",
-        ]
+    def get_package_commands(
+        self, code_package_url, datastore_type, code_package_metadata=None
+    ):
+        # HACK: We want to keep forward compatibility with compute layers so that
+        # they can still call get_package_commands and NOT pass any metadata. If
+        # there is no additional information, we *assume* that it is the default
+        # used.
+        if code_package_metadata is None:
+            code_package_metadata = json.dumps(
+                {
+                    "version": 0,
+                    "archive_format": "tgz",
+                    "mfcontent_version": 1,
+                }
+            )
+        cmds = (
+            [
+                BASH_MFLOG,
+                BASH_FLUSH_LOGS,
+                "mflog 'Setting up task environment.'",
+                self._get_install_dependencies_cmd(datastore_type),
+                "mkdir metaflow",
+                "cd metaflow",
+                "mkdir .metaflow",  # mute local datastore creation log
+                "i=0; while [ $i -le 5 ]; do "
+                "mflog 'Downloading code package...'; "
+                + self._get_download_code_package_cmd(code_package_url, datastore_type)
+                + " && mflog 'Code package downloaded.' && break; "
+                "sleep 10; i=$((i+1)); "
+                "done",
+                "if [ $i -gt 5 ]; then "
+                "mflog 'Failed to download code package from %s "
+                "after 6 tries. Exiting...' && exit 1; "
+                "fi" % code_package_url,
+            ]
+            + MetaflowPackage.get_extract_commands(
+                code_package_metadata, "job.tar", dest_dir="."
+            )
+            + [
+                "export %s=%s:$(printenv %s)" % (k, v.replace('"', '\\"'), k)
+                for k, v in MetaflowPackage.get_post_extract_env_vars(
+                    code_package_metadata, dest_dir="."
+                ).items()
+            ]
+            + [
+                "mflog 'Task is starting.'",
+                "flush_mflogs",
+            ]
+        )
         return cmds
 
     def get_environment_info(self, include_ext_info=False):

metaflow/metaflow_version.py

@@ -11,7 +11,7 @@ import subprocess
 from os import path, name, environ, listdir
 
 from metaflow.extension_support import update_package_info
-from metaflow.info_file import CURRENT_DIRECTORY, read_info_file
+from metaflow.meta_files import read_info_file
 
 
 # True/False correspond to the value `public`` in get_version