Support applying Lightweight Python components in Pipeline SDK #770

GeorgesLorre · 2024-01-10T15:34:46Z

No description provided.

tests/component/test_python_component.py

src/fondant/pipeline/pipeline.py

GeorgesLorre · 2024-01-10T15:39:54Z

src/fondant/pipeline/pipeline.py

+                    "additionalProperties": True,
+                },
+            }
+            component_spec = ComponentSpec(spec_dict)


This could possibly be refactored by creating and operation_spec directly

This is currently one of the main jobs of the ComponentOp. If we would do this here, we would need to do it in different places for the docker components as well. I think it makes more sense to do this in the ComponentOp than in the Dataset class.

GeorgesLorre · 2024-01-10T15:41:20Z

src/fondant/pipeline/pipeline.py

@@ -181,6 +179,22 @@ def __init__(

        self.resources = resources or Resources()

+    @classmethod
+    def from_component_yaml(cls, path, **kwargs):


this will be the entrypoint for reusable and custom components (the ones that the python_component does not cover)

GeorgesLorre · 2024-01-10T15:42:07Z

src/fondant/component/image.py

+
+
+@dataclass
+class Image:


This needs way more functionality to actually make the images with extra dependencies and the script available

Isn't that the job of the runners?

Yes but I guess we'll have shared functionality between runners that could be abstracted here

GeorgesLorre · 2024-01-10T15:43:44Z

src/fondant/component/component.py

+
+        return PythonComponent
+
+    return wrapper


It works but I'm not completely convinced this is the way forward. Essentially we subclass the Component and add an image property with the details provided in the decorator arguments.

The other possibility would be to make the original component an attribute of the PythonComponent. It seems to be a trade-off between deep inheritance and having the PythonComponent act like a regular component.

Since I don't see any specific downsides to the deep inheritance in this case, I would be inclined to follow your subclass proposal here.

GeorgesLorre · 2024-01-10T15:44:46Z

I can recommend this talk on python decorators: https://www.youtube.com/watch?v=MjHpMCIvwsY

GeorgesLorre · 2024-01-10T15:47:22Z

src/fondant/pipeline/pipeline.py

-
+        self.name = name
+        self.image = image
+        self.component_spec = component_spec


I need to check further but maybe we can do without the component_spec here and use the operation_spec along with the image from this point on.

I would propose to unpack the component_spec already, and accept the following arguments here:

name

image

consumes

produces

arguments

That would also make it easier to test in Python code I believe.

Problem then is that we have consumes, produces, and arguments twice, with different meaning. Not sure yet how to solve this in the cleanest way.

See my comment below on a possible solution where we still accept a component_spec here, but create the ComponentSpec based on arguments instead of a dict.

RobbeSneyders

Thanks @GeorgesLorre, starting to look good!

tests/component/test_python_component.py

RobbeSneyders · 2024-01-11T09:28:49Z

tests/component/test_python_component.py

+    pipeline = Pipeline(name="dummy-pipeline", base_path="./data")
+
+    @python_component(
+        base_image="python:3.8-slim-buster",


Here we'll need to provide a base docker image with Fondant installed, right?

yes we need a sane default which we ideally configure in 1 place

RobbeSneyders · 2024-01-11T09:29:44Z

src/fondant/component/image.py

+
+
+@dataclass
+class Image:


Isn't that the job of the runners?

src/fondant/component/component.py

RobbeSneyders · 2024-01-11T09:34:47Z

src/fondant/component/component.py

+
+        return PythonComponent
+
+    return wrapper


The other possibility would be to make the original component an attribute of the PythonComponent. It seems to be a trade-off between deep inheritance and having the PythonComponent act like a regular component.

Since I don't see any specific downsides to the deep inheritance in this case, I would be inclined to follow your subclass proposal here.

src/fondant/pipeline/pipeline.py

RobbeSneyders · 2024-01-11T09:39:00Z

src/fondant/pipeline/pipeline.py

-
+        self.name = name
+        self.image = image
+        self.component_spec = component_spec


I would propose to unpack the component_spec already, and accept the following arguments here:

name

image

consumes

produces

arguments

That would also make it easier to test in Python code I believe.

Problem then is that we have consumes, produces, and arguments twice, with different meaning. Not sure yet how to solve this in the cleanest way.

RobbeSneyders · 2024-01-11T09:41:14Z

src/fondant/pipeline/pipeline.py

+                    "additionalProperties": True,
+                },
+            }
+            component_spec = ComponentSpec(spec_dict)


This is currently one of the main jobs of the ComponentOp. If we would do this here, we would need to do it in different places for the docker components as well. I think it makes more sense to do this in the ComponentOp than in the Dataset class.

RobbeSneyders · 2024-01-11T09:44:05Z

src/fondant/pipeline/pipeline.py

+            spec_dict = {
+                "name": name,
+                "description": "This is an example component",
+                "image": "example_component:latest",
+                "produces": {
+                    "additionalProperties": True,
+                },
+                "consumes": {
+                    "additionalProperties": True,
+                },
+            }


I'm not a fan of creating a dict here.

Possible solution, which would also solve my comment on the ComponentOp init arguments above, is that we refactor the ComponentSpec class to take separate init arguments instead:

class ComponentSpec: def __init__(self, name, image, consumes, produces, arguments): ... @classmethod def from_dict(cls, spec): name, image, consumes, produces, arguments = unpack(spec) return cls(name, image, consumes, produces, arguments)

RobbeSneyders · 2024-01-11T09:44:43Z

src/fondant/pipeline/pipeline.py

-
+        self.name = name
+        self.image = image
+        self.component_spec = component_spec


See my comment below on a possible solution where we still accept a component_spec here, but create the ComponentSpec based on arguments instead of a dict.

RobbeSneyders · 2024-01-11T16:35:18Z

tests/component/test_python_component.py

This is testing the pipeline SDK, so should probably move to tests/pipeline

Actually, this is the case for the python_component and Image as well. I think these should be part of the pipeline SDK.

mrchtr

I like it! Looks really promising and I think adds it is great for the usability and quick experiments.

The decorator give use some additional power that we should use imo. I would try to add useful validation and hints to guarantee a low entry barrier for users.

We can probably tackle this in follow PRs :)

mrchtr · 2024-01-12T09:58:28Z

tests/pipeline/test_python_component.py

+        def __init__(self, **kwargs):
+            pass
+
+        def load(self) -> dd.DataFrame:


Probably to detailed for the first draft but I think we should already consider thinking about such edge cases. I guess it would make sense to implement a validation of the class. Looking at the interface it is possible to use any class that inherits from the BaseComponent as ref. Even if we restrict the interface to Load, Transform, Write (see my comment above) a user might run into issue if he doesn't implement the correct function.

Practical speaking, if I create for instance following class:

class CreateData(...): def some_function(): ...

without the implementation of a load method it would be nice get a validation/exception message when I try to execute the pipeline, that the required load function wasn't implemented.
I think we could add this validation to the decorator implementation.

mrchtr · 2024-01-12T09:59:00Z

src/fondant/pipeline/lightweight_component.py

+class Image:
+    base_image: t.Optional[str] = "fondant:latest"
+    extra_requires: t.Optional[t.List[str]] = None
+    script: t.Optional[str] = None


I guess the script is needed for the sagemaker runner?

For every runner, since the script is not baked into the image. In the docker & KfP runner, we'll probably include the script in the entrypoint instead of uploading it separately though.

mrchtr · 2024-01-12T09:59:25Z

src/fondant/pipeline/lightweight_component.py

+    def wrapper(cls):
+        image = Image(
+            base_image=base_image,
+            extra_requires=extra_requires,


The script parameter is not used here. Is this correct?

I added it on my branch:
feature/python-component...feature/python-in-docker

mrchtr · 2024-01-12T12:48:26Z

src/fondant/pipeline/pipeline.py

@@ -325,7 +339,7 @@ def register_operation(

    def read(
        self,
-        name_or_path: t.Union[str, Path],
+        ref: t.Union[str, Path, t.Type["BaseComponent"]],


I think we can restrict the types here. Read should only allow LoadComponent, apply a XTransformComponent and write a DaskWriteComponent.

mrchtr · 2024-01-12T13:11:54Z

tests/pipeline/test_python_component.py

+                },
+                index=pd.Index(["a", "b", "c"], name="id"),
+            )
+            return dd.from_pandas(df, npartitions=1)


I would try to add a validation for the return type as well. If a user don't returns a dask or pandas dataframe in the custom components, it will not work. It's not related to your code, we have the same requirement already for the custom components. But a here we are lowering the entry barrier and I can imagine that people will not read the documentation in detail anymore.

So I think it might be good if we could catch somehow if it is a wrong return type. I guess it should be possible within the decorator implementation as well.

mrchtr · 2024-01-12T13:43:30Z

The decorator give use some additional power that we should use imo. I would try to add useful validation and hints to guarantee a low entry barrier for users.

To follow up on this - I've played a bit around with the decorator and validations. Maybe it helps.

from typing import TypeVar, Callable

import pandas as pd

T = TypeVar('T')
REQUIRED_METHODS = {"load": [pd.DataFrame]}

def LightWeightComponent(cls):
    class Wrapped(cls):
        def __init__(self, *args, **kwargs):
            for method, types in REQUIRED_METHODS.items():
                attr = self._valid_method_exists(method)
                for _type in types:
                    self._validate_return_type(_type, attr)

        def _valid_method_exists(self, name):
            if not hasattr(self, name):
                raise AttributeError(
                    f"{self.__class__.__name__} Function is missing: {name}")
            else:
                return self.__getattribute__(name)

        def _validate_return_type(self, expected_type: type, func: Callable[..., T]):
            return_type = func.__annotations__["return"]
            if not return_type == expected_type:
                raise AttributeError(
                    f"{return_type} is wrong return type in: {func}. Expected {expected_type}")


    return Wrapped

@LightWeightComponent
class BaseComponent:
    def load(self) -> str:
        pass

@LightWeightComponent
class SubClass(BaseComponent):
    def load(self) -> pd.DataFrame:
        pass

@LightWeightComponent
class SecondSubClass(BaseComponent):
    def load(self) -> str: # Through error during initialisation
        return "0"

if __name__ == '__main__':
    SubClass()
    SecondSubClass()

RobbeSneyders · 2024-01-15T10:12:07Z

+1 on the validation, would really help to give immediate feedback instead of only at runtime. Would keep it out of this PR though, so we can merge it and work in parallel on the remaining work.

RobbeSneyders · 2024-01-15T10:16:29Z

src/fondant/core/component_spec.py

+        return cls(component_spec_dict)
+
+    @classmethod
+    def from_args(


I'm still in favor of making this the __init__ method, and storing the different parts of the spec as different attributes instead of converting to a spec.

mrchtr · 2024-01-15T11:56:40Z

Would keep it out of this PR though, so we can merge it and work in parallel on the remaining work.

Yes indeed. I guess we could even collect some more validations. This were just the first edge cases which came into my mind.

RobbeSneyders

Thanks @GeorgesLorre.

Only some small comments remaining. Looks good overall, looking forward to merge this once the tests are fixed 🚀

src/fondant/core/component_spec.py

src/fondant/pipeline/pipeline.py

GeorgesLorre · 2024-01-16T16:04:12Z

Would keep it out of this PR though, so we can merge it and work in parallel on the remaining work.

Yes indeed. I guess we could even collect some more validations. This were just the first edge cases which came into my mind.

Thx for the great ideas @mrchtr Will take this up in subsequent PR's!

RobbeSneyders

Thanks @GeorgesLorre!

Some minor comments, but let's tackle those in a follow-up PR.

RobbeSneyders · 2024-01-16T19:07:35Z

src/fondant/pipeline/lightweight_component.py

+
+class PythonComponent:
+    @classmethod
+    def image(cls) -> Image:


Any reason for this not to be a (class)property?

RobbeSneyders · 2024-01-16T19:12:53Z

src/fondant/pipeline/pipeline.py

+
+            component_spec = ComponentSpec(
+                name,
+                image.base_image,  # TODO: revisit


What needs to be revisited about this?

RobbeSneyders · 2024-01-16T19:15:03Z

src/fondant/pipeline/pipeline.py

+        if inspect.isclass(ref) and issubclass(ref, PythonComponent):
+            name = ref.__name__
+            image = ref.image()
+            description = ref.__doc__ or "python component"
+
+            component_spec = ComponentSpec(
+                name,
+                image.base_image,  # TODO: revisit
+                description=description,
+                consumes={"additionalProperties": True},
+                produces={"additionalProperties": True},
+            )
+
+            operation = ComponentOp(
+                name,
+                image,
+                component_spec,
+                produces=produces,
+                arguments=arguments,
+                input_partition_rows=input_partition_rows,
+                resources=resources,
+                cache=cache,
+                cluster_type=cluster_type,
+                client_kwargs=client_kwargs,
+            )
+
+        else:
+            operation = ComponentOp.from_component_yaml(
+                ref,
+                produces=produces,
+                arguments=arguments,
+                input_partition_rows=input_partition_rows,
+                resources=resources,
+                cache=cache,
+                cluster_type=cluster_type,
+                client_kwargs=client_kwargs,
+            )


This code is almost identically repeated three times. Would be good if we can abstract this away.

RobbeSneyders · 2024-01-16T19:16:58Z

src/fondant/pipeline/lightweight_component.py

+
+@dataclass
+class Image:
+    base_image: str = "fondant:latest"


We should link this to the installed version of Fondant once we add the CI/CD for this.

RobbeSneyders · 2024-01-16T19:22:55Z

src/fondant/pipeline/pipeline.py

+        if self._is_custom_component(path):
+            component_dir = Path(path)
+        else:
+            component_dir = self._get_registry_path(str(path))


I guess this change can be reverted? Although I would sill prefer it we can find a way to remove the component_dir argument from the __init__ method.

GeorgesLorre commented Jan 10, 2024

View reviewed changes

GeorgesLorre requested review from RobbeSneyders, PhilippeMoussalli and mrchtr January 10, 2024 15:44

GeorgesLorre commented Jan 10, 2024

View reviewed changes

RobbeSneyders reviewed Jan 11, 2024

View reviewed changes

GeorgesLorre force-pushed the feature/python-component branch from 6c23966 to 308e59d Compare January 11, 2024 22:16

mrchtr reviewed Jan 12, 2024

View reviewed changes

RobbeSneyders reviewed Jan 15, 2024

View reviewed changes

src/fondant/core/component_spec.py Outdated Show resolved Hide resolved

src/fondant/pipeline/pipeline.py Outdated Show resolved Hide resolved

RobbeSneyders mentioned this pull request Jan 16, 2024

Add validation to lightweight python component decorator #783

Closed

GeorgesLorre force-pushed the feature/python-component branch from e16e6b1 to e9c53c6 Compare January 16, 2024 16:02

GeorgesLorre marked this pull request as ready for review January 16, 2024 16:03

GeorgesLorre added 5 commits January 16, 2024 17:16

Poc for python components support

17da6e1

Cleanup and move things around

acbd84a

Add support for applying transforms

b13bee6

Fix mypy and tests

6036089

Fix sample test

e137632

GeorgesLorre force-pushed the feature/python-component branch from 118e658 to e137632 Compare January 16, 2024 16:16

GeorgesLorre changed the title ~~Poc for python components support~~ Python components support Jan 16, 2024

RobbeSneyders approved these changes Jan 16, 2024

View reviewed changes

RobbeSneyders linked an issue Jan 16, 2024 that may be closed by this pull request

Support applying Lightweight Python components in Pipeline SDK #750

Closed

RobbeSneyders changed the title ~~Python components support~~ Support applying Lightweight Python components in Pipeline SDK Jan 16, 2024

RobbeSneyders merged commit abcd36f into main Jan 16, 2024
8 of 9 checks passed

Support applying Lightweight Python components in Pipeline SDK #770

Support applying Lightweight Python components in Pipeline SDK #770

Conversation

GeorgesLorre commented Jan 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

GeorgesLorre commented Jan 10, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RobbeSneyders left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

GeorgesLorre Jan 11, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mrchtr left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mrchtr commented Jan 12, 2024

RobbeSneyders commented Jan 15, 2024

Choose a reason for hiding this comment

mrchtr commented Jan 15, 2024

RobbeSneyders left a comment

Choose a reason for hiding this comment

GeorgesLorre commented Jan 16, 2024

RobbeSneyders left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

GeorgesLorre commented Jan 10, 2024 •

edited

Loading

GeorgesLorre Jan 11, 2024 •

edited

Loading