GCP Dataflow作业使用自定义PTransform时出现依赖包未找到错误

问题根因 #

这个问题是Apache Beam Python SDK在分布式运行时的序列化机制导致的：

• 当你使用DataflowRunner提交作业时，所有在流水线运行阶段需要用到的自定义类（包括你写的LimitVolume复合PTransform、里面嵌套的DoFn、lambda函数等）都会被序列化，然后分发到GCP的Worker节点上反序列化执行。
• 当你把逻辑直接写在run()方法里时，作业提交阶段就会把这部分逻辑直接展开为Beam原生的转换算子，不需要序列化任何自定义PTransform类到Worker节点，所以不会触发依赖查找的问题。
• 而保留自定义PTransform时，Worker节点反序列化加载这个类的时候，会依赖全局作用域的导入声明，如果你的导入位置不对、或者主会话上下文没有被同步到Worker，就会出现NameError找不到依赖的报错。

可行解决方案 #

你可以任选以下一种方式修复：

• 调整导入语句位置
  把所有用到的依赖（比如import arrow、import time、import apache_beam as beam等）全部放在main.py的最顶部全局作用域，不要放在run()函数内部，也不要放在if __name__ == "__main__"的判断块里，确保Worker加载类的时候能直接找到这些导入的包。
• 提交作业时添加--save_main_session参数
  在提交Dataflow作业的命令中加上--save_main_session参数，这个参数会把你本地主模块的整个会话上下文（包括所有全局导入的依赖、自定义的类）全部序列化后同步到所有Worker节点，反序列化后就能正常找到对应的依赖。
• 在使用依赖的位置内部导入
  如果不想调整全局导入，可以在用到依赖的函数内部做导入，比如你的lambda用到time包的话，可以改成：

lambda message: beam.window.TimestampedValue(message, __import__('time').time())

如果是自定义DoFn的话，在process方法内部导入对应的包即可。

• 打包为标准Python包提交
  如果后续自定义转换比较多，建议把项目打包为标准Python包，编写对应的setup.py，提交作业时用--setup_file=./setup.py参数代替--requirements_file，这样所有自定义类和依赖都会被正确打包安装到Worker节点，从根本上避免序列化相关的依赖问题。

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