fix potential TaskFailedListener's callback won't be called (#6612)

there is possibility that onJobStart of TaskFailedListener won't be called, if
the job is submitted before the other thread adds addSparkListener.

detail can be found at https://github.com/dmlc/xgboost/pull/6019#issuecomment-760937628
This commit is contained in:
Bobby Wang 2021-01-21 14:20:32 +08:00 committed by GitHub
parent f8bb678c67
commit 9d2832a3a3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -613,8 +613,12 @@ object XGBoost extends Serializable {
}
}
sparkJobThread.setUncaughtExceptionHandler(tracker)
val trackerReturnVal = parallelismTracker.execute {
sparkJobThread.start()
val trackerReturnVal = parallelismTracker.execute(tracker.waitFor(0L))
tracker.waitFor(0L)
}
logger.info(s"Rabit returns with exit code $trackerReturnVal")
val (booster, metrics) = postTrackerReturnProcessing(trackerReturnVal,
boostersAndMetrics, sparkJobThread)