PySpark pandas_udf执行模型推理时重复运行两次的问题排查

环境与场景 #

使用EMR 5.33.1集群，Python 3.7、PySpark 2.4.7版本，将基于Pandas训练的Keras模型通过PySpark做分布式推理，已用pandas_udf封装推理函数：

@pandas_udf(returnType=DoubleType())
def run_inf(*cols):
    print("run inf initiated")
    X= pd.concat(cols, axis=1)
    model.set_weights(bc_model_weight.value)
    predictions = model.predict(X)
    return pd.Series(predictions.reshape(-1))

执行推理代码如下：

bc_model_weight = sc.broadcast(model.get_weights())
pdf = spark.read.parquet(path_to_file).select(*all_cols)
feature_cols = [f for f in all_cols if f.endswith("feature")]
pdf = pdf.withColumn('pred', run_inf(*feature_cols))
pdf.show()

问题现象 #

执行后发现run inf initiated打印两次，但通过explain()查看物理计划仅显示一次run_inf调用，已确认集群仅使用一个执行器，且UDF内打印的主机名一致，无法定位重复执行原因。

可能原因与解决方法 #

1. Spark Action操作的采样触发 #

show()方法默认会先对数据做采样预览，可能触发一次UDF调用，之后正式执行查询时再调用一次，导致日志重复。

• 解决方式：替换show()为collect()，或通过创建临时表后查询的方式避免采样：

# 用collect()替代show()，手动控制展示条数
results = pdf.collect()
for row in results[:20]:
    print(row)

# 或创建临时表后查询
pdf.createOrReplaceTempView("prediction_results")
spark.sql("SELECT * FROM prediction_results LIMIT 20").show()

2. 模型初始化逻辑重复执行 #

当前UDF每次调用都会执行model.set_weights，且开头的打印语句未做重复执行判断，即使同一执行器内的批次处理，也可能因Spark任务调度触发多次调用。

• 解决方式：在UDF内添加初始化判断，确保模型权重仅设置一次：

model_initialized = False

@pandas_udf(returnType=DoubleType())
def run_inf(*cols):
    global model_initialized
    if not model_initialized:
        print("run inf initiated")
        model.set_weights(bc_model_weight.value)
        model_initialized = True
    X = pd.concat(cols, axis=1)
    predictions = model.predict(X)
    return pd.Series(predictions.reshape(-1))

3. EMR Spark配置的影响 #

EMR默认的Spark配置中，自适应执行（spark.sql.adaptive.enabled）等特性可能导致任务拆分或重复调度，可尝试关闭自适应执行验证：

spark.conf.set("spark.sql.adaptive.enabled", "false")

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