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

2018-01-24 19:31:38 -08:00 · 2018-01-24 19:31:38 -08:00 · 5ef684641b
commit 5ef684641b
parent f87802f00c
1 changed files with 3 additions and 2 deletions
--- 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)
    }
  }