Add support to run lightweight python components in docker runner

examples/sample_pipeline/pipeline.py

@@ -1,11 +1,14 @@
 # This file contains a sample pipeline. Loading data from a parquet file,
 # using the load_from_parquet component, chain a custom dummy component, and use
 # the reusable chunking component
+import os
 from pathlib import Path
 
+import pandas as pd
 import pyarrow as pa
 
-from fondant.pipeline import Pipeline
+from fondant.component import PandasTransformComponent
+from fondant.pipeline import Pipeline, lightweight_component
 
 BASE_PATH = Path("./.artifacts").resolve()
 BASE_PATH.mkdir(parents=True, exist_ok=True)
@@ -32,8 +35,28 @@
     "./components/dummy_component",
 )
 
-dataset.apply(
+dataset = dataset.apply(
     "chunk_text",
     arguments={"chunk_size": 10, "chunk_overlap": 2},
     consumes={"text": "text_data"},
 )
+
+
+@lightweight_component(
+    base_image="python:3.8",
+    extra_requires=[
+        f"fondant[component]@git+https://github.com/ml6team/fondant@"
+        f"{os.environ.get('FONDANT_VERSION', 'main')}",
+    ],
+)
+class CalculateChunkLength(PandasTransformComponent):
+    def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
+        dataframe["chunk_length"] = dataframe["text"].apply(len)
+        return dataframe
+
+
+_ = dataset.apply(
+    ref=CalculateChunkLength,
+    consumes={"text": pa.string()},
+    produces={"chunk_length": pa.int32()},
+)

examples/sample_pipeline/run.sh

@@ -27,6 +27,8 @@ trap cleanup EXIT
 # Bind local data directory to pipeline
 data_dir=$(readlink -f "data")
 
+export FONDANT_VERSION=$GIT_HASH
+
 # Run pipeline
 poetry run fondant run local pipeline.py \
   --extra-volumes $data_dir:/data --build-arg FONDANT_VERSION=$GIT_HASH

src/fondant/pipeline/compiler.py

@@ -3,6 +3,7 @@
 import os
 import shlex
 import tempfile
+import textwrap
 import typing as t
 from abc import ABC, abstractmethod
 from dataclasses import asdict
@@ -16,6 +17,7 @@
 from fondant.pipeline import (
     VALID_ACCELERATOR_TYPES,
     VALID_VERTEX_ACCELERATOR_TYPES,
+    Image,
     Pipeline,
 )
 
@@ -45,6 +47,37 @@ def log_unused_configurations(self, **kwargs):
                     f" for runner `{self.__class__.__name__}`.",
                 )
 
+    @staticmethod
+    def _build_entrypoint(image: Image) -> t.List[str]:
+        """Build the entrypoint to execute the provided image."""
+        if not image.script:
+            # Not a lightweight python component
+            return ["fondant", "execute", "main"]
+
+        command = ""
+        if image.extra_requires:
+            requirements = "\n".join(image.extra_requires)
+            command += textwrap.dedent(
+                f"""\
+                printf {shlex.quote(requirements)} > 'requirements.txt'
+                python3 -m pip install -r requirements.txt
+            """,
+            )
+
+        command += textwrap.dedent(
+            f"""\
+            printf {shlex.quote(image.script)} > 'main.py'
+            fondant execute main "$@"
+        """,
+        )
+
+        return [
+            "sh",
+            "-ec",
+            command,
+            "--",  # All arguments provided after this will be passed to `fondant execute main`
+        ]
+
 
 class DockerCompiler(Compiler):
     """Compiler that creates a docker-compose spec from a pipeline."""
@@ -157,6 +190,8 @@ def _generate_spec(
 
             logger.info(f"Compiling service for {component_name}")
 
+            entrypoint = self._build_entrypoint(component_op.image)
+
             # add metadata argument to command
             command = ["--metadata", metadata.to_json()]
 
@@ -204,6 +239,7 @@ def _generate_spec(
             )
 
             services[component_name] = {
+                "entrypoint": entrypoint,
                 "command": command,
                 "depends_on": depends_on,
                 "volumes": volumes,

src/fondant/pipeline/lightweight_component.py

@@ -1,14 +1,24 @@
+import inspect
+import itertools
+import textwrap
 import typing as t
 from dataclasses import dataclass
 from functools import wraps
 
+from fondant.component import Component
+
 
 @dataclass
 class Image:
     base_image: str = "fondant:latest"
     extra_requires: t.Optional[t.List[str]] = None
     script: t.Optional[str] = None
 
+    def __post_init__(self):
+        if self.base_image is None:
+            # TODO: link to Fondant version
+            self.base_image = "fondant:latest"
+
 
 class PythonComponent:
     @classmethod
@@ -23,12 +33,12 @@ def lightweight_component(
     """Decorator to enable a python component."""
 
     def wrapper(cls):
-        kwargs = {}
-        if base_image:
-            kwargs["base_image"] = base_image
-        if extra_requires:
-            kwargs["extra_requires"] = extra_requires
-        image = Image(**kwargs)
+        script = build_python_script(cls)
+        image = Image(
+            base_image=base_image,
+            extra_requires=extra_requires,
+            script=script,
+        )
 
         # updated=() is needed to prevent an attempt to update the class's __dict__
         @wraps(cls, updated=())
@@ -40,3 +50,47 @@ def image(cls) -> Image:
         return AppliedPythonComponent
 
     return wrapper
+
+
+def build_python_script(component_cls: t.Type[Component]) -> str:
+    """Build a self-contained python script for the provided component class, which will act as
+    the `src/main.py` script to execute the component.
+    """
+    imports_source = textwrap.dedent(
+        """\
+        from typing import *
+        import typing as t
+
+        import dask.dataframe as dd
+        import fondant
+        import pandas as pd
+        from fondant.component import *
+        from fondant.core import *
+    """,
+    )
+
+    component_source = inspect.getsource(component_cls)
+    component_source = textwrap.dedent(component_source)
+    component_source_lines = component_source.split("\n")
+
+    # Removing possible decorators (can be multiline) until the class
+    # definition is found
+    component_source_lines = list(
+        itertools.dropwhile(
+            lambda x: not x.startswith("class"),
+            component_source_lines,
+        ),
+    )
+
+    if not component_source_lines:
+        msg = (
+            f'Failed to dedent and clean up the source of function "{component_cls.__name__}". '
+            f"Its probably not properly indented."
+        )
+        raise ValueError(
+            msg,
+        )
+
+    component_source = "\n".join(component_source_lines)
+
+    return "\n\n".join([imports_source, component_source])

src/fondant/pipeline/pipeline.py

@@ -235,12 +235,11 @@ def _configure_caching_from_image_tag(
 
         return cache
 
-    def dockerfile_path(self, path: t.Union[str, Path]) -> t.Optional[Path]:
-        if self._is_custom_component(path):
-            component_dir = Path(path)
-        else:
-            component_dir = self._get_registry_path(str(path))
-        docker_path = component_dir / "Dockerfile"
+    @property
+    def dockerfile_path(self) -> t.Optional[Path]:
+        if not self.component_dir:
+            return None
+        docker_path = self.component_dir / "Dockerfile"
         return docker_path if docker_path.exists() else None
 
     @staticmethod

tests/pipeline/test_python_component.py

@@ -1,21 +1,62 @@
+import textwrap
+
 import dask.dataframe as dd
 import pandas as pd
 import pyarrow as pa
 from fondant.component import DaskLoadComponent, PandasTransformComponent
 from fondant.pipeline import Pipeline, lightweight_component
 
 
-def test_lightweight_component():
-    pipeline = Pipeline(name="dummy-pipeline", base_path="./data")
+def test_build_python_script():
+    @lightweight_component()
+    class CreateData(DaskLoadComponent):
+        def load(self) -> dd.DataFrame:
+            df = pd.DataFrame(
+                {
+                    "x": [1, 2, 3],
+                    "y": [4, 5, 6],
+                },
+                index=pd.Index(["a", "b", "c"], name="id"),
+            )
+            return dd.from_pandas(df, npartitions=1)
+
+    assert CreateData.image().script == textwrap.dedent(
+        """\
+        from typing import *
+        import typing as t
+
+        import dask.dataframe as dd
+        import fondant
+        import pandas as pd
+        from fondant.component import *
+        from fondant.core import *
+
+
+        class CreateData(DaskLoadComponent):
+            def load(self) -> dd.DataFrame:
+                df = pd.DataFrame(
+                    {
+                        "x": [1, 2, 3],
+                        "y": [4, 5, 6],
+                    },
+                    index=pd.Index(["a", "b", "c"], name="id"),
+                )
+                return dd.from_pandas(df, npartitions=1)
+    """,
+    )
+
+
+def test_lightweight_component(tmp_path_factory):
+    pipeline = Pipeline(
+        name="dummy-pipeline",
+        base_path="./data",
+    )
 
     @lightweight_component(
         base_image="python:3.8-slim-buster",
         extra_requires=["pandas", "dask"],
     )
     class CreateData(DaskLoadComponent):
-        def __init__(self, **kwargs):
-            pass
-
         def load(self) -> dd.DataFrame:
             df = pd.DataFrame(
                 {
@@ -44,5 +85,4 @@ def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
         ref=AddN,
         produces={"x": pa.int32(), "y": pa.int32()},
         consumes={"x": pa.int32(), "y": pa.int32()},
-        arguments={"n": 1},
     )