Spark ML模型（SparkXGBRegressor等）在Databricks部署推理问题咨询

核心问题解答 #

是否仅Scikit-learn模型支持推理部署？ #

不是，Databricks支持多种Spark ML模型的部署，包括SparkXGBRegressor、RandomForestRegressor这类原生Spark ML模型，并非只有Scikit-learn模型可以部署。

部署Spark模型的其他方式 #

• 通过MLflow部署为Spark UDF：将MLflow记录的Spark模型注册为UDF，可在Spark SQL或DataFrame中直接调用，适合批量推理场景。
• 部署为Databricks Model Serving端点：将MLflow注册的Spark模型部署为实时REST API端点，Databricks会自动管理Spark集群资源，支持低延迟在线推理。
• 导出为标准格式部署：将Spark Pipeline模型导出为PMML或ONNX格式，使用兼容的推理引擎（如ONNX Runtime）部署，适配跨平台场景。

代码问题修复 #

问题1：仅记录最后一次运行模型而非最优模型 #

你的Hyperopt调参逻辑未跟踪最优模型，且hp.choice返回的是参数索引而非实际值，导致最终训练用的参数可能错误。修复后代码可在调参过程中实时更新最优模型，并正确解析参数值。

问题2：MLflow自动日志警告解决 #

警告原因是SparkXGBRegressor不在MLflow自动日志的默认允许列表中，两种解决方式：

• 关闭自动日志，改用代码中已有的mlflow.spark.log_model()显式记录模型（推荐，更可控）。
• 在Spark配置中添加spark.mlflow.pysparkml.autolog.logModelAllowlistFile，指定包含ml.dmlc.xgboost4j.scala.spark.XGBoostRegressor的允许列表文件路径（每行一个模型类名）。

完整修正代码 #

#-------------------------------------------------------XGBoost-------------------------------------------------------------------------
#train_df=train_df.limit(188123)
from pyspark.ml.feature import StringIndexer, VectorAssembler
from pyspark.ml import Pipeline
from xgboost.spark import SparkXGBRegressor
from pyspark.ml.evaluation import RegressionEvaluator
import numpy as np
from mlflow.models.signature import infer_signature
from hyperopt import hp
#vec_assembler = VectorAssembler(inputCols=train_df.columns[1:], outputCol="features")

xgb = SparkXGBRegressor(num_workers=1, label_col="price", missing=0.0)

pipeline = Pipeline(stages=[ordinal_encoder, vec_assembler, xgb])
regression_evaluator = RegressionEvaluator(predictionCol="prediction", labelCol="price")

# 初始化最优模型和RMSE跟踪
best_rmse = float('inf')
best_model = None

def objective_function(params):    
    # 解析Hyperopt返回的索引为实际参数值
    max_depth = np.arange(5, 15, dtype=int)[params["max_depth"]]
    n_estimators = np.arange(50, 80, dtype=int)[params["n_estimators"]]

    with mlflow.start_run(nested=True):
        estimator = pipeline.copy({xgb.max_depth: max_depth, xgb.n_estimators: n_estimators})
        model = estimator.fit(train_df)

        preds = model.transform(test_df)
        rmse = regression_evaluator.evaluate(preds)
        mlflow.log_metric("rmse", rmse)
        mlflow.log_param("max_depth", max_depth)
        mlflow.log_param("n_estimators", n_estimators)

        # 更新全局最优模型
        nonlocal best_rmse, best_model
        if rmse < best_rmse:
            best_rmse = rmse
            best_model = model

    return rmse


search_space = {
    "max_depth" : hp.choice('max_depth', np.arange(5, 15, dtype=int)),
     "n_estimators": hp.choice('n_estimators', np.arange(50, 80, dtype=int))
}

from hyperopt import fmin, tpe, Trials
import numpy as np
import mlflow
import mlflow.spark
import mlflow.sklearn
# 关闭自动日志，避免非允许列表模型的警告
# mlflow.pyspark.ml.autolog(log_models=True)


num_evals = 1
trials = Trials()
best_hyperparam = fmin(fn=objective_function, 
                       space=search_space,
                       algo=tpe.suggest, 
                       max_evals=num_evals,
                       trials=trials,
                       rstate=np.random.default_rng(42))

# 用最优参数在全量数据上重新训练并记录模型
with mlflow.start_run():
    # 解析最优参数的实际值
    best_max_depth = np.arange(5, 15, dtype=int)[best_hyperparam["max_depth"]]
    best_n_estimators = np.arange(50, 80, dtype=int)[best_hyperparam["n_estimators"]]
    estimator = pipeline.copy({xgb.max_depth: best_max_depth, xgb.n_estimators: best_n_estimators})
    combined_df = train_df.union(test_df) # 合并训练和验证数据

    pipeline_model = estimator.fit(combined_df)
    pred_df = pipeline_model.transform(test_df)
    rmse = regression_evaluator.evaluate(pred_df)

    # 记录最终模型的参数和指标
    mlflow.log_param("maxdepth", best_max_depth)
    mlflow.log_param("n_estimators", best_n_estimators)
    mlflow.log_metric("rmse", rmse)
    mlflow.spark.log_model(pipeline_model, "model", input_example=test_df.select(old_cols_list).limit(1).toPandas())

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