MLflow部署PythonModel触发SparkContext RuntimeError求助

你遇到的错误是因为load_context方法在模型序列化/加载过程中可能被分发到worker节点执行，而SparkContext仅能在driver端使用，无法在worker的代码中引用。当你在load_context里调用get_identifier_information()执行Spark查询时，Spark尝试在worker端调用SparkContext，导致触发报错。

• 提前在driver端预加载数据并转为本地可序列化结构
  在调用mlflow.pyfunc.log_model之前，先在driver节点完成Spark查询，将结果转为pandas DataFrame、字典或本地文件这类可序列化的本地结构，避免在模型的load_context或predict方法中执行Spark操作。

  示例代码：

  # 1. 在driver端提前查询并转换数据
  def get_identifier_information():
      # 仅在driver端执行Spark查询
      spark_df = spark.sql("SELECT * FROM third_party_products")
      # 转为pandas DataFrame（本地可序列化结构）
      return spark_df.toPandas()
  
  # 预加载数据并保存为本地csv文件
  identifier_df = get_identifier_information()
  identifier_df.to_csv("identifier_data.csv", index=False)
  
  # 2. 修改模型类，读取本地预加载的数据
  class EquipmentEmbeddingEndpoint(mlflow.pyfunc.PythonModel):
      def load_context(self, context):
          # 从artifacts读取预保存的本地数据
          self.identifier_data = pd.read_csv(context.artifacts["identifier_data"])
  
      def predict(self, context, model_input):
          # 使用本地的self.identifier_data进行业务处理
          # ...
  
  # 3. 日志模型时传入本地数据作为artifact
  mlflow.pyfunc.log_model(
      artifact_path="equipment_embedding_model",
      python_model=EquipmentEmbeddingEndpoint(),
      artifacts={"identifier_data": "identifier_data.csv"}
  )
  

• 禁止在模型方法中执行Spark操作
  load_context和predict方法可能在非Spark环境（比如MLflow模型服务、本地）或worker节点运行，此时没有可用的SparkContext，绝对不要在这些方法中调用spark.sql、创建SparkSession等操作。

• 若需依赖Spark数据，考虑替代方案
  如果业务逻辑必须依赖Spark分布式数据集，建议使用Spark UDF封装逻辑，而非mlflow.pyfunc模型；或者确保模型仅在driver端运行，但这种场景不适用于大多数模型部署场景。

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