如何将utils.py中的sayHello函数作为Metaflow步骤函数使用？

我在utils.py文件中定义了sayHello函数，同时有一个运行简单Metaflow的flow.py文件。我希望从utils.py导入sayHello并将其作为flow.py中的步骤函数使用，是否可行？由于类中需要self参数，这可能存在难度。若可行，如何将上一步的输出传入该函数并传递给下一步？

以下是我的尝试代码：

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

直接将外部函数作为Metaflow步骤不可行，因为Metaflow的步骤方法必须是绑定到FlowSpec类实例的方法，必须接收self参数。但可以通过以下方式实现需求：

1. 类内包装调用外部函数 #

在FlowSpec子类中定义带@step装饰器的方法，在方法内部调用导入的外部函数，同时满足Metaflow对步骤方法的要求。

2. 传递上一步输出并传递给下一步 #

如果需要将上一步的输出传入外部函数，只需在包装步骤中通过self获取上一步的变量，调用外部函数后，将结果赋值给self的属性，即可传递给后续步骤。

修改后的完整代码 #

utils.py（无需修改，可扩展参数）

def sayHello(name="world"):
    print(f"hello {name}")
    return f"hello {name}"

flow.py

from metaflow import FlowSpec, step
from utils import sayHello

class BranchFlow(FlowSpec):
    @step
    def start(self):
        # 定义要传递给后续步骤的参数
        self.greeting_target = "Metaflow"
        self.next(self.a, self.b, self.sayHello_step)

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

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

    # 包装外部函数的步骤方法
    @step
    def sayHello_step(self):
        # 获取上一步传递的变量
        target = self.greeting_target
        # 调用外部函数并接收返回值
        self.greeting_result = sayHello(target)
        print(f"步骤内接收返回值：{self.greeting_result}")
        # 将结果传递给join步骤
        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))
        # 读取sayHello_step分支的结果
        if hasattr(inputs, 'sayHello_step'):
            print('sayHello返回结果：%s' % inputs.sayHello_step.greeting_result)
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == '__main__':
     BranchFlow()    

关键说明 #

• sayHello_step是类内的标准步骤方法，满足Metaflow对self参数的要求，内部调用外部的sayHello函数。
• 通过self.greeting_target传递上一步的输入给sayHello_step，调用外部函数后用self.greeting_result存储结果，后续join步骤可通过inputs.sayHello_step.greeting_result获取该值。
• 运行流程后，join步骤会打印所有分支的计算结果，包括sayHello_step的返回值。

内容的提问来源于stack exchange，提问作者code_monkey