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

[jvm-packages] fix safe execution (#4046)

Browse Source
This commit is contained in:
Nan Zhu 2019-01-05 19:45:37 -08:00 committed by GitHub
parent 6a569b8cd9
commit e290ec9a80
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
jvm-packages/xgboost4j-spark/src/main/scala/org/apache/spark/SparkParallelismTracker.scala
View File

@ -20,7 +20,7 @@ import java.net.URL
import org.apache.commons.logging.LogFactory
import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
import org.apache.spark.scheduler.{SparkListener, SparkListenerExecutorRemoved, SparkListenerTaskEnd}
import org.codehaus.jackson.map.ObjectMapper
import scala.collection.JavaConverters._
import scala.concurrent.ExecutionContext.Implicits.global
@ -98,9 +98,11 @@ class SparkParallelismTracker(
*/
def execute[T](body: => T): T = {
if (timeout <= 0) {
logger.info("starting training without setting timeout for waiting for resources")
body
} else {
try {
logger.info(s"starting training with timeout set as $timeout ms for waiting for resources")
waitForCondition(numAliveCores >= requestedCores, timeout)
} catch {
case _: TimeoutException =>
@ -119,9 +121,14 @@ private class ErrorInXGBoostTraining(msg: String) extends ControlThrowable {
private[spark] class TaskFailedListener extends SparkListener {
override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
taskEnd.reason match {
case reason: TaskFailedReason =>
throw new ErrorInXGBoostTraining(s"ExecutorLost during XGBoost Training: " +
s"${reason.toErrorString}")
case taskEnd: SparkListenerTaskEnd =>
if (taskEnd.reason.isInstanceOf[TaskFailedReason]) {
throw new ErrorInXGBoostTraining(s"TaskFailed during XGBoost Training: " +
s"${taskEnd.reason}")
}
case executorRemoved: SparkListenerExecutorRemoved =>
throw new ErrorInXGBoostTraining(s"Executor lost during XGBoost Training: " +
s"${executorRemoved.reason}")
case _ =>
}
}