BlacklistTracker takes a SchedulerBackend as input

Author: José Hiram Soltren

core/src/main/scala/org/apache/spark/scheduler/BlacklistTracker.scala

@@ -21,9 +21,10 @@ import java.util.concurrent.atomic.AtomicReference
 
 import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet}
 
-import org.apache.spark.{SparkConf, SparkContext}
+import org.apache.spark.SparkConf
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config
+import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend
 import org.apache.spark.util.{Clock, SystemClock, Utils}
 
 /**
@@ -50,10 +51,11 @@ import org.apache.spark.util.{Clock, SystemClock, Utils}
 private[scheduler] class BlacklistTracker (
     private val listenerBus: LiveListenerBus,
     conf: SparkConf,
+    scheduler: Option[CoarseGrainedSchedulerBackend],
     clock: Clock = new SystemClock()) extends Logging {
 
-  def this(sc: SparkContext) = {
-    this(sc.listenerBus, sc.getConf)
+  def this(sched: CoarseGrainedSchedulerBackend) = {
+    this(sched.listenerBus, sched.conf, Some(sched))
   }
 
   BlacklistTracker.validateBlacklistConfs(conf)

core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala

@@ -32,6 +32,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config
 import org.apache.spark.scheduler.SchedulingMode.SchedulingMode
 import org.apache.spark.scheduler.TaskLocality.TaskLocality
+import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend
 import org.apache.spark.scheduler.local.LocalSchedulerBackend
 import org.apache.spark.storage.BlockManagerId
 import org.apache.spark.util.{AccumulatorV2, ThreadUtils, Utils}
@@ -54,7 +55,7 @@ import org.apache.spark.util.{AccumulatorV2, ThreadUtils, Utils}
 private[spark] class TaskSchedulerImpl private[scheduler](
     val sc: SparkContext,
     val maxTaskFailures: Int,
-    blacklistTrackerOpt: Option[BlacklistTracker],
+    var blacklistTrackerOpt: Option[BlacklistTracker],
     isLocal: Boolean = false)
   extends TaskScheduler with Logging
 {
@@ -63,14 +64,14 @@ private[spark] class TaskSchedulerImpl private[scheduler](
     this(
       sc,
       sc.conf.get(config.MAX_TASK_FAILURES),
-      TaskSchedulerImpl.maybeCreateBlacklistTracker(sc))
+      null)
   }
 
   def this(sc: SparkContext, maxTaskFailures: Int, isLocal: Boolean) = {
     this(
       sc,
       maxTaskFailures,
-      TaskSchedulerImpl.maybeCreateBlacklistTracker(sc),
+      null,
       isLocal = isLocal)
   }
 
@@ -148,6 +149,7 @@ private[spark] class TaskSchedulerImpl private[scheduler](
 
   def initialize(backend: SchedulerBackend) {
     this.backend = backend
+    blacklistTrackerOpt = TaskSchedulerImpl.maybeCreateBlacklistTracker(backend)
     // temporarily set rootPool name to empty
     rootPool = new Pool("", schedulingMode, 0, 0)
     schedulableBuilder = {
@@ -701,11 +703,16 @@ private[spark] object TaskSchedulerImpl {
     retval.toList
   }
 
-  private def maybeCreateBlacklistTracker(sc: SparkContext): Option[BlacklistTracker] = {
-    if (BlacklistTracker.isBlacklistEnabled(sc.conf)) {
-      Some(new BlacklistTracker(sc))
-    } else {
-      None
+  private def maybeCreateBlacklistTracker(sched: SchedulerBackend): Option[BlacklistTracker] = {
+    sched match {
+      case backend: CoarseGrainedSchedulerBackend =>
+        if (BlacklistTracker.isBlacklistEnabled(backend.conf)) {
+          Some(new BlacklistTracker(backend))
+        } else {
+          None
+        }
+      case _ =>
+        None
     }
   }
 

core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala

@@ -49,7 +49,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
   protected val totalCoreCount = new AtomicInteger(0)
   // Total number of executors that are currently registered
   protected val totalRegisteredExecutors = new AtomicInteger(0)
-  protected val conf = scheduler.sc.conf
+  val conf = scheduler.sc.conf
   private val maxRpcMessageSize = RpcUtils.maxMessageSizeBytes(conf)
   private val defaultAskTimeout = RpcUtils.askRpcTimeout(conf)
   // Submit tasks only after (registered resources / total expected resources)
@@ -73,7 +73,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
   @GuardedBy("CoarseGrainedSchedulerBackend.this")
   private var numPendingExecutors = 0
 
-  private val listenerBus = scheduler.sc.listenerBus
+  val listenerBus = scheduler.sc.listenerBus
 
   // Executors we have requested the cluster manager to kill that have not died yet; maps
   // the executor ID to whether it was explicitly killed by the driver (and thus shouldn't

core/src/test/scala/org/apache/spark/scheduler/BlacklistTrackerSuite.scala

@@ -41,7 +41,7 @@ class BlacklistTrackerSuite extends SparkFunSuite with BeforeAndAfterEach with M
     scheduler = mockTaskSchedWithConf(conf)
 
     clock.setTime(0)
-    blacklist = new BlacklistTracker(null, conf, clock)
+    blacklist = new BlacklistTracker(null, conf, null, clock)
   }
 
   override def afterEach(): Unit = {
@@ -90,7 +90,7 @@ class BlacklistTrackerSuite extends SparkFunSuite with BeforeAndAfterEach with M
 
     clock.setTime(0)
     listenerBusMock = mock[LiveListenerBus]
-    blacklist = new BlacklistTracker(listenerBusMock, conf, clock)
+    blacklist = new BlacklistTracker(listenerBusMock, conf, null, clock)
   }
 
   test("Blacklisting individual tasks and checking for SparkListenerEvents") {

core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala

@@ -426,7 +426,7 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
     val clock = new ManualClock
     // We don't directly use the application blacklist, but its presence triggers blacklisting
     // within the taskset.
-    val blacklistTrackerOpt = Some(new BlacklistTracker(null, conf, clock))
+    val blacklistTrackerOpt = Some(new BlacklistTracker(null, conf, null, clock))
     val manager = new TaskSetManager(sched, taskSet, 4, blacklistTrackerOpt, clock)
 
     {