@flow decorator

[davidwaroquiers]

Hello everyone,

Just one idea that came up from presenting jobflow and giving trainings on it to several people, either from academy or from industry. In all cases, people were confused that the make method of a job maker needs to use an @job decorator (actually, it is not mandatory, you can always return a Job object directly but most of the time you'd use the decorator) while the make method of a flow maker does not have any. I think it would be nice to have something where we can put an @flow decorator with some functionality. What I would see as useful is to be able to return the output directly (exactly the same as for jobs) and then the decorator creates itself the Flow, with the list of inner jobs/flows and the output set to the returned value. There might be other things that could be done by this decorator if needed/desired.

Quick example usage:

@job
def add(a,b):
    return a+b

@job
def multiply(a,b):
    return a*b

@dataclass
class AddMultiply(Maker):
    name: str = 'double_multiply'

    @flow
    def make(self, a, b, c):
        j1 = add(a, b)
        j2 = multiply(j1.output, c)
        return {"result": j2.output, "random_info": 3.14}

In this example, the decorator means that the make method returns a Flow object equivalent to:

Flow([j1, j2], output={"result": j2.output, "random_info": 3.14})

I already discussed and exchanged a bit about this with @utf but I wanted to share the idea here so that other people can comment. As I already played a bit to make a proof-of concept implementation, I thought I'd just share it.

The proof-of-concept implementation only works within a single file (no time currently to try to fix it such that it can be importable from jobflow, but I can explain to whoever is interested what is the actual issue). Also if anyone has a better idea on how to implement this, do not hesitate to propose :)

Pinging @arosen93, @computron, @gpetretto as you all might be interested to give your opinions ;-)

POC implementation (in principle can be run directly if jobflow is installed):

import inspect
import textwrap
import ast
from jobflow import Flow, job, Job, run_locally, Maker
from dataclasses import dataclass, field


class MyTransformer(ast.NodeTransformer):
    """Node transformer to be able to return the local variables of a
    function or method together with its output."""
    def my_visit(self, node):
        for i, child in enumerate(node.body):
            if isinstance(child, ast.Return):
                # Create the node to assign a locals_ variable containing the local variables within the function,
                # i.e. "locals_ = locals()". Note that this does not add the node or expression to the ast tree yet.
                new_node = ast.Assign(
                    targets=[ast.Name(id='locals_', ctx=ast.Store())],
                    value=ast.Call(func=ast.Name(id='locals', ctx=ast.Load()), args=[], keywords=[])
                )
                # Insert the "locals_ = locals()" Assign node to the ast tree.
                node.body.insert(i, new_node)
                # Default fix indentation and line numbers
                ast.fix_missing_locations(new_node)
                # Modify the return value so that it returns the local variables and the return value as a 2-tuple
                child.value = ast.Tuple(elts=[ast.Name(id='locals_'), child.value], ctx=ast.Store())
                break
            elif hasattr(child, 'body'):
                self.my_visit(child)

    def visit_FunctionDef(self, node):
        self.my_visit(node)
        for dec in node.decorator_list:
            if isinstance(dec, ast.Name) and dec.id == 'flow':
                node.decorator_list.remove(dec)
            elif isinstance(dec, ast.Call) and dec.func.id == 'flow':
                node.decorator_list.remove(dec)
        return node


def flow(method=None):
    """Decorator for Flows.

    Implementation note:
    We first change dynamically the function/method so that it returns the variables
    local to the function. This is done using the above ast.NodeTransformer.
    The function/method is changed so that it returns a 2-tuple consisting of the local
    variables and the initial return value.
    Example:
    def mymethod(a):
        return a
    will be changed to:
    def mymethod(a):
        locals_ = locals()
        return (locals_, a)
    Then we use the locals in the wrapper and extract the list of jobs and flows
    defined in the function/method.
    """
    # Get the source code of the method
    source = inspect.getsource(method)

    # Needed for methods
    source = textwrap.dedent(source)

    # Use the ast module to parse the source code
    tree = ast.parse(source)

    # Apply the transformer to the AST
    transformer = MyTransformer()
    new_tree = transformer.visit(tree)

    # Generate the modified source code
    modified_source = ast.unparse(new_tree)

    # Execute the code and update the example_method function
    exec(modified_source)
    new_meth = locals()[method.__name__]

    def wrapper(*args, **kwargs):
        method_locals, method_return = new_meth(*args, **kwargs)
        steps = []
        for name, value in method_locals.items():
            if isinstance(value, (Job, Flow)):
                steps.append(value)
        return Flow(steps, output=method_return)
    return wrapper


@dataclass
class MultiplyMaker(Maker):
    name: str = 'multiply'
    multiple: int = 1

    @job
    def make(self, number):
        return number * self.multiple


@dataclass
class AddMaker(Maker):
    name: str = 'add'

    @job
    def make(self, a, b):
        return a + b


@dataclass
class AddMultiply(Maker):
    name: str = 'double_multiply'
    multiply_maker_1: MultiplyMaker = field(default_factory=lambda: MultiplyMaker(name='multiply_1', multiple=1))
    multiple_maker_2: MultiplyMaker = field(default_factory=lambda: MultiplyMaker(name='multiply_2', multiple=1))

    @flow
    def make(self, a, b):
        j1 = self.multiply_maker_1.make(a)
        j2 = self.multiple_maker_2.make(b)
        add_job = AddMaker().make(j1.output, j2.output)
        return add_job.output


add_multiply_maker = AddMultiply()
myflow = add_multiply_maker.make(2, 5)

outputs = run_locally(myflow)
print(outputs)

[Andrew-S-Rosen]

I am strongly in favor of this, in part because this is basically what many other workflow engines do (e.g. take Prefect, as an example). I think it is more intuitive to be able to use a @flow decorator in Jobflow, and the proposed idea here makes a lot of sense to me (although I have admittedly not gone through the implementation in much depth). This may also help with adding other "managers" to jobflow in the future, if that ultimately is of interest. For what it's worth, I have been working on some mini projects related to interoperability between workflow engines, and the manual construction of the Flow object has been one of my major stumbling blocks when it comes to linking Jobflow with other tools.

[Andrew-S-Rosen]

As a follow-up question to @davidwaroquiers: how would you handle dynamic flows with a Response(replace=Flow()) kind of approach? That seems tricky. Is it do-able with something like this?

[davidwaroquiers]

As such I don't think so as Flows are not executed contrary to jobs (and to what is done in prefect I think using async functions). Now maybe it is possible to make some kind of workaround because it would indeed be very useful. In any case, I think this topic is worth discussing together. Also pinging @mkhorton as you might be interested in this discussion.

[davidwaroquiers]

Just to add up on my previous comment, something similar to what is done when a job is replaced by an entire Flow, there is a final job added automatically by jobflow, which contains the output and has the same uuid as the initial job that has been replaced. Something like that may be thought to address this. Anyway, I think it's worth discussing that with everyone interested and maybe a brilliant idea will come out :)

[Andrew-S-Rosen]

Don't know that I'll be able to come up with anything better than you but am certainly happy to continue the brainstorming! I really like the idea. 😄

[utf]

It should be possible to make @flow and @job interact properly such that if they are combined they default to returning a Response(replace=...)

E.g.,

@job
def dynamic_job(maker):
    job1 = maker.make()
    job2 = maker.make(job1.output)
    flow = Flow([job1, job2], output=job2.output)
    return Response(replace=flow)

would become:

@flow
@job
def dynamic_job(maker):
    job1 = maker.make()
    job2 = maker.make(job1.output)
    return job2.output

Alternatively, we could make a new decorator @subflow which does this.

[Andrew-S-Rosen]

Ah, that's a very good point!

[utf]

I like this idea. Although I think this will encourage people to shorten the code from:

@flow
def make(a, b, maker):
    j1 = maker.make(a)
    j2 = maker.make(b)
    add_job = AddMaker().make(j1.output, j2.output)
    return add_job.output

to

@flow
def make(a, b, maker):
    j1 = maker.make(a)
    j2 = maker.make(b)
    return AddMaker().make(j1.output, j2.output).output

I.e., the add_job will not appear in locals(). So we should try and enable this sort of behaviour to be captured correctly.

[davidwaroquiers]

Indeed, we should check all possible use cases (there might be other ones, e.g. maybe list comprehensions should be handled ?). In any case, there might also be an alternative solution/implementation to achieve this (which may ease things ?). Notwithstanding the fact that it is currently not working if it is imported from jobflow (I tried). So it currently works as such here but there is definitely some work to be done. I only shared the above as a short POC to showcase the idea.

[davidwaroquiers]

Ok I went back to this and tried to make it work. I think the current solution I had implemented (outside of jobflow) won't work except with quite some big uncertainties and/or vulnerabilities :/
Looking for other options ... if anyone has ideas, don't hesitate :-)

[Andrew-S-Rosen]

Aw that's a bummer. This would be such an awesome feature, and I keep thinking about how nice it would be. 😅 But I am not skilled enough to implement it myself so don't have any suggestions from my end.

[davidwaroquiers]

Could not agree more :) I gave another try with a different approach here: #547
Waiting for feedback before making (or trying to make) this robust and mergeable :)