Add @package_sources for sibling source directories

metaflow/plugins/__init__.py

@@ -180,6 +180,7 @@
 ]
 
 TL_PLUGINS_DESC = [
+    ("package_sources", ".package_sources.package_sources"),
     ("yaml_parser", ".parsers.yaml_parser"),
     ("requirements_txt_parser", ".pypi.parsers.requirements_txt_parser"),
     ("namespaced_event_name", ".namespaced_events.namespaced_event_name"),

metaflow/plugins/package_sources.py

@@ -0,0 +1,173 @@
+import inspect
+import os
+from typing import Iterable, Tuple
+
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import DEFAULT_PACKAGE_SUFFIXES
+from metaflow.packaging_sys import ContentType
+from metaflow.packaging_sys.utils import suffix_filter, walk
+from metaflow.user_decorators.user_flow_decorator import FlowMutator
+
+
+class package_sources(FlowMutator):
+    """Include additional files or directories in a flow's code package.
+
+    Relative source paths are resolved from the directory containing the flow
+    file, not from the current working directory. By default, each source is
+    placed in the code package under its basename.
+
+    Parameters
+    ----------
+    sources : path-like, (path-like, path-like), or iterable of these
+        A source file or directory, a ``(source, arcname)`` pair, or multiple
+        source specifications. Directories are traversed recursively.
+
+        A source may be absolute or relative to the flow file. The optional
+        ``arcname`` in a pair specifies where that source is placed inside the
+        code package.
+
+        Use a list to specify exactly two sources without archive names;
+        a two-item tuple is interpreted as ``(source, arcname)``.
+    arcname : path-like, optional
+        Destination for a single source inside the code package. It must be a
+        safe relative path and cannot be absolute, ``.``, or contain ``..``.
+        For multiple sources, specify archive paths with ``(source, arcname)``
+        pairs instead.
+    suffixes : iterable of str or comma-separated str, optional
+        File suffixes to include. Leading dots are optional and matching is
+        case-insensitive. The default is ``DEFAULT_PACKAGE_SUFFIXES``
+        (``.py,.R,.RDS`` by default).
+
+        Providing this argument replaces the default suffix set; it does not
+        extend it.
+
+    Raises
+    ------
+    MetaflowException
+        If a source does not exist, an archive path is unsafe, or ``arcname``
+        is used with multiple sources.
+
+    Examples
+    --------
+    Given this project layout::
+
+        project/
+        ├── flows/
+        │   └── train.py
+        └── src/
+            └── forecasting/
+                └── __init__.py
+
+    Package ``forecasting`` at the archive root so it remains importable as
+    ``import forecasting`` during remote execution::
+
+        from metaflow import FlowSpec, package_sources
+
+        @package_sources("../src/forecasting")
+        class TrainFlow(FlowSpec):
+            ...
+
+    Package multiple sources, assigning a custom archive location to one of
+    them and including JSON files::
+
+        @package_sources(
+            [
+                "../shared",
+                ("../generated/client", "vendor/client"),
+            ],
+            suffixes=[".py", ".json"],
+        )
+        class TrainFlow(FlowSpec):
+            ...
+    """
+
+    def init(self, sources, arcname=None, suffixes=None):
+        specs = self._source_specs(sources)
+        if arcname is not None:
+            if len(specs) != 1:
+                raise MetaflowException(
+                    "arcname can only be used with a single package source"
+                )
+            specs = ((specs[0][0], self._validate_arcname(arcname)),)
+
+        self._sources = specs
+        self._file_filter = suffix_filter(self._suffixes(suffixes))
+
+    @classmethod
+    def _source_specs(cls, sources):
+        if isinstance(sources, (str, os.PathLike)) or cls._is_path_arcname_pair(
+            sources
+        ):
+            sources = (sources,)
+
+        specs = []
+        for source in sources:
+            if cls._is_path_arcname_pair(source):
+                specs.append((os.fspath(source[0]), cls._validate_arcname(source[1])))
+            else:
+                specs.append((os.fspath(source), None))
+        return tuple(specs)
+
+    @staticmethod
+    def _is_path_arcname_pair(value):
+        return (
+            isinstance(value, tuple)
+            and len(value) == 2
+            and isinstance(value[0], (str, os.PathLike))
+        )
+
+    @staticmethod
+    def _validate_arcname(arcname):
+        arcname = os.fspath(arcname)
+        normalized = os.path.normpath(arcname)
+        if (
+            normalized == "."
+            or os.path.isabs(normalized)
+            or ".." in arcname.replace("\\", "/").split("/")
+        ):
+            raise MetaflowException(
+                "package_sources arcname must be a relative path inside the code package"
+            )
+        return normalized
+
+    @staticmethod
+    def _suffixes(suffixes):
+        if suffixes is None:
+            suffixes = DEFAULT_PACKAGE_SUFFIXES.split(",")
+        elif isinstance(suffixes, str):
+            suffixes = suffixes.split(",")
+        return tuple(
+            suffix if suffix.startswith(".") else "." + suffix
+            for suffix in (suffix.strip() for suffix in suffixes)
+            if suffix
+        )
+
+    def add_to_package(self) -> Iterable[Tuple[str, str, ContentType]]:
+        flow_file = inspect.getfile(self._flow_cls)
+        flow_dir = os.path.dirname(os.path.abspath(flow_file))
+
+        for source, arcname in self._sources:
+            source_path = source
+            if not os.path.isabs(source_path):
+                source_path = os.path.join(flow_dir, source_path)
+            source_path = os.path.realpath(source_path)
+
+            if not os.path.exists(source_path):
+                raise MetaflowException(
+                    "package_sources source does not exist: %s" % source
+                )
+
+            root_arcname = arcname or os.path.basename(os.path.normpath(source_path))
+            if os.path.isfile(source_path):
+                if self._file_filter(os.path.basename(source_path)):
+                    yield (source_path, root_arcname, ContentType.USER_CONTENT)
+                continue
+
+            for file_path, rel_arcname in walk(
+                source_path + os.sep, file_filter=self._file_filter
+            ):
+                yield (
+                    file_path,
+                    os.path.join(root_arcname, rel_arcname),
+                    ContentType.USER_CONTENT,
+                )

test/unit/test_package_sources.py

@@ -0,0 +1,119 @@
+import importlib
+import os
+import shutil
+import sys
+from unittest import mock
+
+import pytest
+
+from metaflow import package_sources
+from metaflow.exception import MetaflowException
+from metaflow.packaging_sys import ContentType
+from metaflow.plugins.package_sources import package_sources as package_sources_plugin
+
+
+def _write(path, content=""):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(content)
+
+
+def _mutator(sources, flow_file, **kwargs):
+    mutator = package_sources.__new__(package_sources)
+    mutator._flow_cls = mock.Mock()
+    mutator.init(sources, **kwargs)
+    return mutator, mock.patch(
+        "metaflow.plugins.package_sources.inspect.getfile",
+        return_value=os.fspath(flow_file),
+    )
+
+
+def test_package_sources_is_available_at_top_level():
+    assert package_sources is package_sources_plugin
+
+
+def test_packages_sibling_source_relative_to_flow_file(tmp_path):
+    flow_file = tmp_path / "flows" / "flow.py"
+    source = tmp_path / "shared"
+    _write(flow_file)
+    _write(source / "__init__.py")
+    _write(source / "helper.py")
+    _write(source / "data.json")
+
+    mutator, getfile = _mutator("../shared", flow_file)
+    with getfile:
+        results = list(mutator.add_to_package())
+
+    assert {result[1] for result in results} == {
+        os.path.join("shared", "__init__.py"),
+        os.path.join("shared", "helper.py"),
+    }
+    assert all(result[2] == ContentType.USER_CONTENT for result in results)
+
+
+def test_src_layout_package_is_importable_after_packaging(tmp_path, monkeypatch):
+    flow_file = tmp_path / "flows" / "flow.py"
+    source = tmp_path / "src" / "forecasting"
+    package_root = tmp_path / "package"
+    _write(flow_file)
+    _write(source / "__init__.py", "VALUE = 'packaged'\n")
+
+    mutator, getfile = _mutator("../src/forecasting", flow_file)
+    with getfile:
+        results = list(mutator.add_to_package())
+
+    for file_path, archive_path, _ in results:
+        destination = package_root / archive_path
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copyfile(file_path, destination)
+
+    monkeypatch.syspath_prepend(os.fspath(package_root))
+    sys.modules.pop("forecasting", None)
+    try:
+        module = importlib.import_module("forecasting")
+        assert module.VALUE == "packaged"
+    finally:
+        sys.modules.pop("forecasting", None)
+
+
+def test_supports_multiple_sources_arcnames_and_suffixes(tmp_path):
+    flow_file = tmp_path / "flows" / "flow.py"
+    first = tmp_path / "first"
+    second = tmp_path / "second"
+    _write(flow_file)
+    _write(first / "helper.py")
+    _write(first / "config.json")
+    _write(second / "model.py")
+
+    mutator, getfile = _mutator(
+        [("../first", "vendor/first"), "../second"],
+        flow_file,
+        suffixes=[".py", ".json"],
+    )
+    with getfile:
+        results = list(mutator.add_to_package())
+
+    assert {result[1] for result in results} == {
+        os.path.join("vendor", "first", "helper.py"),
+        os.path.join("vendor", "first", "config.json"),
+        os.path.join("second", "model.py"),
+    }
+
+
+@pytest.mark.parametrize("arcname", [".", "../outside", "nested/../outside", "/tmp"])
+def test_rejects_unsafe_arcnames(arcname):
+    with pytest.raises(MetaflowException, match="relative path"):
+        package_sources._validate_arcname(arcname)
+
+
+def test_rejects_arcname_with_multiple_sources(tmp_path):
+    with pytest.raises(MetaflowException, match="single package source"):
+        _mutator(["../first", "../second"], tmp_path / "flow.py", arcname="pkg")
+
+
+def test_missing_source_raises(tmp_path):
+    flow_file = tmp_path / "flows" / "flow.py"
+    _write(flow_file)
+    mutator, getfile = _mutator("../missing", flow_file)
+
+    with getfile, pytest.raises(MetaflowException, match="source does not exist"):
+        list(mutator.add_to_package())