Databricks中MLFlow联用Spark记录模型报RuntimeError求解

mlflow.pyfunc.log_model执行时会通过pickle序列化传入的模型实例，序列化过程会遍历模型实例关联的所有引用。当前predict方法直接引用了全局作用域的SparkSession对象，该对象内部持有不可被序列化的SparkContext，Spark本身禁止在序列化、worker执行场景下引用SparkContext，直接触发SPARK-5063错误。
该报错和predict的实际运行逻辑无关，只要模型类的引用链上存在SparkContext，序列化阶段就会失败，这也是注释掉log_model、或删除predict中Spark相关代码后流程可正常运行的原因。

方案1：解耦推理与IO逻辑（推荐） #

模型类仅保留纯推理逻辑，将Spark写表的IO操作移到模型调用外层，从根源上避免模型持有Spark相关引用，符合MLflow模型的设计规范，后续模型部署、复用不会被绑定固定的下游逻辑：

class SomeModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, input):
        # 仅保留推理计算逻辑，直接返回预测结果pandas对象
        pandas_df = pd.DataFrame(...insert predictions here...)
        return pandas_df

# 训练、日志记录、推理、写表流程
with mlflow.start_run(run_name="SomeModel_run"):
    model = SomeModel()
    # 模型无Spark引用，序列化无报错
    mlflow.pyfunc.log_model("somemodel", python_model=model)
    
    # 外层执行推理后再处理写表
    pred_result = model.predict(None, input_data)
    spark_df = spark.createDataFrame(pred_result)
    spark_df.write.saveAsTable('tablename', mode='append')

方案2：predict内部懒加载SparkSession（兼容现有逻辑） #

如果需要保留predict内部写表的逻辑，不要在全局作用域、类属性层面引用Spark对象，改为在predict方法执行时动态获取当前活动的SparkSession。这种写法下，序列化模型时不会扫描到方法内部的局部变量引用，只有实际执行predict时才会建立Spark连接，不会触发序列化报错：

class SomeModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, input):
        # 方法内部动态获取SparkSession，不会被序列化捕获
        from pyspark.sql import SparkSession
        spark = SparkSession.getActiveSession()

        # 执行推理逻辑
        pandas_df = pd.DataFrame(...insert predictions here...)
        spark_df = spark.createDataFrame(pandas_df)
        spark_df.write.saveAsTable('tablename', mode='append')

# 日志记录逻辑无需调整
with mlflow.start_run(run_name="SomeModel_run"):
    model = SomeModel()
    mlflow.pyfunc.log_model("somemodel", python_model=model)

注意不要将spark对象赋值为类属性（如self.spark = SparkSession...），否则还是会被序列化扫描到SparkContext引用触发报错。

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