Databricks使用MLflow注册模型预测后调用.show()报错如何解决

你注册的model1是Spark MLlib类型的MLflow模型，当使用mlflow.pyfunc.spark_udf加载生成预测UDF时：

• UDF会分发到Spark Worker节点执行
• Spark类型模型的pyfunc加载逻辑会默认尝试在Worker节点本地启动独立SparkContext初始化模型，而Databricks的Worker节点不允许启动独立SparkContext，Java网关启动失败就触发该报错。

方案1：直接加载Spark原生模型执行预测（推荐） #

不需要生成UDF，直接调用Spark模型的原生transform方法即可，完全适配Databricks PySpark运行环境，修改后代码如下：

import mlflow.spark

model_details = mlflow.tracking.MlflowClient().get_latest_versions('model1', stages=['staging'])[0]
# 加载Spark原生模型
model = mlflow.spark.load_model(model_uri=model_details.source)
input_df = sdf.drop('progression')
# 调用transform生成预测结果
df = model.transform(input_df)
# 正常展示结果
df.show(truncate=False)

方案2：注册非Spark类型的模型 #

如果必须使用pyfunc UDF做预测，你可以在训练模型并记录到MLflow时，存储非Spark类型的模型（比如sklearn、xgboost等原生Python风味模型），这类模型加载为pyfunc UDF时不会尝试启动SparkContext，就不会触发该报错。

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