Allow FlowDecorators, FlowMutators and StepMutators to add_to_package

romain-intel · romain-intel · commit 44dfacaaa59f · 2026-04-07T00:03:13.000-07:00
This provides an easier way to build things that can modify what
gets packaged with the .tar.gz that gets sent remotely to
execute
diff --git a/metaflow/decorators.py b/metaflow/decorators.py
@@ -261,6 +261,41 @@ def get_top_level_options(self):
         """
         return []
 
+    def add_to_package(self):
+        """
+        Called to add custom files needed by this flow decorator. 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 []
+
 
 # compare this to parameters.add_custom_parameters
 def add_decorator_options(cmd):
diff --git a/metaflow/package/__init__.py b/metaflow/package/__init__.py
@@ -615,6 +615,16 @@ def _add_tuple(path_tuple):
             elif file_type == ContentType.OTHER_CONTENT:
                 self._mfcontent.add_other_file(file_path, file_name)
 
+        # flow decorators
+        for decos in self._flow._flow_decorators.values():
+            for deco in decos:
+                for path_tuple in deco.add_to_package():
+                    path_tuple = _check_tuple(path_tuple)
+                    if path_tuple is None:
+                        continue
+                    _add_tuple(path_tuple)
+
+        # step decorators
         for step in self._flow:
             for deco in step.decorators:
                 for path_tuple in deco.add_to_package():
@@ -630,6 +640,27 @@ def _add_tuple(path_tuple):
                 continue
             _add_tuple(path_tuple)
 
+        # flow mutators
+        for mutator in self._flow._flow_mutators:
+            for path_tuple in mutator.add_to_package():
+                path_tuple = _check_tuple(path_tuple)
+                if path_tuple is None:
+                    continue
+                _add_tuple(path_tuple)
+
+        # step mutators (deduplicated across steps)
+        seen_step_mutators = set()
+        for step in self._flow:
+            for mutator in step.config_decorators:
+                if id(mutator) in seen_step_mutators:
+                    continue
+                seen_step_mutators.add(id(mutator))
+                for path_tuple in mutator.add_to_package():
+                    path_tuple = _check_tuple(path_tuple)
+                    if path_tuple is None:
+                        continue
+                    _add_tuple(path_tuple)
+
     def _user_code_tuples(self):
         if R.use_r():
             # the R working directory
diff --git a/metaflow/user_decorators/user_flow_decorator.py b/metaflow/user_decorators/user_flow_decorator.py
@@ -262,3 +262,38 @@ def mutate(
             A representation of this flow
         """
         return None
+
+    def add_to_package(self):
+        """
+        Called to add custom files needed by this flow mutator. 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 []
diff --git a/metaflow/user_decorators/user_step_decorator.py b/metaflow/user_decorators/user_step_decorator.py
@@ -747,3 +747,38 @@ def mutate(
             A representation of this step
         """
         return None
+
+    def add_to_package(self):
+        """
+        Called to add custom files needed by this step mutator. 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 []
diff --git a/test/unit/test_add_to_package.py b/test/unit/test_add_to_package.py
@@ -0,0 +1,206 @@
+"""Tests for _add_addl_files in MetaflowPackage.
+
+Exercises the add_to_package hooks on FlowDecorators, FlowMutators,
+StepDecorators, and StepMutators, as well as deduplication and error handling.
+"""
+
+import os
+import tempfile
+from types import ModuleType
+from unittest.mock import MagicMock, call
+
+import pytest
+
+from metaflow.package import (
+    MetaflowPackage,
+    NonUniqueFileNameToFilePathMappingException,
+)
+from metaflow.packaging_sys import ContentType
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_step(decorators=None, config_decorators=None):
+    step = MagicMock()
+    step.decorators = decorators or []
+    step.config_decorators = config_decorators or []
+    return step
+
+
+def _make_flow(steps, flow_decorators=None, flow_mutators=None):
+    flow = MagicMock()
+    # The flow may be iterated multiple times (step decos + step mutators),
+    # so return a fresh iterator each time.
+    flow.__iter__ = lambda self: iter(steps)
+    flow._flow_decorators = flow_decorators or {}
+    flow._flow_mutators = flow_mutators or []
+    return flow
+
+
+def _make_environment(tuples=None):
+    env = MagicMock()
+    env.add_to_package.return_value = tuples or []
+    return env
+
+
+def _make_mfcontent():
+    return MagicMock()
+
+
+def _make_deco(tuples):
+    """Create a mock decorator-like object with an add_to_package method."""
+    deco = MagicMock()
+    deco.add_to_package.return_value = tuples
+    return deco
+
+
+def _call_add_addl_files(flow, environment, mfcontent):
+    """Call _add_addl_files on a bare MetaflowPackage instance."""
+    pkg = object.__new__(MetaflowPackage)
+    pkg._flow = flow
+    pkg._environment = environment
+    pkg._mfcontent = mfcontent
+    pkg._add_addl_files()
+    return mfcontent
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+def test_flow_decorator_add_to_package():
+    """Flow decorator's add_to_package files are added."""
+    with tempfile.NamedTemporaryFile(suffix=".py") as f:
+        deco = _make_deco([(f.name, "flow_deco_file.py", ContentType.CODE_CONTENT)])
+        flow = _make_flow(
+            steps=[_make_step()],
+            flow_decorators={"my_deco": [deco]},
+        )
+        mc = _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
+        mc.add_code_file.assert_called_once_with(
+            os.path.realpath(f.name), "flow_deco_file.py"
+        )
+
+
+def test_flow_mutator_add_to_package_module():
+    """Flow mutator's add_to_package with MODULE_CONTENT calls add_module."""
+    import json
+
+    mutator = _make_deco([(json, None, ContentType.MODULE_CONTENT)])
+    flow = _make_flow(steps=[_make_step()], flow_mutators=[mutator])
+    mc = _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
+    mc.add_module.assert_called_once_with(json)
+
+
+def test_step_mutator_deduplicated_across_steps():
+    """Same StepMutator instance on two steps: add_to_package called once."""
+    with tempfile.NamedTemporaryFile(suffix=".py") as f:
+        mutator = _make_deco([(f.name, "shared.py", ContentType.CODE_CONTENT)])
+        step1 = _make_step(config_decorators=[mutator])
+        step2 = _make_step(config_decorators=[mutator])
+        flow = _make_flow(steps=[step1, step2])
+        mc = _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
+        assert mutator.add_to_package.call_count == 1
+        mc.add_code_file.assert_called_once()
+
+
+def test_step_mutator_distinct_instances():
+    """Two different StepMutator instances: both called."""
+    with tempfile.NamedTemporaryFile(suffix=".py") as f1, tempfile.NamedTemporaryFile(
+        suffix=".py"
+    ) as f2:
+        m1 = _make_deco([(f1.name, "file1.py", ContentType.CODE_CONTENT)])
+        m2 = _make_deco([(f2.name, "file2.py", ContentType.CODE_CONTENT)])
+        step = _make_step(config_decorators=[m1, m2])
+        flow = _make_flow(steps=[step])
+        mc = _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
+        assert m1.add_to_package.call_count == 1
+        assert m2.add_to_package.call_count == 1
+        assert mc.add_code_file.call_count == 2
+
+
+def test_legacy_two_tuple_defaults_to_code_content():
+    """A 2-tuple (file_path, arcname) is treated as CODE_CONTENT."""
+    with tempfile.NamedTemporaryFile(suffix=".py") as f:
+        deco = _make_deco([(f.name, "legacy.py")])
+        step = _make_step(decorators=[deco])
+        flow = _make_flow(steps=[step])
+        mc = _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
+        mc.add_code_file.assert_called_once_with(os.path.realpath(f.name), "legacy.py")
+
+
+def test_non_unique_filename_raises():
+    """Different file paths for the same arcname raises an exception."""
+    with tempfile.NamedTemporaryFile(suffix=".py") as f1, tempfile.NamedTemporaryFile(
+        suffix=".py"
+    ) as f2:
+        d1 = _make_deco([(f1.name, "same_name.py", ContentType.CODE_CONTENT)])
+        d2 = _make_deco([(f2.name, "same_name.py", ContentType.CODE_CONTENT)])
+        step = _make_step(decorators=[d1, d2])
+        flow = _make_flow(steps=[step])
+        with pytest.raises(NonUniqueFileNameToFilePathMappingException):
+            _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
+
+
+def test_module_content_deduplicated():
+    """Same module returned by two decorators: add_module called once."""
+    import json
+
+    d1 = _make_deco([(json, None, ContentType.MODULE_CONTENT)])
+    d2 = _make_deco([(json, None, ContentType.MODULE_CONTENT)])
+    step = _make_step(decorators=[d1, d2])
+    flow = _make_flow(steps=[step])
+    mc = _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
+    mc.add_module.assert_called_once_with(json)
+
+
+def test_other_content_type():
+    """OTHER_CONTENT files are passed to add_other_file."""
+    with tempfile.NamedTemporaryFile(suffix=".yaml") as f:
+        deco = _make_deco([(f.name, "config.yaml", ContentType.OTHER_CONTENT)])
+        step = _make_step(decorators=[deco])
+        flow = _make_flow(steps=[step])
+        mc = _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
+        mc.add_other_file.assert_called_once_with(
+            os.path.realpath(f.name), "config.yaml"
+        )
+
+
+def test_ordering_flow_decorators_before_step_decorators():
+    """Flow decorators are processed before step decorators.
+
+    We verify the flow decorator's add_to_package is called before the step
+    decorator's by checking call order on a shared mock recorder.
+    """
+    call_order = []
+
+    def make_recording_deco(label, tuples):
+        deco = MagicMock()
+
+        def record():
+            call_order.append(label)
+            return tuples
+
+        deco.add_to_package = record
+        return deco
+
+    with tempfile.NamedTemporaryFile(suffix=".py") as f1, tempfile.NamedTemporaryFile(
+        suffix=".py"
+    ) as f2:
+        flow_deco = make_recording_deco(
+            "flow_deco", [(f1.name, "flow_file.py", ContentType.CODE_CONTENT)]
+        )
+        step_deco = make_recording_deco(
+            "step_deco", [(f2.name, "step_file.py", ContentType.CODE_CONTENT)]
+        )
+        step = _make_step(decorators=[step_deco])
+        flow = _make_flow(
+            steps=[step],
+            flow_decorators={"fd": [flow_deco]},
+        )
+        _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
+        assert call_order == ["flow_deco", "step_deco"]
