ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/lint.py

@@ -52,7 +52,7 @@ def check_reserved_words(graph):
     msg = "Step name *%s* is a reserved word. Choose another name for the " "step."
     for node in graph:
         if node.name in RESERVED:
-            raise LintWarn(msg % node.name)
+            raise LintWarn(msg % node.name, node.func_lineno, node.source_file)
 
 
 @linter.ensure_fundamentals
@@ -76,9 +76,9 @@ def check_that_end_is_end(graph):
     node = graph["end"]
 
     if node.has_tail_next or node.invalid_tail_next:
-        raise LintWarn(msg0, node.tail_next_lineno)
+        raise LintWarn(msg0, node.tail_next_lineno, node.source_file)
     if node.num_args > 1:
-        raise LintWarn(msg1, node.tail_next_lineno)
+        raise LintWarn(msg1, node.tail_next_lineno, node.source_file)
 
 
 @linter.ensure_fundamentals
@@ -90,7 +90,7 @@ def check_step_names(graph):
     )
     for node in graph:
         if re.search("[^a-z0-9_]", node.name) or node.name[0] == "_":
-            raise LintWarn(msg.format(node), node.func_lineno)
+            raise LintWarn(msg.format(node), node.func_lineno, node.source_file)
 
 
 @linter.ensure_fundamentals
@@ -108,11 +108,11 @@ def check_num_args(graph):
     msg2 = "Step *{0.name}* is missing the 'self' argument."
     for node in graph:
         if node.num_args > 2:
-            raise LintWarn(msg0.format(node), node.func_lineno)
+            raise LintWarn(msg0.format(node), node.func_lineno, node.source_file)
         elif node.num_args == 2 and node.type != "join":
-            raise LintWarn(msg1.format(node), node.func_lineno)
+            raise LintWarn(msg1.format(node), node.func_lineno, node.source_file)
         elif node.num_args == 0:
-            raise LintWarn(msg2.format(node), node.func_lineno)
+            raise LintWarn(msg2.format(node), node.func_lineno, node.source_file)
 
 
 @linter.ensure_static_graph
@@ -125,7 +125,7 @@ def check_static_transitions(graph):
     )
     for node in graph:
         if node.type != "end" and not node.has_tail_next:
-            raise LintWarn(msg.format(node), node.func_lineno)
+            raise LintWarn(msg.format(node), node.func_lineno, node.source_file)
 
 
 @linter.ensure_static_graph
@@ -134,11 +134,17 @@ 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."
+        "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, numbers or config expressions "
+        "(self.config.key_name), not variables."
     )
     for node in graph:
         if node.type != "end" and node.has_tail_next and node.invalid_tail_next:
-            raise LintWarn(msg.format(node), node.tail_next_lineno)
+            raise LintWarn(msg.format(node), node.tail_next_lineno, node.source_file)
 
 
 @linter.ensure_static_graph
@@ -151,7 +157,11 @@ def check_unknown_transitions(graph):
     for node in graph:
         unknown = [n for n in node.out_funcs if n not in graph]
         if unknown:
-            raise LintWarn(msg.format(node, step=unknown[0]), node.tail_next_lineno)
+            raise LintWarn(
+                msg.format(node, step=unknown[0]),
+                node.tail_next_lineno,
+                node.source_file,
+            )
 
 
 @linter.ensure_acyclicity
@@ -165,9 +175,13 @@ def check_for_acyclicity(graph):
 
     def check_path(node, seen):
         for n in node.out_funcs:
+            if node.type == "split-switch" and n == node.name:
+                continue
             if n in seen:
                 path = "->".join(seen + [n])
-                raise LintWarn(msg.format(path), node.tail_next_lineno)
+                raise LintWarn(
+                    msg.format(path), node.tail_next_lineno, node.source_file
+                )
             else:
                 check_path(graph[n], seen + [n])
 
@@ -195,7 +209,7 @@ def check_for_orphans(graph):
     orphans = nodeset - seen
     if orphans:
         orphan = graph[list(orphans)[0]]
-        raise LintWarn(msg.format(orphan), orphan.func_lineno)
+        raise LintWarn(msg.format(orphan), orphan.func_lineno, orphan.source_file)
 
 
 @linter.ensure_static_graph
@@ -226,26 +240,48 @@ 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:
+                if node.type == "split-switch" and n == node.name:
+                    continue
+                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)
-                raise LintWarn(msg0.format(roots=roots))
+                raise LintWarn(
+                    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]
-                if len(node.in_funcs) != len(split_roots):
-                    paths = ", ".join(node.in_funcs)
+
+                # Resolve each incoming function to its root branch from the split.
+                resolved_branches = set(
+                    graph[n].split_branches[-1] for n in node.in_funcs
+                )
+
+                # compares the set of resolved branches against the expected branches
+                # from the split.
+                if len(resolved_branches) != len(
+                    split_roots
+                ) or resolved_branches ^ set(split_roots):
+                    paths = ", ".join(resolved_branches)
                     roots = ", ".join(split_roots)
                     raise LintWarn(
                         msg1.format(
                             node, paths=paths, num_roots=len(split_roots), roots=roots
                         ),
                         node.func_lineno,
+                        node.source_file,
                     )
             else:
-                raise LintWarn(msg2.format(node), node.func_lineno)
+                raise LintWarn(msg2.format(node), node.func_lineno, node.source_file)
 
             # check that incoming steps come from the same lineage
             # (no cross joins)
@@ -256,14 +292,56 @@ def check_split_join_balance(graph):
                     return tuple(graph[n].split_parents)
 
             if not all_equal(map(parents, node.in_funcs)):
-                raise LintWarn(msg3.format(node), node.func_lineno)
+                raise LintWarn(msg3.format(node), node.func_lineno, node.source_file)
+        else:
+            new_stack = split_stack
 
         for n in node.out_funcs:
+            if node.type == "split-switch" and n == node.name:
+                continue
             traverse(graph[n], new_stack)
 
     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):
@@ -276,7 +354,9 @@ def check_empty_foreaches(graph):
         if node.type == "foreach":
             joins = [n for n in node.out_funcs if graph[n].type == "join"]
             if joins:
-                raise LintWarn(msg.format(node, join=joins[0]))
+                raise LintWarn(
+                    msg.format(node, join=joins[0]), node.func_lineno, node.source_file
+                )
 
 
 @linter.ensure_static_graph
@@ -290,7 +370,22 @@ def check_parallel_step_after_next(graph):
         if node.parallel_foreach and not all(
             graph[out_node].parallel_step for out_node in node.out_funcs
         ):
-            raise LintWarn(msg.format(node))
+            raise LintWarn(msg.format(node), node.func_lineno, node.source_file)
+
+
+@linter.ensure_static_graph
+@linter.check
+def check_join_followed_by_parallel_step(graph):
+    msg = (
+        "An @parallel step should be followed by a join step. Step *{0}* is called "
+        "after an @parallel step but is not a join step. Please add an extra `inputs` "
+        "argument to the step."
+    )
+    for node in graph:
+        if node.parallel_step and not graph[node.out_funcs[0]].type == "join":
+            raise LintWarn(
+                msg.format(node.out_funcs[0]), node.func_lineno, node.source_file
+            )
 
 
 @linter.ensure_static_graph
@@ -305,7 +400,9 @@ def check_parallel_foreach_calls_parallel_step(graph):
             for node2 in graph:
                 if node2.out_funcs and node.name in node2.out_funcs:
                     if not node2.parallel_foreach:
-                        raise LintWarn(msg.format(node, node2))
+                        raise LintWarn(
+                            msg.format(node, node2), node.func_lineno, node.source_file
+                        )
 
 
 @linter.ensure_non_nested_foreach
@@ -318,4 +415,26 @@ def check_nested_foreach(graph):
     for node in graph:
         if node.type == "foreach":
             if any(graph[p].type == "foreach" for p in node.split_parents):
-                raise LintWarn(msg.format(node))
+                raise LintWarn(msg.format(node), node.func_lineno, node.source_file)
+
+
+@linter.ensure_static_graph
+@linter.check
+def check_ambiguous_joins(graph):
+    for node in graph:
+        if node.type == "join":
+            problematic_parents = [
+                p_name
+                for p_name in node.in_funcs
+                if graph[p_name].type == "split-switch"
+            ]
+            if problematic_parents:
+                msg = (
+                    "A conditional path cannot lead directly to a join step.\n"
+                    "In your conditional step(s) {parents}, one or more of the possible paths transition directly to the join step {join_name}.\n"
+                    "As a workaround, please introduce an intermediate, unconditional step on that specific path before joining."
+                ).format(
+                    parents=", ".join("*%s*" % p for p in problematic_parents),
+                    join_name="*%s*" % node.name,
+                )
+                raise LintWarn(msg, node.func_lineno, node.source_file)

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 → metadata_provider}/heartbeat.py

@@ -1,11 +1,12 @@
+import json
 import time
+from threading import Thread
+
 import requests
-import json
 
-from threading import Thread
-from metaflow.sidecar import MessageTypes, Message
-from metaflow.metaflow_config import SERVICE_HEADERS
 from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import SERVICE_HEADERS
+from metaflow.sidecar import Message, MessageTypes
 
 HB_URL_KEY = "hb_url"
 
@@ -51,14 +52,29 @@ class MetadataHeartBeat(object):
                 time.sleep(frequency_secs)
                 retry_counter = 0
             except HeartBeatException as e:
+                print(e)
                 retry_counter = retry_counter + 1
-                time.sleep(4**retry_counter)
+                time.sleep(1.5**retry_counter)
 
     def _heartbeat(self):
         if self.hb_url is not None:
-            response = requests.post(
-                url=self.hb_url, data="{}", headers=self.headers.copy()
-            )
+            try:
+                response = requests.post(
+                    url=self.hb_url, data="{}", headers=self.headers.copy()
+                )
+            except requests.exceptions.ConnectionError as e:
+                raise HeartBeatException(
+                    "HeartBeat request (%s) failed" " (ConnectionError)" % (self.hb_url)
+                )
+            except requests.exceptions.Timeout as e:
+                raise HeartBeatException(
+                    "HeartBeat request (%s) failed" " (Timeout)" % (self.hb_url)
+                )
+            except requests.exceptions.RequestException as e:
+                raise HeartBeatException(
+                    "HeartBeat request (%s) failed"
+                    " (RequestException) %s" % (self.hb_url, str(e))
+                )
             # Unfortunately, response.json() returns a string that we need
             # to cast to json; however when the request encounters an error
             # the return type is a json blob :/

metaflow/{metadata → metadata_provider}/metadata.py

@@ -5,6 +5,7 @@ import time
 from collections import namedtuple
 from itertools import chain
 
+from typing import List
 from metaflow.exception import MetaflowInternalError, MetaflowTaggingError
 from metaflow.tagging_util import validate_tag
 from metaflow.util import get_username, resolve_identity_as_tuple, is_stringish
@@ -76,6 +77,12 @@ class ObjectOrder:
 
 @with_metaclass(MetadataProviderMeta)
 class MetadataProvider(object):
+    TYPE = None
+
+    @classmethod
+    def metadata_str(cls):
+        return "%s@%s" % (cls.TYPE, cls.INFO)
+
     @classmethod
     def compute_info(cls, val):
         """
@@ -623,6 +630,25 @@ class MetadataProvider(object):
             sys_info["r_version"] = env["r_version_code"]
         return sys_info
 
+    def _get_git_info_as_dict(self):
+        git_info = {}
+        # NOTE: For flows executing remotely, we want to read from the INFO file of the code package that contains
+        # information on the original environment that deployed the flow.
+        # Otherwise git related info will be missing, as the repository is not part of the codepackage.
+        from metaflow.packaging_sys import MetaflowCodeContent
+
+        env = MetaflowCodeContent.get_info() or self._environment.get_environment_info()
+        for key in [
+            "repo_url",
+            "branch_name",
+            "commit_sha",
+            "has_uncommitted_changes",
+        ]:
+            if key in env and env[key]:
+                git_info[key] = env[key]
+
+        return git_info
+
     def _get_system_tags(self):
         """Convert system info dictionary into a list of system tags"""
         return [
@@ -653,19 +679,78 @@ 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)],
                 )
             )
+        # Add script name as metadata
+        script_name = self._environment.get_environment_info()["script"]
+        metadata.append(
+            MetaDatum(
+                field="script-name",
+                value=script_name,
+                type="script-name",
+                tags=["attempt_id:{0}".format(attempt)],
+            )
+        )
+        # And add git metadata
+        git_info = self._get_git_info_as_dict()
+        if git_info:
+            metadata.append(
+                MetaDatum(
+                    field="git-info",
+                    value=json.dumps(git_info),
+                    type="git-info",
+                    tags=["attempt_id:{0}".format(attempt)],
+                )
+            )
         if metadata:
             self.register_metadata(run_id, step_name, task_id, metadata)
 
+    @classmethod
+    def filter_tasks_by_metadata(
+        cls,
+        flow_name: str,
+        run_id: str,
+        step_name: str,
+        field_name: str,
+        pattern: str,
+    ) -> List[str]:
+        """
+        Filter tasks by metadata field and pattern, returning task pathspecs that match criteria.
+
+        Parameters
+        ----------
+        flow_name : str
+            Flow name, that the run belongs to.
+        run_id: str
+            Run id, together with flow_id, that identifies the specific Run whose tasks to query
+        step_name: str
+            Step name to query tasks from
+        field_name: str
+            Metadata field name to query
+        pattern: str
+            Pattern to match in metadata field value
+
+        Returns
+        -------
+        List[str]
+            List of task pathspecs that satisfy the query
+        """
+        raise NotImplementedError()
+
     @staticmethod
     def _apply_filter(elts, filters):
         if filters is None: