Import a function from one python file and use that function as a step function in metaflow

Question

I have a function sayHello defined in one python file utils.py. I also have a python file flow.py that runs a simple metaflow. I want to import sayHello from utils.py and use sayHello as a step function in flow.py. Is it possible to do that? This might be hard because we need self in the class. If it's possible, one question is that we how do pass output from the previous step into the function and pass the output to next step. The following is my attempt.

#utils.py
def sayHello():
    print("hello world")

#flow.py
from metaflow import FlowSpec, step, Parameter
from metaflow import Metaflow
from utils import sayHello

def function(p):
    return p

class BranchFlow(FlowSpec):
    @step
    def start(self):
        self.next(self.a, self.b, self.sayHello)

    @step
    def a(self):
        self.x = 1
        self.next(self.join)

    @step
    def b(self):
        self.x = 2
        self.next(self.join)


    @step
    def join(self, inputs):
        print('a is %s' % inputs.a.x)
        print('b is %s' % inputs.b.x)
        print('total is %d' % sum(input.x for input in inputs))
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == '__main__':
     BranchFlow()

"Is it possible to do that?" Yes, almost anything you can think of is possible. — Code-Apprentice, Aug 19 '22 at 16:08
"This might be hard because we need self in the class." What do you mean? — Code-Apprentice, Aug 19 '22 at 16:09
"how do pass output from the previous step into the function and pass the output to next step" you probably need to return a value rather than printing it. — Code-Apprentice, Aug 19 '22 at 16:09
I want to use sayHello function as a step function. The step function always comes with self — code_monkey, Aug 19 '22 at 16:10
I think you mean that a step function must have `self` as a parameter. This is because as I say in my answer below that the step function must be a method of the class. — Code-Apprentice, Aug 19 '22 at 20:56

score 2 · Accepted Answer · answered Aug 19 '22 at 16:17

According to the documentation for FlowSpec.next(), the method you pass to self.next() must be a member of the current class and decorated with @step.

Probably the easiest way to use another function from another module is to wrap it in a member function:

class BranchFlow(FlowSpec):
    # ...

    @step
    def sayHello(self):
        sayHello()

Now you can do self.next(self.sayHello).

score 1 · Answer 2 · answered Aug 19 '22 at 16:16

If you need the sayHello method to reference self then add a parameter to the sayHello method if the utils.py. I believe you are importing the method correctly, so that looks good. Just change it slightly:

#utils.py
def sayHello(self, ...other_args_if_needed):
    # Assuming `self` has a `getName` getter method.
    print(f"hello world, I am {self.getName()}")

Now you can import the sayHello method and it can be used in the class object. Just remember you need to have getter and setter methods already defined. self might just be an empty dictionary when the class is initialized, so sayHello would throw an exception if the getName wasn't defined on self.

Import a function from one python file and use that function as a step function in metaflow

2 Answers2