DataBricks训练XGBoost模型时Spark Driver崩溃求助

1. 先修正代码语法错误 #

你的代码里max_depth=5enter code here存在明显语法错误，多余的反引号和enter code here`会导致代码执行失败，直接引发Driver异常。修正后应为：

xgb = SparkXGBClassifier(
    features_col="features",
    label_col=target_column,
    prediction_col="prediction",
    num_round=50,
    max_depth=5,
    missing=0.0,
    num_workers=2,
    kill_spark_context_on_worker_failure=False
)

2. 拆分Pipeline验证特征工程阶段 #

直接运行完整Pipeline很难定位问题，先单独验证特征工程输出是否异常：

# 仅运行特征工程步骤
feature_pipeline = Pipeline(stages=indexers + encoders + [imputer] + [assembler])
transformed_df = feature_pipeline.fit(train_df).transform(train_df)

# 检查特征向量的基本情况
transformed_df.select("features").describe().show()
# 查看单条数据的特征维度
transformed_df.select(size("features")).show()

如果这一步就出现Driver崩溃，说明问题出在特征工程：

• 检查是否有高基数类别特征：比如某个类别列有几万甚至几十万不同取值，OneHotEncoder后生成极多特征，导致特征向量体积过大，耗尽Driver内存。
• 检查数值特征的量级：如果某列数值范围极大（比如1e9级别），可能导致特征向量在内存中占用过多空间。

3. 调整Spark Driver内存配置 #

虽然集群总内存为56GB，但默认Driver内存分配可能不足。在Databricks集群配置页面：

• 如果是单节点集群（Driver与Worker合一），将Driver Memory调整为40GB左右（预留部分内存给系统进程）；
• 如果是多节点集群，确保Driver Memory至少分配16GB以上，避免内存不足引发OOM导致Driver重启。

4. 优化XGBoost参数缓解计算压力 #

当前参数可能导致Driver或Worker计算负载过高，尝试调整以下参数：

xgb = SparkXGBClassifier(
    features_col="features",
    label_col=target_column,
    prediction_col="prediction",
    num_round=20,  # 先减少迭代轮次测试
    max_depth=4,   # 降低树深度减少计算量
    missing=0.0,
    num_workers=2,
    kill_spark_context_on_worker_failure=False,
    colsample_bytree=0.8,  # 每棵树随机采样80%的特征
    subsample=0.8,         # 每棵树随机采样80%的数据
    cache_size="1024"      # 限制内存缓存大小
)

5. 查看详细Driver日志定位根因 #

你提供的日志仅显示Driver重启流程，需要查看更详细的错误日志：

• 进入Databricks集群页面，点击"Driver logs"；
• 查看stdout和stderr日志，查找OutOfMemoryError、XGBoost相关报错等关键信息，这是定位问题最直接的方式。

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