Databricks中multiprocessing Pool.map内存溢出及性能下降问题求助

核心问题点 #

1. 全局序列化器的跨进程复用风险
   代码中avro_serializer是主进程初始化的全局对象，子进程通过fork继承后，可能存在非进程安全问题，导致每个子进程复用该对象时产生内存泄漏或资源占用无法释放。Avro序列化器通常并非为多进程共享设计，跨进程复用易引发内部状态累积。

2. 全量数据拉取至驱动节点
   df.collect()将整个DataFrame的所有记录加载到驱动节点内存，复杂结构的1万条数据本身会占用大量内存。后续多进程处理时，每个子进程还会复制数据副本，进一步加剧内存压力。

3. 冗余的JSON序列化/反序列化操作
   json.loads(avro_serializer.to_json(row_dict))先将Avro对象转为JSON字符串，再解析为字典，这一步完全多余——直接返回avro_serializer.to_json(row_dict)的字符串结果即可，重复操作会额外消耗CPU和内存。

4. Databricks环境与multiprocessing的适配冲突
   Databricks驱动节点为容器化管理，multiprocessing的fork模式可能与容器内存隔离机制冲突，导致进程退出后内存无法被容器正确回收，出现每次运行后交换内存持续上升的情况。

具体优化方案 #

1. 子进程内独立初始化序列化器 #

避免全局序列化器跨进程共享，改为在子进程任务中初始化，或通过进程池初始化器传入Schema，确保每个进程拥有独立实例：

def init_worker(schema_dict):
    global avro_serializer
    avro_schema = avro.schema.SchemaFromJSONData(schema_dict, avro.schema.Names())
    avro_serializer = AvroJsonSerializer(avro_schema)

def create_json_avro_encoding(row):
    row_dict = row.asDict(True)
    # 移除冗余的json.loads
    return avro_serializer.to_json(row_dict)

# 初始化进程池时传入Schema
with Pool(pool_cnt, initializer=init_worker, initargs=(avro_schema_dict,), maxtasksperchild=1) as pool:
    json_data_ret = pool.map(create_json_avro_encoding, records)

2. 改用Spark分布式UDF替代multiprocessing #

放弃驱动节点的多进程处理，利用Spark分布式计算能力将编码逻辑放到Executor节点执行，彻底规避驱动内存瓶颈：

from pyspark.sql.types import StringType

@F.udf(returnType=StringType())
def avro_encode_udf(row_json):
    row_dict = json.loads(row_json)
    avro_schema = avro.schema.SchemaFromJSONData(avro_schema_dict, avro.schema.Names())
    avro_serializer = AvroJsonSerializer(avro_schema)
    return avro_serializer.to_json(row_dict)

# 分布式执行编码逻辑
result_df = df.withColumn("avro_json", avro_encode_udf(F.to_json(F.struct("*"))))
# 按需收集结果，避免全量加载
json_data_ret = result_df.select("avro_json").rdd.map(lambda x: x[0]).collect()

3. 强制内存回收与进程池调优 #

• 处理完成后显式清理变量并触发垃圾回收：

  import gc
  
  del records
  del json_data_ret
  gc.collect()
  

• 进一步降低进程池数量至CPU核心数的1/4（如8核用2进程），避免驱动节点资源过载。

4. 移除冗余的JSON解析步骤 #

直接返回avro_serializer.to_json(row_dict)的字符串结果，无需通过json.loads转为字典，减少内存占用与CPU消耗。

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