Apache Beam+Dataflow中Lambda调用函数触发NameError问题求助

问题分析 #

1. NameError原因：Lambda表达式中引用的convert_proto_to_dict函数在Dataflow分布式worker节点上无法被序列化找到——Beam执行时会将代码序列化传输到worker，Lambda内部的外部函数引用如果没有被正确包含在序列化范围内，就会触发未定义错误。
2. 直接传递ProtoBuf类的问题：直接将protobuf_schema_pb2.Message作为参数传给beam.Map时，ProtoBuf类对象无法被正确序列化传输到worker，导致兼容性问题。

解决方案 #

方案1：通过类路径字符串动态导入ProtoBuf类（基于beam.Map） #

修改转换函数，使其接受ProtoBuf类的路径字符串，在函数内部动态导入类，避免序列化类对象：

def convert_proto_to_dict(data, schema_class_path):
    import importlib
    from google.protobuf.json_format import MessageToDict

    # 解析类路径，拆分模块名和类名
    module_name, class_name = schema_class_path.rsplit('.', 1)
    module = importlib.import_module(module_name)
    schema_class = getattr(module, class_name)
    
    message = schema_class()
    message.ParseFromString(data)   
    return MessageToDict(message, preserving_proto_field_name=True) 

Pipeline调用代码修改为：

with beam.Pipeline(options=pipeline_options) as p:
    data = (
        p
        | beam.io.ReadFromPubSub(subscription='projects/abc/subscriptions/abc-sub')
        | beam.Map(convert_proto_to_dict, 'protobuf_schema_pb2.Message')
    )

方案2：自定义可序列化DoFn（更灵活） #

自定义DoFn，在worker节点初始化时动态导入ProtoBuf类，彻底规避序列化问题：

import apache_beam as beam
from google.protobuf.json_format import MessageToDict
import importlib

class ConvertProtoToDict(beam.DoFn):
    def __init__(self, schema_class_path):
        self.schema_class_path = schema_class_path
        self.schema_class = None

    def setup(self):
        # 在worker节点上执行初始化，动态导入ProtoBuf类
        module_name, class_name = self.schema_class_path.rsplit('.', 1)
        module = importlib.import_module(module_name)
        self.schema_class = getattr(module, class_name)

    def process(self, element):
        message = self.schema_class()
        message.ParseFromString(element)
        yield MessageToDict(message, preserving_proto_field_name=True)

Pipeline调用代码：

with beam.Pipeline(options=pipeline_options) as p:
    data = (
        p
        | beam.io.ReadFromPubSub(subscription='projects/abc/subscriptions/abc-sub')
        | beam.ParDo(ConvertProtoToDict('protobuf_schema_pb2.Message'))
    )

额外注意事项 #

• 确保worker环境安装了对应版本的protobuf库，可通过requirements.txt或setup.py指定依赖。
• 如果protobuf_schema_pb2是本地生成的模块，需将其包含在Dataflow作业的打包文件中，通过--setup_file参数指定安装配置。

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