Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for instance/class/static methods #13944

Merged
merged 7 commits into from
Jun 13, 2024
Merged

Add support for instance/class/static methods #13944

merged 7 commits into from
Jun 13, 2024

Conversation

jlowin
Copy link
Member

@jlowin jlowin commented Jun 11, 2024

This PR adds support for decorating instance/class/static methods as @flows or @tasks.

Today, methods aren't supported because of how Prefect handles the implicit self argument; you'll get an error like ParameterBindError: Error binding parameters for function 'instance_method': missing a required argument: 'self'. Function 'instance_method' has signature 'self' but received args: () and kwargs: []..

This is because the function signature (used for validating parameters) demands an explicit self call arg that will never be supplied by the user.

The solution (for both flows and tasks) has two parts:

  1. Implement the descriptor protocol on Flows and Tasks by adding a __get__ method. When an instance method is accessed on its parent instance, the __get__ method is called with the instance as its first argument. It is expected to return the bound method. We use that arg to create and return a copy of the Flow object which has that instance bound to its own fn as a __prefect_self__ attribute.
  2. Modify Prefect's get_call_parameters function (which transforms user-supplied arguments into a full set of callargs) to look for __prefect_self__ and, if found, prepend it to the front of the callargs.

This solution means:

  1. Users never have to supply an implicit bound instance, just like normal Python
  2. Prefect receives the implicit instance as a "normal" parameter, meaning that all input-based systems, such as caching, respect its existance.

@jlowin jlowin requested a review from a team as a code owner June 11, 2024 21:33
@zzstoatzz
Copy link
Contributor

fyi: I was curious about BaseModel here, so I tried, and it gets mad because the resulting method is not annotated with a type

In [1]: from pydantic import BaseModel

In [2]: from prefect import flow

In [3]: class Foo(BaseModel):
   ...:     x: int
   ...:     @flow
   ...:     def f(self):
   ...:         print(self.x)
   
PydanticUserError: A non-annotated attribute was detected: `f = <prefect.flows.Flow object at 0x106e84950>`. All model fields require a type annotation; if `f` is not meant to be a field, you may be able to resolve this error by annotating it as a `ClassVar` or updating `model_config['ignored_types']`.

here's how it can work

def test_flow_supports_instance_methods_with_basemodel(self):
    class Foo(pydantic.BaseModel):
        model_config = pydantic.ConfigDict(ignored_types=(Flow,))
        x: int = 5

        @flow
        def instance_method(self):
            return self.x

    assert Foo().instance_method() == 5
    assert isinstance(Foo().instance_method, Flow)

Copy link
Member

@cicdw cicdw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clever - can you add some tests to confirm that self is properly provided to task input cache policies? (or really any hook that engages with task inputs, could also be result_storage_key)

@jlowin
Copy link
Member Author

jlowin commented Jun 12, 2024

@cicdw's request demonstrated that the original approach did not satisfy Prefect's internal parameter machinery since it was too implicit.

I've updated the PR to reflect a new, far simpler, approach. Instead of returning a bound instance from the descriptor method, we store the bound instance on __prefect_self__ in a way that let's Prefect's parameter functions detect and restore it.

Copy link
Member

@cicdw cicdw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

@jlowin jlowin merged commit 65a7a23 into main Jun 13, 2024
26 checks passed
@jlowin jlowin deleted the methods branch June 13, 2024 20:12
@abrookins
Copy link
Collaborator

Very cool!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

4 participants