Add 'include' parameter to merge_artifacts for easier control of artifacts to merge (#287)

Co-authored-by: Romain Cledat <[email protected]>

Author: romain-intel

metaflow/exception.py

@@ -105,9 +105,16 @@ class CommandException(MetaflowException):
 class MetaflowDataMissing(MetaflowException):
     headline = "Data missing"
 
-class MergeArtifactsException(MetaflowException):
+class UnhandledInMergeArtifactsException(MetaflowException):
     headline = "Unhandled artifacts in merge"
 
     def __init__(self, msg, unhandled):
-        super(MergeArtifactsException, self).__init__(msg)
+        super(UnhandledInMergeArtifactsException, self).__init__(msg)
+        self.artifact_names = unhandled
+
+class MissingInMergeArtifactsException(MetaflowException):
+    headline = "Missing artifacts in merge"
+
+    def __init__(self, msg, unhandled):
+        super(MissingInMergeArtifactsException, self).__init__(msg)
         self.artifact_names = unhandled

metaflow/flowspec.py

@@ -6,7 +6,8 @@
 
 from . import cmd_with_io
 from .parameters import Parameter
-from .exception import MetaflowException, MetaflowInternalError, MergeArtifactsException
+from .exception import MetaflowException, MetaflowInternalError, \
+    MissingInMergeArtifactsException, UnhandledInMergeArtifactsException
 from .graph import FlowGraph
 
 # For Python 3 compatibility
@@ -253,7 +254,7 @@ def _find_input(self, stack_index=None):
                                                                   frame.index + 1))
             return self._cached_input[stack_index]
 
-    def merge_artifacts(self, inputs, exclude=[]):
+    def merge_artifacts(self, inputs, exclude=[], include=[]):
         """
         Merge the artifacts coming from each merge branch (from inputs)
 
@@ -288,40 +289,68 @@ def merge_artifacts(self, inputs, exclude=[]):
         inputs : List[Steps]
             Incoming steps to the join point
         exclude : List[str], optional
+            If specified, do not consider merging artifacts with a name in `exclude`.
+            Cannot specify if `include` is also specified
+        include : List[str], optional
+            If specified, only merge artifacts specified. Cannot specify if `exclude` is
+            also specified
 
         Raises
         ------
         MetaflowException
             This exception is thrown if this is not called in a join step
-        MergeArtifactsException
+        UnhandledInMergeArtifactsException
             This exception is thrown in case of unresolved conflicts
+        MissingInMergeArtifactsException
+            This exception is thrown in case an artifact specified in `include cannot
+            be found
         """
         node = self._graph[self._current_step]
         if node.type != 'join':
             msg = "merge_artifacts can only be called in a join and step *{step}* "\
                   "is not a join".format(step=self._current_step)
             raise MetaflowException(msg)
+        if len(exclude) > 0 and len(include) > 0:
+            msg = "`exclude` and `include` are mutually exclusive in merge_artifacts"
+            raise MetaflowException(msg)
 
         to_merge = {}
         unresolved = []
         for inp in inputs:
             # available_vars is the list of variables from inp that should be considered
-            available_vars = ((var, sha) for var, sha in inp._datastore.items()
-                              if (var not in exclude) and (not hasattr(self, var)))
+            if include:
+                available_vars = ((var, sha) for var, sha in inp._datastore.items()
+                                  if (var in include) and (not hasattr(self, var)))
+            else:
+                available_vars = ((var, sha) for var, sha in inp._datastore.items()
+                                if (var not in exclude) and (not hasattr(self, var)))
             for var, sha in available_vars:
                 _, previous_sha = to_merge.setdefault(var, (inp, sha))
                 if previous_sha != sha:
                     # We have a conflict here
                     unresolved.append(var)
-
+        # Check if everything in include is present in to_merge
+        missing = []
+        for v in include:
+            if v not in to_merge and not hasattr(self, v):
+                missing.append(v)
         if unresolved:
             # We have unresolved conflicts so we do not set anything and error out
             msg = "Step *{step}* cannot merge the following artifacts due to them "\
                   "having conflicting values:\n[{artifacts}].\nTo remedy this issue, "\
                   "be sure to explictly set those artifacts (using "\
                   "self.<artifact_name> = ...) prior to calling merge_artifacts."\
                   .format(step=self._current_step, artifacts=', '.join(unresolved))
-            raise MergeArtifactsException(msg, unresolved)
+            raise UnhandledInMergeArtifactsException(msg, unresolved)
+        if missing:
+            msg = "Step *{step}* specifies that [{include}] should be merged but "\
+                  "[{missing}] are not present.\nTo remedy this issue, make sure "\
+                  "that the values specified in only come from at least one branch"\
+                  .format(
+                      step=self._current_step,
+                      include=', '.join(include),
+                      missing=', '.join(missing))
+            raise MissingInMergeArtifactsException(msg, missing)
         # If things are resolved, we go and fetch from the datastore and set here
         for var, (inp, _) in to_merge.items():
             setattr(self, var, getattr(inp, var))

test/core/tests/merge_artifacts.py

@@ -22,10 +22,20 @@ def modify_things(self):
     @steps(0, ['join'], required=True)
     def merge_things(self, inputs):
         from metaflow.current import current
-        from metaflow.exception import MergeArtifactsException
+        from metaflow.exception import UnhandledInMergeArtifactsException, MetaflowException
 
         # Test to make sure non-merged values are reported
-        assert_exception(lambda: self.merge_artifacts(inputs), MergeArtifactsException)
+        assert_exception(lambda: self.merge_artifacts(inputs), UnhandledInMergeArtifactsException)
+
+        # Test to make sure nothing is set if failed merge_artifacts
+        assert(not hasattr(self, 'non_modified_passdown'))
+        assert(not hasattr(self, 'manual_merge_required'))
+
+        # Test to make sure that only one of exclude/include is used
+        assert_exception(lambda: self.merge_artifacts(
+            inputs,
+            exclude=['ignore_me'],
+            include=['non_modified_passdown']), MetaflowException)
 
         # Test to make sure nothing is set if failed merge_artifacts
         assert(not hasattr(self, 'non_modified_passdown'))

test/core/tests/merge_artifacts_include.py

@@ -0,0 +1,52 @@
+from metaflow_test import MetaflowTest, ExpectationFailed, steps
+
+class MergeArtifactsIncludeTest(MetaflowTest):
+    PRIORITY = 1
+
+    @steps(0, ['start'])
+    def start(self):
+        self.non_modified_passdown = 'a'
+        self.modified_to_same_value = 'b'
+        self.manual_merge_required = 'c'
+        self.ignore_me = 'd'
+
+    @steps(2, ['linear'])
+    def modify_things(self):
+        # Set to different things
+        from metaflow.current import current
+        self.manual_merge_required = current.task_id
+        self.ignore_me = current.task_id
+        self.modified_to_same_value = 'e'
+        assert_equals(self.non_modified_passdown, 'a')
+
+    @steps(0, ['join'], required=True)
+    def merge_things(self, inputs):
+        from metaflow.current import current
+        from metaflow.exception import MissingInMergeArtifactsException
+
+        self.manual_merge_required = current.task_id
+        # Test to see if we raise an exception if include specifies non-merged things
+        assert_exception(lambda: self.merge_artifacts(
+            inputs, include=['manual_merge_required', 'foobar']), MissingInMergeArtifactsException)
+
+        # Test to make sure nothing is set if failed merge_artifacts
+        assert(not hasattr(self, 'non_modified_passdown'))
+
+        # Merge include non_modified_passdown
+        self.merge_artifacts(inputs, include=['non_modified_passdown'])
+
+        # Ensure that everything we expect is passed down
+        assert_equals(self.non_modified_passdown, 'a')
+        assert_equals(self.manual_merge_required, current.task_id)
+        assert(not hasattr(self, 'ignore_me'))
+        assert(not hasattr(self, 'modified_to_same_value'))
+
+    @steps(0, ['end'])
+    def end(self):
+        # Check that all values made it through
+        assert_equals(self.non_modified_passdown, 'a')
+        assert(hasattr(self, 'manual_merge_required'))
+
+    @steps(3, ['all'])
+    def step_all(self):
+        assert_equals(self.non_modified_passdown, 'a')