[Breaking][jvm-packages] Use barrier execution mode (#7836)

With the introduction of the barrier execution mode. we don't need to kill SparkContext when some xgboost tasks failed. Instead, Spark will handle the errors for us. So in this PR, `killSparkContextOnWorkerFailure` parameter is deleted.

jvm-packages/xgboost4j-spark/src/test/resources/log4j.properties

@@ -1 +1 @@
-log4j.logger.org.apache.spark=ERROR
+log4j.logger.org.apache.spark=ERROR

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/ExternalCheckpointManagerSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ Copyright (c) 2014-2022 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -19,7 +19,7 @@ package ml.dmlc.xgboost4j.scala.spark
 import java.io.File
 
 import ml.dmlc.xgboost4j.scala.{Booster, DMatrix, ExternalCheckpointManager, XGBoost => SXGBoost}
-import org.scalatest.{FunSuite, Ignore}
+import org.scalatest.FunSuite
 import org.apache.hadoop.fs.{FileSystem, Path}
 
 class ExternalCheckpointManagerSuite extends FunSuite with TmpFolderPerSuite with PerTest {

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/FeatureSizeValidatingSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ Copyright (c) 2014-2022 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -16,10 +16,8 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.XGBoostError
 import org.apache.spark.Partitioner
 import org.apache.spark.ml.feature.VectorAssembler
-import org.apache.spark.sql.SparkSession
 import org.scalatest.FunSuite
 import org.apache.spark.sql.functions._
 
@@ -53,7 +51,7 @@ class FeatureSizeValidatingSuite extends FunSuite with PerTest {
       "objective" -> "binary:logistic",
       "num_round" -> 5, "num_workers" -> 2, "use_external_memory" -> true, "missing" -> 0)
     import DataUtils._
-    val sparkSession = SparkSession.builder().getOrCreate()
+    val sparkSession = ss
     import sparkSession.implicits._
     val repartitioned = sc.parallelize(Synthetic.trainWithDiffFeatureSize, 2)
       .map(lp => (lp.label, lp)).partitionBy(

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/MissingValueHandlingSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ Copyright (c) 2014-2022 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -16,14 +16,14 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.XGBoostError
 import org.apache.spark.ml.feature.VectorAssembler
 import org.apache.spark.ml.linalg.Vectors
 import org.apache.spark.sql.DataFrame
 import org.scalatest.FunSuite
-
 import scala.util.Random
 
+import org.apache.spark.SparkException
+
 class MissingValueHandlingSuite extends FunSuite with PerTest {
   test("dense vectors containing missing value") {
     def buildDenseDataFrame(): DataFrame = {
@@ -113,7 +113,7 @@ class MissingValueHandlingSuite extends FunSuite with PerTest {
     val inputDF = vectorAssembler.transform(testDF).select("features", "label")
     val paramMap = List("eta" -> "1", "max_depth" -> "2",
       "objective" -> "binary:logistic", "missing" -> -1.0f, "num_workers" -> 1).toMap
-    intercept[XGBoostError] {
+    intercept[SparkException] {
       new XGBoostClassifier(paramMap).fit(inputDF)
     }
   }
@@ -140,7 +140,7 @@ class MissingValueHandlingSuite extends FunSuite with PerTest {
     inputDF.show()
     val paramMap = List("eta" -> "1", "max_depth" -> "2",
       "objective" -> "binary:logistic", "missing" -> -1.0f, "num_workers" -> 1).toMap
-    intercept[XGBoostError] {
+    intercept[SparkException] {
       new XGBoostClassifier(paramMap).fit(inputDF)
     }
   }

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/ParameterSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ Copyright (c) 2014-2022 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -16,9 +16,9 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.XGBoostError
-import org.scalatest.{BeforeAndAfterAll, FunSuite, Ignore}
+import org.scalatest.{BeforeAndAfterAll, FunSuite}
 
+import org.apache.spark.SparkException
 import org.apache.spark.ml.param.ParamMap
 
 class ParameterSuite extends FunSuite with PerTest with BeforeAndAfterAll {
@@ -40,28 +40,16 @@ class ParameterSuite extends FunSuite with PerTest with BeforeAndAfterAll {
     assert(xgbCopy2.MLlib2XGBoostParams("eval_metric").toString === "logloss")
   }
 
-  private def waitForSparkContextShutdown(): Unit = {
-    var totalWaitedTime = 0L
-    while (!ss.sparkContext.isStopped && totalWaitedTime <= 120000) {
-      Thread.sleep(10000)
-      totalWaitedTime += 10000
-    }
-    assert(ss.sparkContext.isStopped === true)
-  }
-
   test("fail training elegantly with unsupported objective function") {
     val paramMap = Map("eta" -> "0.1", "max_depth" -> "6", "silent" -> "1",
       "objective" -> "wrong_objective_function", "num_class" -> "6", "num_round" -> 5,
       "num_workers" -> numWorkers)
     val trainingDF = buildDataFrame(MultiClassification.train)
     val xgb = new XGBoostClassifier(paramMap)
-    try {
-      val model = xgb.fit(trainingDF)
-    } catch {
-      case e: Throwable => // swallow anything
-    } finally {
-      waitForSparkContextShutdown()
+    intercept[SparkException] {
+      xgb.fit(trainingDF)
     }
+
   }
 
   test("fail training elegantly with unsupported eval metrics") {
@@ -70,12 +58,8 @@ class ParameterSuite extends FunSuite with PerTest with BeforeAndAfterAll {
       "num_workers" -> numWorkers, "eval_metric" -> "wrong_eval_metrics")
     val trainingDF = buildDataFrame(MultiClassification.train)
     val xgb = new XGBoostClassifier(paramMap)
-    try {
-      val model = xgb.fit(trainingDF)
-    } catch {
-      case e: Throwable => // swallow anything
-    } finally {
-      waitForSparkContextShutdown()
+    intercept[SparkException] {
+      xgb.fit(trainingDF)
     }
   }
 

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/PerTest.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ Copyright (c) 2014-2022 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -19,7 +19,7 @@ package ml.dmlc.xgboost4j.scala.spark
 import java.io.File
 
 import ml.dmlc.xgboost4j.{LabeledPoint => XGBLabeledPoint}
-import org.apache.spark.{SparkConf, SparkContext, TaskFailedListener}
+import org.apache.spark.SparkContext
 import org.apache.spark.sql._
 import org.scalatest.{BeforeAndAfterEach, FunSuite}
 
@@ -40,32 +40,16 @@ trait PerTest extends BeforeAndAfterEach { self: FunSuite =>
       .appName("XGBoostSuite")
       .config("spark.ui.enabled", false)
       .config("spark.driver.memory", "512m")
+      .config("spark.barrier.sync.timeout", 10)
       .config("spark.task.cpus", 1)
 
   override def beforeEach(): Unit = getOrCreateSession
 
   override def afterEach() {
-    TaskFailedListener.sparkContextShutdownLock.synchronized {
-      if (currentSession != null) {
-        // this synchronization is mostly for the tests involving SparkContext shutdown
-        // for unit test involving the sparkContext shutdown there are two different events sequence
-        // 1. SparkContext killer is executed before afterEach, in this case, before SparkContext
-        // is fully stopped, afterEach() will block at the following code block
-        // 2. SparkContext killer is executed afterEach, in this case, currentSession.stop() in will
-        // block to wait for all msgs in ListenerBus get processed. Because currentSession.stop()
-        // has been called, SparkContext killer will not take effect
-        while (TaskFailedListener.killerStarted) {
-          TaskFailedListener.sparkContextShutdownLock.wait()
-        }
-        currentSession.stop()
-        cleanExternalCache(currentSession.sparkContext.appName)
-        currentSession = null
-      }
-      if (TaskFailedListener.sparkContextKiller != null) {
-        TaskFailedListener.sparkContextKiller.interrupt()
-        TaskFailedListener.sparkContextKiller = null
-      }
-      TaskFailedListener.killerStarted = false
+    if (currentSession != null) {
+      currentSession.stop()
+      cleanExternalCache(currentSession.sparkContext.appName)
+      currentSession = null
     }
   }
 

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/PersistenceSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014,2021 by Contributors
+ Copyright (c) 2014-2022 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostConfigureSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ Copyright (c) 2014-2022 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -16,10 +16,8 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.Rabit
 import ml.dmlc.xgboost4j.scala.{Booster, DMatrix}
 
-import scala.collection.JavaConverters._
 import org.apache.spark.sql._
 import org.scalatest.FunSuite
 

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostGeneralSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ Copyright (c) 2014-2022 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -16,13 +16,12 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.XGBoostError
 import scala.util.Random
 
 import ml.dmlc.xgboost4j.{LabeledPoint => XGBLabeledPoint}
 import ml.dmlc.xgboost4j.scala.DMatrix
 
-import org.apache.spark.TaskContext
+import org.apache.spark.{SparkException, TaskContext}
 import org.scalatest.FunSuite
 
 import org.apache.spark.ml.feature.VectorAssembler
@@ -375,13 +374,14 @@ class XGBoostGeneralSuite extends FunSuite with TmpFolderPerSuite with PerTest {
 
   test("throw exception for empty partition in trainingset") {
     val paramMap = Map("eta" -> "0.1", "max_depth" -> "6", "silent" -> "1",
-      "objective" -> "multi:softmax", "num_class" -> "2", "num_round" -> 5,
-      "num_workers" -> numWorkers, "tree_method" -> "auto")
+      "objective" -> "binary:logistic", "num_class" -> "2", "num_round" -> 5,
+      "num_workers" -> numWorkers, "tree_method" -> "auto", "allow_non_zero_for_missing" -> true)
     // The Dmatrix will be empty
-    val trainingDF = buildDataFrame(Seq(XGBLabeledPoint(1.0f, 1, Array(), Array())))
+    val trainingDF = buildDataFrame(Seq(XGBLabeledPoint(1.0f, 4,
+      Array(0, 1, 2, 3), Array(0, 1, 2, 3))))
     val xgb = new XGBoostClassifier(paramMap)
-    intercept[XGBoostError] {
-      val model = xgb.fit(trainingDF)
+    intercept[SparkException] {
+      xgb.fit(trainingDF)
     }
   }
 

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRabitRegressionSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ Copyright (c) 2014-2022 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -16,14 +16,15 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.{Rabit, XGBoostError}
-import ml.dmlc.xgboost4j.scala.{Booster, DMatrix}
-import org.apache.spark.TaskFailedListener
-import org.apache.spark.SparkException
+import ml.dmlc.xgboost4j.java.Rabit
+import ml.dmlc.xgboost4j.scala.Booster
 import scala.collection.JavaConverters._
+
 import org.apache.spark.sql._
 import org.scalatest.FunSuite
 
+import org.apache.spark.SparkException
+
 class XGBoostRabitRegressionSuite extends FunSuite with PerTest {
   val predictionErrorMin = 0.00001f
   val maxFailure = 2;
@@ -33,15 +34,6 @@ class XGBoostRabitRegressionSuite extends FunSuite with PerTest {
     .config("spark.kryo.classesToRegister", classOf[Booster].getName)
     .master(s"local[${numWorkers},${maxFailure}]")
 
-  private def waitAndCheckSparkShutdown(waitMiliSec: Int): Boolean = {
-    var totalWaitedTime = 0L
-    while (!ss.sparkContext.isStopped && totalWaitedTime <= waitMiliSec) {
-      Thread.sleep(10)
-      totalWaitedTime += 10
-    }
-    return ss.sparkContext.isStopped
-  }
-
   test("test classification prediction parity w/o ring reduce") {
     val training = buildDataFrame(Classification.train)
     val testDF = buildDataFrame(Classification.test)
@@ -91,14 +83,11 @@ class XGBoostRabitRegressionSuite extends FunSuite with PerTest {
   }
 
   test("test rabit timeout fail handle") {
-    // disable spark kill listener to verify if rabit_timeout take effect and kill tasks
-    TaskFailedListener.killerStarted = true
-
     val training = buildDataFrame(Classification.train)
     // mock rank 0 failure during 8th allreduce synchronization
     Rabit.mockList = Array("0,8,0,0").toList.asJava
 
-    try {
+    intercept[SparkException] {
       new XGBoostClassifier(Map(
         "eta" -> "0.1",
         "max_depth" -> "10",
@@ -108,39 +97,7 @@ class XGBoostRabitRegressionSuite extends FunSuite with PerTest {
         "num_workers" -> numWorkers,
         "rabit_timeout" -> 0))
         .fit(training)
-    } catch {
-      case e: Throwable => // swallow anything
-    } finally {
-      // assume all tasks throw exception almost same time
-      // 100ms should be enough to exhaust all retries
-      assert(waitAndCheckSparkShutdown(100) == true)
-      TaskFailedListener.killerStarted = false
     }
   }
 
-  test("test SparkContext should not be killed ") {
-    cancel("For some reason, sparkContext can't cancel the job locally in the CI env," +
-      "which will be resolved when introducing barrier mode")
-    val training = buildDataFrame(Classification.train)
-    // mock rank 0 failure during 8th allreduce synchronization
-    Rabit.mockList = Array("0,8,0,0").toList.asJava
-
-    try {
-      new XGBoostClassifier(Map(
-        "eta" -> "0.1",
-        "max_depth" -> "10",
-        "verbosity" -> "1",
-        "objective" -> "binary:logistic",
-        "num_round" -> 5,
-        "num_workers" -> numWorkers,
-        "kill_spark_context_on_worker_failure" -> false,
-        "rabit_timeout" -> 0))
-        .fit(training)
-    } catch {
-      case e: Throwable => // swallow anything
-    } finally {
-      // wait 3s to check if SparkContext is killed
-      assert(waitAndCheckSparkShutdown(3000) == false)
-    }
-  }
 }

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressorSuite.scala

@@ -21,7 +21,6 @@ import ml.dmlc.xgboost4j.scala.{DMatrix, XGBoost => ScalaXGBoost}
 import org.apache.spark.ml.linalg.{Vector, Vectors}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.{DataFrame, Row}
-import org.apache.spark.sql.types._
 import org.scalatest.FunSuite
 
 import org.apache.spark.ml.feature.VectorAssembler

jvm-packages/xgboost4j-spark/src/test/scala/org/apache/spark/SparkParallelismTrackerSuite.scala

@@ -1,151 +0,0 @@
-/*
- Copyright (c) 2014 by Contributors
-
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.
- */
-
-package org.apache.spark
-
-import org.scalatest.FunSuite
-import _root_.ml.dmlc.xgboost4j.scala.spark.PerTest
-import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.SparkSession
-
-import scala.math.min
-
-class SparkParallelismTrackerSuite extends FunSuite with PerTest {
-
-  val numParallelism: Int = min(Runtime.getRuntime.availableProcessors(), 4)
-
-  override protected def sparkSessionBuilder: SparkSession.Builder = SparkSession.builder()
-    .master(s"local[${numParallelism}]")
-    .appName("XGBoostSuite")
-    .config("spark.ui.enabled", true)
-    .config("spark.driver.memory", "512m")
-    .config("spark.task.cpus", 1)
-
-  private def waitAndCheckSparkShutdown(waitMiliSec: Int): Boolean = {
-    var totalWaitedTime = 0L
-    while (!ss.sparkContext.isStopped && totalWaitedTime <= waitMiliSec) {
-      Thread.sleep(100)
-      totalWaitedTime += 100
-    }
-    ss.sparkContext.isStopped
-  }
-
-  test("tracker should not affect execution result when timeout is not larger than 0") {
-    val nWorkers = numParallelism
-    val rdd: RDD[Int] = sc.parallelize(1 to nWorkers)
-    val tracker = new SparkParallelismTracker(sc, 10000, nWorkers)
-    val disabledTracker = new SparkParallelismTracker(sc, 0, nWorkers)
-    assert(tracker.execute(rdd.sum()) == rdd.sum())
-    assert(disabledTracker.execute(rdd.sum()) == rdd.sum())
-  }
-
-  test("tracker should throw exception if parallelism is not sufficient") {
-    val nWorkers = numParallelism * 3
-    val rdd: RDD[Int] = sc.parallelize(1 to nWorkers)
-    val tracker = new SparkParallelismTracker(sc, 1000, nWorkers)
-    intercept[IllegalStateException] {
-      tracker.execute {
-        rdd.map { i =>
-          // Test interruption
-          Thread.sleep(Long.MaxValue)
-          i
-        }.sum()
-      }
-    }
-  }
-
-  test("tracker should throw exception if parallelism is not sufficient with" +
-    " spark.task.cpus larger than 1") {
-    sc.conf.set("spark.task.cpus", "2")
-    val nWorkers = numParallelism
-    val rdd: RDD[Int] = sc.parallelize(1 to nWorkers)
-    val tracker = new SparkParallelismTracker(sc, 1000, nWorkers)
-    intercept[IllegalStateException] {
-      tracker.execute {
-        rdd.map { i =>
-          // Test interruption
-          Thread.sleep(Long.MaxValue)
-          i
-        }.sum()
-      }
-    }
-  }
-
-  test("tracker should not kill SparkContext when killSparkContextOnWorkerFailure=false") {
-    val nWorkers = numParallelism
-    val tracker = new SparkParallelismTracker(sc, 0, nWorkers, false)
-    val rdd: RDD[Int] = sc.parallelize(1 to nWorkers, nWorkers)
-    try {
-      tracker.execute {
-        rdd.map { i =>
-          val partitionId = TaskContext.get().partitionId()
-          if (partitionId == 0) {
-            throw new RuntimeException("mocking task failing")
-          }
-          i
-        }.sum()
-      }
-    } catch {
-      case e: Exception => // catch the exception
-    } finally {
-      // wait 3s to check if SparkContext is killed
-      assert(waitAndCheckSparkShutdown(3000) == false)
-    }
-  }
-
-  test("tracker should cancel the correct job when killSparkContextOnWorkerFailure=false") {
-    val nWorkers = 2
-    val tracker = new SparkParallelismTracker(sc, 0, nWorkers, false)
-    val rdd: RDD[Int] = sc.parallelize(1 to 10, nWorkers)
-    val thread = new TestThread(sc)
-    thread.start()
-    try {
-      tracker.execute {
-        rdd.map { i =>
-          Thread.sleep(100)
-          val partitionId = TaskContext.get().partitionId()
-          if (partitionId == 0) {
-            throw new RuntimeException("mocking task failing")
-          }
-          i
-        }.sum()
-      }
-    } catch {
-      case e: Exception => // catch the exception
-    } finally {
-      thread.join(8000)
-      // wait 3s to check if SparkContext is killed
-      assert(waitAndCheckSparkShutdown(3000) == false)
-    }
-  }
-
-  private[this] class TestThread(sc: SparkContext) extends Thread {
-    override def run(): Unit = {
-      var sum: Double = 0.0f
-      try {
-        val rdd = sc.parallelize(1 to 4, 2)
-        sum = rdd.mapPartitions(iter => {
-          // sleep 2s to ensure task is alive when cancelling other jobs
-          Thread.sleep(2000)
-          iter
-        }).sum()
-      } finally {
-        // get the correct result
-        assert(sum.toInt == 10)
-      }
-    }
-  }
-}