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()

Asked By: code_monkey

||

Source

Answer 1

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.

Answered By: lua_python_java

Answer 2

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).

Answered By: Code-Apprentice

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

Question:

Answers: