Databricks中dill序列化scikit-learn pipeline报错求助

问题 #

在Databricks平台上使用scikit-learn（版本0.24.2）构建机器学习pipeline，训练完成后执行以下代码序列化保存拟合后的pipeline对象至DBFS路径：

import dill
f = open("/dbfs/HR_pipe.p", mode='wb')
dill.dump(fitted,f)

其中fitted = main_pipeline.fit(X_train, y_train)，执行时触发SparkContext引用异常，报错回溯如下：

---------------------------------------------------------------------------
Exception Traceback (most recent call last)
<command-97102145541341> in <module>
      1 import dill
      2 f = open("/dbfs/HR_pipe.p", mode='wb')
----> 3 dill.dump(fitted,f)
...
Exception: It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transformation. SparkContext can only be used on the driver, not in code that it run on workers. For more information, see SPARK-5063.

解决方案 #

原因分析 #

你的scikit-learn pipeline中可能包含了依赖SparkContext的组件（比如自定义Transformer调用了Spark API），或训练过程中意外将SparkContext引用嵌入到了pipeline对象中，导致dill序列化时触发Spark的分布式执行限制。

具体解决方法 #

• 清理pipeline中的Spark依赖
  检查pipeline内所有Transformer/Estimator，确保拟合后的组件没有保留SparkContext、Spark DataFrame或RDD的引用。如果有自定义组件，修改代码移除这些引用，只保留本地数据结构（如numpy数组、pandas DataFrame）。

• 先本地序列化再复制到DBFS
  避免直接写入DBFS时的分布式上下文干扰，先将pipeline序列化到本地临时文件，再复制到DBFS：

  import dill
  import shutil
  import os
  
  # 序列化到本地临时路径
  local_path = "/tmp/HR_pipe.p"
  with open(local_path, 'wb') as f:
      dill.dump(fitted, f)
  
  # 复制到DBFS
  shutil.copy(local_path, "/dbfs/HR_pipe.p")
  # 可选：删除本地临时文件
  os.remove(local_path)
  

• 改用scikit-learn官方推荐的joblib序列化
  joblib对scikit-learn组件的兼容性更好，能避免部分dill的序列化问题：

  import joblib
  import shutil
  
  # 本地序列化
  local_path = "/tmp/HR_pipe.joblib"
  joblib.dump(fitted, local_path)
  
  # 复制到DBFS
  shutil.copy(local_path, "/dbfs/HR_pipe.joblib")
  

• 验证并清理pipeline对象
  打印pipeline各组件的属性，检查是否存在Spark相关对象并手动清理：

  for step in fitted.steps:
      print(f"Step {step[0]}: {dir(step[1])}")
  

内容的提问来源于stack exchange，提问作者antonio mario Arrizza