python - 使用自定义 ptransform 时使用数据流运行器构建光束管道时出现递归错误

Question

我使用 Beam 在本地构建并设法运行了一个令人满意的管道，我已准备好将作业发送到 DataFlow。

我计划只使用 save_main_session 管道选项来 pickle 我的 session ，但是我在尝试这样做时遇到了递归错误。经过几次试验和错误后，我设法将范围缩小到我使用装饰器定义 ptransform_fn 的方式。

请在下面找到一个最小的可重现示例

# my_script.py

from typing import Set
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions, SetupOptions
from apache_beam.transforms.ptransform import ptransform_fn


@ptransform_fn
def my_function(pcoll):
    return pcoll | beam.Create([1])


if __name__ == "__main__":
    options = PipelineOptions()
    options.view_as(SetupOptions).save_main_session = True
    with beam.Pipeline(options=options) as p:
        p | my_function()

完整的回溯很长，但以RecursionError: maximum recursion depth exceeded while calling a Python object

结尾

(请注意，这是启用此错误的 save_main_session=True)选项，因此我可以使用本地运行程序运行此 python -m my_script 并将遇到递归错误)

由于 ptransform_fn 实际上使 my_function 以“非 pythonic”方式运行(在没有定义参数的情况下调用)，pickler 库似乎有一个这个问题。

所以我最后的问题是:

• 这是预期的行为吗？如果我想使用 save_main_session ，是否应该坚持定义继承自 Beam.PTransform 的类？
• 是否有一种简单的方法(如设置管道选项)能够 pickle 这个脚本并在数据流上运行它？

Answer 1

save_main_session 本质上有点脆弱；对于任何重要的事情，我建议将逻辑放在一个命名模块中，该模块可以导入到您的主脚本(以及您的工作程序)中。