Fixed SparkParallelTracker to work with Spark2.3 (dmlc#3062)

machine4life · Jan 25, 2018 · 5ef6846 · 5ef6846
1 parent f87802f
commit 5ef6846
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/jvm-packages/xgboost4j-spark/src/main/scala/org/apache/spark/SparkParallelismTracker.scala b/jvm-packages/xgboost4j-spark/src/main/scala/org/apache/spark/SparkParallelismTracker.scala
@@ -76,11 +76,12 @@ class SparkParallelismTracker(
   }
 
   private[this] def safeExecute[T](body: => T): T = {
-    sc.listenerBus.listeners.add(0, new TaskFailedListener)
+    val listener = new TaskFailedListener;
+    sc.addSparkListener(listener)
     try {
       body
     } finally {
-      sc.listenerBus.listeners.remove(0)
+      sc.listenerBus.removeListener(listener)
     }
   }