ob-metaflow 2.16.8.2rc2__py2.py3-none-any.whl → 2.17.0.1__py2.py3-none-any.whl

metaflow/graph.py

@@ -68,8 +68,6 @@ class DAGNode(object):
         self.has_tail_next = False
         self.invalid_tail_next = False
         self.num_args = 0
-        self.switch_cases = {}
-        self.condition = None
         self.foreach_param = None
         self.num_parallel = 0
         self.parallel_foreach = False
@@ -85,56 +83,6 @@ class DAGNode(object):
     def _expr_str(self, expr):
         return "%s.%s" % (expr.value.id, expr.attr)
 
-    def _parse_switch_dict(self, dict_node):
-        switch_cases = {}
-
-        if isinstance(dict_node, ast.Dict):
-            for key, value in zip(dict_node.keys, dict_node.values):
-                case_key = None
-
-                # handle string literals
-                if isinstance(key, ast.Str):
-                    case_key = key.s
-                elif isinstance(key, ast.Constant) and isinstance(key.value, str):
-                    case_key = key.value
-                elif isinstance(key, ast.Attribute):
-                    if isinstance(key.value, ast.Attribute) and isinstance(
-                        key.value.value, ast.Name
-                    ):
-                        # This handles self.config.some_key
-                        if key.value.value.id == "self":
-                            config_var = key.value.attr
-                            config_key = key.attr
-                            case_key = f"config:{config_var}.{config_key}"
-                        else:
-                            return None
-                    else:
-                        return None
-
-                # handle variables or other dynamic expressions - not allowed
-                elif isinstance(key, ast.Name):
-                    return None
-                else:
-                    # can't statically analyze this key
-                    return None
-
-                if case_key is None:
-                    return None
-
-                # extract the step name from the value
-                if isinstance(value, ast.Attribute) and isinstance(
-                    value.value, ast.Name
-                ):
-                    if value.value.id == "self":
-                        step_name = value.attr
-                        switch_cases[case_key] = step_name
-                    else:
-                        return None
-                else:
-                    return None
-
-        return switch_cases if switch_cases else None
-
     def _parse(self, func_ast, lineno):
         self.num_args = len(func_ast.args.args)
         tail = func_ast.body[-1]
@@ -156,38 +104,7 @@ class DAGNode(object):
             self.has_tail_next = True
             self.invalid_tail_next = True
             self.tail_next_lineno = lineno + tail.lineno - 1
-
-            # Check if first argument is a dictionary (switch case)
-            if (
-                len(tail.value.args) == 1
-                and isinstance(tail.value.args[0], ast.Dict)
-                and any(k.arg == "condition" for k in tail.value.keywords)
-            ):
-                # This is a switch statement
-                switch_cases = self._parse_switch_dict(tail.value.args[0])
-                condition_name = None
-
-                # Get condition parameter
-                for keyword in tail.value.keywords:
-                    if keyword.arg == "condition":
-                        if isinstance(keyword.value, ast.Str):
-                            condition_name = keyword.value.s
-                        elif isinstance(keyword.value, ast.Constant) and isinstance(
-                            keyword.value.value, str
-                        ):
-                            condition_name = keyword.value.value
-                        break
-
-                if switch_cases and condition_name:
-                    self.type = "split-switch"
-                    self.condition = condition_name
-                    self.switch_cases = switch_cases
-                    self.out_funcs = list(switch_cases.values())
-                    self.invalid_tail_next = False
-                    return
-
-            else:
-                self.out_funcs = [e.attr for e in tail.value.args]
+            self.out_funcs = [e.attr for e in tail.value.args]
 
             keywords = dict(
                 (k.arg, getattr(k.value, "s", None)) for k in tail.value.keywords
@@ -234,7 +151,6 @@ class DAGNode(object):
     has_tail_next={0.has_tail_next} (line {0.tail_next_lineno})
     invalid_tail_next={0.invalid_tail_next}
     foreach_param={0.foreach_param}
-    condition={0.condition}
     parallel_step={0.parallel_step}
     parallel_foreach={0.parallel_foreach}
     -> {out}""".format(
@@ -303,8 +219,6 @@ class FlowGraph(object):
             if node.type in ("split", "foreach"):
                 node.split_parents = split_parents
                 split_parents = split_parents + [node.name]
-            elif node.type == "split-switch":
-                node.split_parents = split_parents
             elif node.type == "join":
                 # ignore joins without splits
                 if split_parents:
@@ -345,37 +259,15 @@ class FlowGraph(object):
     def output_dot(self):
         def edge_specs():
             for node in self.nodes.values():
-                if node.type == "split-switch":
-                    # Label edges for switch cases
-                    for case_value, step_name in node.switch_cases.items():
-                        yield (
-                            '{0} -> {1} [label="{2}" color="blue" fontcolor="blue"];'.format(
-                                node.name, step_name, case_value
-                            )
-                        )
-                else:
-                    for edge in node.out_funcs:
-                        yield "%s -> %s;" % (node.name, edge)
+                for edge in node.out_funcs:
+                    yield "%s -> %s;" % (node.name, edge)
 
         def node_specs():
             for node in self.nodes.values():
-                if node.type == "split-switch":
-                    # Hexagon shape for switch nodes
-                    condition_label = (
-                        f"switch: {node.condition}" if node.condition else "switch"
-                    )
-                    yield (
-                        '"{0.name}" '
-                        '[ label = <<b>{0.name}</b><br/><font point-size="9">{condition}</font>> '
-                        '  fontname = "Helvetica" '
-                        '  shape = "hexagon" '
-                        '  style = "filled" fillcolor = "lightgreen" ];'
-                    ).format(node, condition=condition_label)
-                else:
-                    nodetype = "join" if node.num_args > 1 else node.type
-                    yield '"{0.name}"' '[ label = <<b>{0.name}</b> | <font point-size="10">{type}</font>> ' '  fontname = "Helvetica" ' '  shape = "record" ];'.format(
-                        node, type=nodetype
-                    )
+                nodetype = "join" if node.num_args > 1 else node.type
+                yield '"{0.name}"' '[ label = <<b>{0.name}</b> | <font point-size="10">{type}</font>> ' '  fontname = "Helvetica" ' '  shape = "record" ];'.format(
+                    node, type=nodetype
+                )
 
         return (
             "digraph {0.name} {{\n"
@@ -399,8 +291,6 @@ class FlowGraph(object):
                 if node.parallel_foreach:
                     return "split-parallel"
                 return "split-foreach"
-            elif node.type == "split-switch":
-                return "split-switch"
             return "unknown"  # Should never happen
 
         def node_to_dict(name, node):
@@ -435,9 +325,6 @@ class FlowGraph(object):
                 d["foreach_artifact"] = node.foreach_param
             elif d["type"] == "split-parallel":
                 d["num_parallel"] = node.num_parallel
-            elif d["type"] == "split-switch":
-                d["condition"] = node.condition
-                d["switch_cases"] = node.switch_cases
             if node.matching_join:
                 d["matching_join"] = node.matching_join
             return d
@@ -452,8 +339,8 @@ class FlowGraph(object):
                 steps_info[cur_name] = node_dict
                 resulting_list.append(cur_name)
 
-                node_type = node_to_type(cur_node)
-                if node_type in ("split-static", "split-foreach"):
+                if cur_node.type not in ("start", "linear", "join"):
+                    # We need to look at the different branches for this
                     resulting_list.append(
                         [
                             populate_block(s, cur_node.matching_join)
@@ -461,19 +348,8 @@ class FlowGraph(object):
                         ]
                     )
                     cur_name = cur_node.matching_join
-                elif node_type == "split-switch":
-                    all_paths = [
-                        populate_block(s, end_name) for s in cur_node.out_funcs
-                    ]
-                    resulting_list.append(all_paths)
-                    cur_name = end_name
                 else:
-                    # handles only linear, start, and join steps.
-                    if cur_node.out_funcs:
-                        cur_name = cur_node.out_funcs[0]
-                    else:
-                        # handles terminal nodes or when we jump to 'end_name'.
-                        break
+                    cur_name = cur_node.out_funcs[0]
             return resulting_list
 
         graph_structure = populate_block("start", "end")

metaflow/lint.py

@@ -1,6 +1,6 @@
 import re
 from .exception import MetaflowException
-from .util import all_equal, get_split_branch_for_node
+from .util import all_equal
 
 
 class LintWarn(MetaflowException):
@@ -134,13 +134,7 @@ def check_valid_transitions(graph):
     msg = (
         "Step *{0.name}* specifies an invalid self.next() transition. "
         "Make sure the self.next() expression matches with one of the "
-        "supported transition types:\n"
-        "  • Linear: self.next(self.step_name)\n"
-        "  • Fan-out: self.next(self.step1, self.step2, ...)\n"
-        "  • Foreach: self.next(self.step, foreach='variable')\n"
-        "  • Switch: self.next({{\"key\": self.step, ...}}, condition='variable')\n\n"
-        "For switch statements, keys must be string literals or config expressions "
-        "(self.config.key_name), not variables or numbers."
+        "supported transition types."
     )
     for node in graph:
         if node.type != "end" and node.has_tail_next and node.invalid_tail_next:
@@ -238,13 +232,7 @@ def check_split_join_balance(graph):
             new_stack = split_stack
         elif node.type in ("split", "foreach"):
             new_stack = split_stack + [("split", node.out_funcs)]
-        elif node.type == "split-switch":
-            # For a switch, continue traversal down each path with the same stack
-            for n in node.out_funcs:
-                traverse(graph[n], split_stack)
-            return
         elif node.type == "end":
-            new_stack = split_stack
             if split_stack:
                 _, split_roots = split_stack.pop()
                 roots = ", ".join(split_roots)
@@ -252,25 +240,10 @@ def check_split_join_balance(graph):
                     msg0.format(roots=roots), node.func_lineno, node.source_file
                 )
         elif node.type == "join":
-            new_stack = split_stack
             if split_stack:
                 _, split_roots = split_stack[-1]
                 new_stack = split_stack[:-1]
-
-                # Identify the split this join corresponds to from its parentage.
-                split_node_name = node.split_parents[-1]
-
-                # Resolve each incoming function to its root branch from the split.
-                resolved_branches = set()
-                for in_node_name in node.in_funcs:
-                    branch = get_split_branch_for_node(
-                        graph, in_node_name, split_node_name
-                    )
-                    if branch:
-                        resolved_branches.add(branch)
-
-                # compares the set of resolved branches against the expected branches from the split.
-                if len(resolved_branches) != len(split_roots):
+                if len(node.in_funcs) != len(split_roots):
                     paths = ", ".join(node.in_funcs)
                     roots = ", ".join(split_roots)
                     raise LintWarn(
@@ -293,8 +266,6 @@ def check_split_join_balance(graph):
 
             if not all_equal(map(parents, node.in_funcs)):
                 raise LintWarn(msg3.format(node), node.func_lineno, node.source_file)
-        else:
-            new_stack = split_stack
 
         for n in node.out_funcs:
             traverse(graph[n], new_stack)
@@ -302,44 +273,6 @@ def check_split_join_balance(graph):
     traverse(graph["start"], [])
 
 
-@linter.ensure_static_graph
-@linter.check
-def check_switch_splits(graph):
-    """Check conditional split constraints"""
-    msg0 = (
-        "Step *{0.name}* is a switch split but defines {num} transitions. "
-        "Switch splits must define at least 2 transitions."
-    )
-    msg1 = "Step *{0.name}* is a switch split but has no condition variable."
-    msg2 = "Step *{0.name}* is a switch split but has no switch cases defined."
-
-    for node in graph:
-        if node.type == "split-switch":
-            # Check at least 2 outputs
-            if len(node.out_funcs) < 2:
-                raise LintWarn(
-                    msg0.format(node, num=len(node.out_funcs)),
-                    node.func_lineno,
-                    node.source_file,
-                )
-
-            # Check condition exists
-            if not node.condition:
-                raise LintWarn(
-                    msg1.format(node),
-                    node.func_lineno,
-                    node.source_file,
-                )
-
-            # Check switch cases exist
-            if not node.switch_cases:
-                raise LintWarn(
-                    msg2.format(node),
-                    node.func_lineno,
-                    node.source_file,
-                )
-
-
 @linter.ensure_static_graph
 @linter.check
 def check_empty_foreaches(graph):

metaflow/metaflow_environment.py

@@ -203,6 +203,19 @@ class MetaflowEnvironment(object):
                     "mfcontent_version": 1,
                 }
             )
+
+        extra_exports = []
+        for k, v in MetaflowPackage.get_post_extract_env_vars(
+            code_package_metadata, dest_dir="$(pwd)"
+        ).items():
+            if k.endswith(":"):
+                # If the value ends with a colon, we override the existing value
+                extra_exports.append("export %s=%s" % (k[:-1], v))
+            else:
+                extra_exports.append(
+                    "export %s=%s:$(printenv %s)" % (k, v.replace('"', '\\"'), k)
+                )
+
         cmds = (
             [
                 BASH_MFLOG,
@@ -226,12 +239,7 @@ class MetaflowEnvironment(object):
             + 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()
-            ]
+            + extra_exports
             + [
                 "mflog 'Task is starting.'",
                 "flush_mflogs",

metaflow/package/__init__.py

@@ -17,7 +17,6 @@ from ..packaging_sys.utils import suffix_filter, walk
 from ..metaflow_config import DEFAULT_PACKAGE_SUFFIXES
 from ..exception import MetaflowException
 from ..user_configs.config_parameters import dump_config_values
-from ..util import get_metaflow_root
 from .. import R
 
 DEFAULT_SUFFIXES_LIST = DEFAULT_PACKAGE_SUFFIXES.split(",")
@@ -76,12 +75,22 @@ class MetaflowPackage(object):
             from ..user_decorators.user_flow_decorator import FlowMutatorMeta
             from ..user_decorators.user_step_decorator import UserStepDecoratorMeta
 
-            if (
-                m.__name__ in FlowMutatorMeta._import_modules
-                or m.__name__ in UserStepDecoratorMeta._import_modules
-                or hasattr(m, "METAFLOW_PACKAGE")
-            ):
-                return True
+            # Be very defensive here to filter modules in case there are
+            # some badly behaved modules that have weird values for
+            # METAFLOW_PACKAGE_POLICY for example.
+            try:
+                if (
+                    m.__name__ in FlowMutatorMeta._import_modules
+                    or m.__name__ in UserStepDecoratorMeta._import_modules
+                    or (
+                        hasattr(m, "METAFLOW_PACKAGE_POLICY")
+                        and m.METAFLOW_PACKAGE_POLICY == "include"
+                    )
+                ):
+                    return True
+                return False
+            except:
+                return False
 
         if mfcontent is None:
             self._mfcontent = MetaflowCodeContentV1(criteria=_module_selector)
@@ -350,10 +359,10 @@ class MetaflowPackage(object):
         """
         backend = cls.get_backend(pkg_metadata)
         with backend.cls_open(archive) as opened_archive:
-            include_names = MetaflowCodeContent.get_archive_content_names(
+            include_members = MetaflowCodeContent.get_archive_content_members(
                 opened_archive, content_types, backend
             )
-            backend.extract_members(include_names, dest_dir)
+            backend.cls_extract_members(opened_archive, include_members, dest_dir)
 
     def user_tuples(self, timeout: Optional[float] = None):
         # Wait for at least the blob to be formed

metaflow/packaging_sys/__init__.py

@@ -118,9 +118,7 @@ class MetaflowCodeContent:
         return handling_cls.get_filename_impl(mfcontent_info, filename, content_type)
 
     @classmethod
-    def get_env_vars_for_packaged_metaflow(
-        cls, dest_dir: str
-    ) -> Optional[Dict[str, str]]:
+    def get_env_vars_for_packaged_metaflow(cls, dest_dir: str) -> Dict[str, str]:
         """
         Get the environment variables that are needed to run Metaflow when it is
         packaged. This is typically used to set the PYTHONPATH to include the
@@ -128,17 +126,19 @@ class MetaflowCodeContent:
 
         Returns
         -------
-        Optional[Dict[str, str]]
+        Dict[str, str]
             The environment variables that are needed to run Metaflow when it is
-            packaged -- None if there are no such variables (not packaged for example)
+            packaged it present.
         """
-        mfcontent_info = cls._extract_mfcontent_info()
+        mfcontent_info = cls._extract_mfcontent_info(dest_dir)
         if mfcontent_info is None:
             # No MFCONTENT_MARKER file found -- this is not a packaged Metaflow code
             # package so no environment variables to set.
-            return None
+            return {}
         handling_cls = cls._get_mfcontent_class(mfcontent_info)
-        return handling_cls.get_post_extract_env_vars_impl(dest_dir)
+        v = handling_cls.get_post_extract_env_vars_impl(dest_dir)
+        v["METAFLOW_EXTRACTED_ROOT:"] = dest_dir
+        return v
 
     @classmethod
     def get_archive_info(
@@ -216,15 +216,15 @@ class MetaflowCodeContent:
         )
 
     @classmethod
-    def get_archive_content_names(
+    def get_archive_content_members(
         cls,
         archive: Any,
         content_types: Optional[int] = None,
         packaging_backend: Type[PackagingBackend] = TarPackagingBackend,
-    ) -> List[str]:
+    ) -> List[Any]:
         mfcontent_info = cls._extract_archive_mfcontent_info(archive, packaging_backend)
         handling_cls = cls._get_mfcontent_class(mfcontent_info)
-        return handling_cls.get_archive_content_names_impl(
+        return handling_cls.get_archive_content_members_impl(
             mfcontent_info, archive, content_types, packaging_backend
         )
 
@@ -276,7 +276,9 @@ class MetaflowCodeContent:
                 "Invalid package -- unknown version %s in info: %s"
                 % (version_id, cls._mappings)
             )
-        return cls._mappings[version_id].get_post_extract_env_vars_impl(dest_dir)
+        v = cls._mappings[version_id].get_post_extract_env_vars_impl(dest_dir)
+        v["METAFLOW_EXTRACTED_ROOT:"] = dest_dir
+        return v
 
     # Implement the _impl methods in the base subclass (in this file). These need to
     # happen with as few imports as possible to prevent circular dependencies.
@@ -337,14 +339,14 @@ class MetaflowCodeContent:
         raise NotImplementedError("get_archive_filename_impl not implemented")
 
     @classmethod
-    def get_archive_content_names_impl(
+    def get_archive_content_members_impl(
         cls,
         mfcontent_info: Optional[Dict[str, Any]],
         archive: Any,
         content_types: Optional[int] = None,
         packaging_backend: Type[PackagingBackend] = TarPackagingBackend,
-    ) -> List[str]:
-        raise NotImplementedError("get_archive_content_names_impl not implemented")
+    ) -> List[Any]:
+        raise NotImplementedError("get_archive_content_members_impl not implemented")
 
     @classmethod
     def get_post_extract_env_vars_impl(cls, dest_dir: str) -> Dict[str, str]:
@@ -523,19 +525,22 @@ class MetaflowCodeContent:
         return mfcontent_info
 
     @classmethod
-    def _extract_mfcontent_info(cls) -> Optional[Dict[str, Any]]:
-        if "_local" in cls._cached_mfcontent_info:
-            return cls._cached_mfcontent_info["_local"]
+    def _extract_mfcontent_info(
+        cls, target_dir: Optional[str] = None
+    ) -> Optional[Dict[str, Any]]:
+        target_dir = target_dir or "_local"
+        if target_dir in cls._cached_mfcontent_info:
+            return cls._cached_mfcontent_info[target_dir]
 
         mfcontent_info = None  # type: Optional[Dict[str, Any]]
-        if os.path.exists(os.path.join(get_metaflow_root(), MFCONTENT_MARKER)):
-            with open(
-                os.path.join(get_metaflow_root(), MFCONTENT_MARKER),
-                "r",
-                encoding="utf-8",
-            ) as f:
+        if target_dir == "_local":
+            root = os.environ.get("METAFLOW_EXTRACTED_ROOT", get_metaflow_root())
+        else:
+            root = target_dir
+        if os.path.exists(os.path.join(root, MFCONTENT_MARKER)):
+            with open(os.path.join(root, MFCONTENT_MARKER), "r", encoding="utf-8") as f:
                 mfcontent_info = json.load(f)
-        cls._cached_mfcontent_info["_local"] = mfcontent_info
+        cls._cached_mfcontent_info[target_dir] = mfcontent_info
         return mfcontent_info
 
     def get_package_version(self) -> int:
@@ -627,13 +632,13 @@ class MetaflowCodeContentV0(MetaflowCodeContent, version_id=0):
         return None
 
     @classmethod
-    def get_archive_content_names_impl(
+    def get_archive_content_members_impl(
         cls,
         mfcontent_info: Optional[Dict[str, Any]],
         archive: Any,
         content_types: Optional[int] = None,
         packaging_backend: Type[PackagingBackend] = TarPackagingBackend,
-    ) -> List[str]:
+    ) -> List[Any]:
         """
         For V0, we use a static list of known files to classify the content
         """
@@ -649,16 +654,20 @@ class MetaflowCodeContentV0(MetaflowCodeContent, version_id=0):
             "condav2-1.cnd": ContentType.OTHER_CONTENT.value,
         }
         to_return = []
-        for filename in packaging_backend.cls_list_members(archive):
+        for member in packaging_backend.cls_list_members(archive):
+            filename = packaging_backend.cls_member_name(member)
+            added = False
             for prefix, classification in known_prefixes.items():
                 if (
                     prefix[-1] == "/" and filename.startswith(prefix)
                 ) or prefix == filename:
                     if content_types & classification:
-                        to_return.append(filename)
-                elif content_types & ContentType.USER_CONTENT.value:
-                    # Everything else is user content
-                    to_return.append(filename)
+                        to_return.append(member)
+                        added = True
+                        break
+            if not added and content_types & ContentType.USER_CONTENT.value:
+                # Everything else is user content
+                to_return.append(member)
         return to_return
 
     @classmethod
@@ -705,7 +714,7 @@ class MetaflowCodeContentV1Base(MetaflowCodeContent, version_id=1):
         cls, mfcontent_info: Optional[Dict[str, Any]], filename: str, in_archive: bool
     ) -> str:
         if in_archive:
-            return filename
+            return os.path.join(cls._other_dir, filename)
         return os.path.join(get_metaflow_root(), "..", cls._other_dir, filename)
 
     @classmethod
@@ -713,7 +722,7 @@ class MetaflowCodeContentV1Base(MetaflowCodeContent, version_id=1):
         cls, mfcontent_info: Optional[Dict[str, Any]], filename: str, in_archive: bool
     ) -> str:
         if in_archive:
-            return filename
+            return os.path.join(cls._code_dir, filename)
         return os.path.join(get_metaflow_root(), filename)
 
     @classmethod
@@ -832,37 +841,38 @@ class MetaflowCodeContentV1Base(MetaflowCodeContent, version_id=1):
         return None
 
     @classmethod
-    def get_archive_content_names_impl(
+    def get_archive_content_members_impl(
         cls,
         mfcontent_info: Optional[Dict[str, Any]],
         archive: Any,
         content_types: Optional[int] = None,
         packaging_backend: Type[PackagingBackend] = TarPackagingBackend,
-    ) -> List[str]:
+    ) -> List[Any]:
         to_return = []
         module_content = set(mfcontent_info.get("module_files", []))
-        for filename in packaging_backend.cls_list_members(archive):
+        for member in packaging_backend.cls_list_members(archive):
+            filename = packaging_backend.cls_member_name(member)
             if filename.startswith(cls._other_dir) and (
                 content_types & ContentType.OTHER_CONTENT.value
             ):
-                to_return.append(filename)
+                to_return.append(member)
             elif filename.startswith(cls._code_dir):
                 # Special case for marker which is a other content even if in code.
-                if filename == f"{cls._code_dir}/{MFCONTENT_MARKER}":
+                if filename == MFCONTENT_MARKER:
                     if content_types & ContentType.OTHER_CONTENT.value:
-                        to_return.append(filename)
+                        to_return.append(member)
                     else:
                         continue
                 # Here it is either module or code
                 if os.path.join(cls._code_dir, filename) in module_content:
                     if content_types & ContentType.MODULE_CONTENT.value:
-                        to_return.append(filename)
+                        to_return.append(member)
                 elif content_types & ContentType.CODE_CONTENT.value:
-                    to_return.append(filename)
+                    to_return.append(member)
             else:
                 if content_types & ContentType.USER_CONTENT.value:
                     # Everything else is user content
-                    to_return.append(filename)
+                    to_return.append(member)
         return to_return
 
     @classmethod