如何在Databricks结合MLFlow实现超参网格搜索并行训练？

核心思路是利用glom()把每个Spark分区的超参组合打包成列表，再映射一个批量训练评估函数，让每个分区对应一个并行任务，任务内批量处理该分区的所有超参组合，正好匹配你要的3个任务各训练3个模型的需求。

1. 定义批量训练评估函数 #

这个函数接收一组超参组合，遍历每个组合完成训练、评估，并用MLFlow记录所有结果。适配Databricks的MLFlow集成：

import mlflow
import mlflow.xgboost
import xgboost as xgb
from sklearn.metrics import mean_squared_error

def train_batch(params_list):
    # 绑定到指定MLFlow实验（Databricks环境自动关联工作区，无需额外配置URI）
    mlflow.set_experiment("/Shared/XGBoost_Hyperparam_Search")
    
    batch_results = []
    for params in params_list:
        # 解析超参（示例对应x=迭代轮数，y=采样比例）
        num_round, subsample = params
        
        # 替换为你的训练/测试数据：建议用广播变量传递大数据集，避免重复加载
        dtrain = xgb.DMatrix(data=train_data.drop("label"), label=train_data["label"])
        dtest = xgb.DMatrix(data=test_data.drop("label"), label=test_data["label"])
        
        # 组装XGBoost参数
        xgb_config = {
            "objective": "reg:squarederror",
            "subsample": subsample,
            "eval_metric": "rmse"
        }
        
        # 启动MLFlow Run记录
        with mlflow.start_run():
            model = xgb.train(xgb_config, dtrain, num_boost_round=num_round)
            preds = model.predict(dtest)
            rmse = mean_squared_error(test_data["label"], preds, squared=False)
            
            # 记录超参、指标和模型
            mlflow.log_params({"num_round": num_round, "subsample": subsample})
            mlflow.log_metric("rmse", rmse)
            mlflow.xgboost.log_model(model, "xgb_model")
            
            batch_results.append({"params": params, "rmse": rmse})
    return batch_results

2. 并行执行超参搜索 #

用你已有的超参组合生成逻辑，并行化时指定分区数，通过glom()打包分区内的超参，再映射批量函数：

# 生成超参组合（示例为2个参数，实际替换为你的5个参数的所有组合）
paras_combo = [(x, y) for x in [50, 100, 150] for y in [0.8, 0.9, 0.95]]

# 并行执行：3个分区对应3个并行任务，每个任务处理3组超参
parallel_output = (
    sc.parallelize(paras_combo, 3)
    .glom()  # 将每个分区的超参组合打包成列表
    .map(train_batch)  # 每个分区映射批量训练函数
    .collect()  # 收集所有任务的结果
)

# 展开结果（每个任务返回一个结果列表，合并成一维列表）
all_results = [res for batch_res in parallel_output for res in batch_res]

3. 实用注意事项 #

• 大数据集优化：如果训练数据量很大，用sc.broadcast(train_data)广播数据，避免每个任务重复加载，大幅提升效率。
• 容错处理：在train_batch的循环内加try-except块，避免单个超参组合训练失败导致整个任务崩溃。
• 资源配置：根据模型训练的资源开销，调整Spark executor的数量和内存/CPU配额，避免资源瓶颈。

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