在Databricks Feature Store自定义MLFlow PyFunc类加载PipelineModel报错

我正在创建一个自定义PyFunc类以用于Databricks Feature Store，因其模型服务UI和feature store的log_model()方法仅支持PythonModel类。底层模型是执行各类分箱与变换操作的PipelineModel()。

编写的代码 #

import mlflow.pyfunc
from pyspark.ml.functions import vector_to_array
from pyspark.ml import PipelineModel

class custom_model_class(mlflow.pyfunc.PythonModel):

  def __init__(self, model_path, threshold):
    self.model_path = model_path
    self.threshold = threshold
    self.model = None

  def load_context(self, context):
    self.model = PipelineModel.load(self.model_path)

  def predict(self, context, model_input):
    return self.model.transform(model_input).withColumn("prediction_opt_thresh", when(vector_to_array("probability")[1] > lit(self.threshold), 1).otherwise(0))

custom_model = custom_model_class(model_path = pipeline_model_directory, threshold = 0.52)

运行错误 #

RuntimeError: SparkContext should only be created and accessed on the driver., from , line 12.

已尝试的解决方法（均未成功） #

• JVM返回None，表明不存在Spark Session
• 创建Spark Session时触发上述相同错误
• 使用mlflow.pyfunc.load_model（不依赖Spark Session）加载PipelineModel()可运行，但无法返回概率值

请问该如何解决？Feature Store与MLLib能否协同工作？

首先明确：Feature Store完全可以和MLLib协同工作，问题出在PyFunc的运行环境限制——模型服务（或log_model后的推理）是在无Spark Driver的环境中执行，而Spark ML的PipelineModel必须依赖Spark Session运行，直接在PyFunc里调用会因为没有可用的Spark Context而报错。以下是具体可行方案：

方案1：使用MLflow原生的Spark ML模型日志（推荐） #

Databricks Feature Store的log_model()方法原生支持Spark ML模型，无需强行包装成PyFunc。直接用以下方式日志模型即可保留所有功能，包括概率值输出：

import mlflow.spark
from databricks.feature_store import FeatureStoreClient

fs = FeatureStoreClient()
fs.log_model(
    model=pipeline_model,
    artifact_path="spark-pipeline-model",
    flavor=mlflow.spark,
    training_set=training_set,  # 你的Feature Store训练集对象
    registered_model_name="your-registered-model-name"
)

如果需要自定义阈值逻辑，可以在日志模型前修改Pipeline中的分类器阈值，或者在推理阶段通过Spark UDF对输出结果二次处理。

方案2：将Spark Pipeline转换为PyFunc兼容格式（无Spark依赖） #

如果必须使用自定义PyFunc，需要把Spark Pipeline的逻辑转换为MLflow可序列化的纯Python实现：

1. 先将Spark Pipeline模型用mlflow.spark.log_model导出为MLflow模型；
2. 在自定义PyFunc中加载该模型，此时MLflow会自动将其转换为兼容无Spark环境的格式；
3. 直接对返回的Pandas DataFrame处理概率值与阈值：

import mlflow.pyfunc
import pandas as pd

class custom_model_class(mlflow.pyfunc.PythonModel):
    def __init__(self, model_path, threshold):
        self.model_path = model_path
        self.threshold = threshold
        self.model = None

    def load_context(self, context):
        # 加载MLflow序列化后的Spark模型（已转为PyFunc兼容格式）
        self.model = mlflow.pyfunc.load_model(self.model_path)

    def predict(self, context, model_input):
        # 模型返回Pandas DataFrame
        result = self.model.predict(model_input)
        # 提取概率值并应用阈值（假设probability列是数组类型）
        result["prediction_opt_thresh"] = result["probability"].apply(
            lambda x: 1 if x[1] > self.threshold else 0
        )
        return result

注意：该方法要求Pipeline中的所有Transformer（如Bucketizer、StringIndexer等）支持MLflow序列化，自定义Transformer需实现MLflow的序列化接口。

方案3：本地Spark Session（仅测试/特定场景） #

如果你的推理环境允许创建本地Spark Session（比如仅在Driver节点运行），可以在load_context中显式创建，但不推荐用于分布式服务：

from pyspark.sql import SparkSession

def load_context(self, context):
    # 创建本地Spark Session
    spark = SparkSession.builder.master("local[*]").getOrCreate()
    self.model = PipelineModel.load(self.model_path)

Databricks模型服务环境通常禁止创建Spark Context，此方法仅适用于测试场景。

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