Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Change training stage from ResultStage to ShuffleMapStage (#9423)

Browse Source
This commit is contained in:
jinmfeng001 2023-08-03 23:40:04 +08:00 committed by GitHub
parent 1332ff787f
commit 04c99683c3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala
View File

@ -407,7 +407,10 @@ object XGBoost extends Serializable {
}} }}
val (booster, metrics) = boostersAndMetrics.collect()(0) // The repartition step is to make training stage as ShuffleMapStage, so that when one
// of the training task fails the training stage can retry. ResultStage won't retry when
// it fails.
val (booster, metrics) = boostersAndMetrics.repartition(1).collect()(0)
val trackerReturnVal = tracker.waitFor(0L) val trackerReturnVal = tracker.waitFor(0L)
logger.info(s"Rabit returns with exit code $trackerReturnVal") logger.info(s"Rabit returns with exit code $trackerReturnVal")
if (trackerReturnVal != 0) { if (trackerReturnVal != 0) {