Addressed tgraves and mridulm comments

Author: venkata91

core/src/main/java/org/apache/spark/SparkFirehoseListener.java

@@ -162,9 +162,14 @@ public void onSpeculativeTaskSubmitted(SparkListenerSpeculativeTaskSubmitted spe
     onEvent(speculativeTask);
   }
 
-  public void onUnschedulableTaskSet(
-      SparkListenerUnschedulableTaskSet unschedulableTaskSet) {
-    onEvent(unschedulableTaskSet);
+  public void onUnschedulableTaskSetAdded(
+      SparkListenerUnschedulableTaskSetAdded unschedulableTaskSetAdded) {
+    onEvent(unschedulableTaskSetAdded);
+  }
+
+  public void onUnschedulableTaskSetRemoved(
+      SparkListenerUnschedulableTaskSetRemoved unschedulableTaskSetRemoved) {
+    onEvent(unschedulableTaskSetRemoved);
   }
 
   @Override

core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala

@@ -281,7 +281,7 @@ private[spark] class ExecutorAllocationManager(
   private def maxNumExecutorsNeededPerResourceProfile(rpId: Int): Int = {
     val pending = listener.totalPendingTasksPerResourceProfile(rpId)
     val pendingSpeculative = listener.pendingSpeculativeTasksPerResourceProfile(rpId)
-    val numUnschedulables = listener.pendingUnschedulableTaskSetsPerResourceProfile(rpId)
+    val unschedulableTaskSets = listener.pendingUnschedulableTaskSetsPerResourceProfile(rpId)
     val running = listener.totalRunningTasksPerResourceProfile(rpId)
     val numRunningOrPendingTasks = pending + running
     val rp = resourceProfileManager.resourceProfileFromId(rpId)
@@ -290,6 +290,7 @@ private[spark] class ExecutorAllocationManager(
       s" tasksperexecutor: $tasksPerExecutor")
     val maxNeeded = math.ceil(numRunningOrPendingTasks * executorAllocationRatio /
       tasksPerExecutor).toInt
+
     val maxNeededWithSpeculationLocalityOffset =
       if (tasksPerExecutor > 1 && maxNeeded == 1 && pendingSpeculative > 0) {
       // If we have pending speculative tasks and only need a single executor, allocate one more
@@ -299,11 +300,11 @@ private[spark] class ExecutorAllocationManager(
       maxNeeded
     }
 
-    // Since the maxNeededWithSpeculationLocalityOffset already includes the num executors needed
-    // to run unschedulable tasks, we would only try to add when the
-    // maxNeededWithSpeculationLocalityOffset is lesser than the active executors
-    if (numUnschedulables > 0) {
-      val maxNeededForUnschedulables = math.ceil(numUnschedulables * executorAllocationRatio /
+    if (unschedulableTaskSets > 0) {
+      // Request additional executors only if maxNeededWithSpeculationLocalityOffset is less than
+      // current active executors, if not requesting maxNeededWithSpeculationLocalityOffset should
+      // be enough to make the unschedulable tasks schedulable now.
+      val maxNeededForUnschedulables = math.ceil(unschedulableTaskSets * executorAllocationRatio /
         tasksPerExecutor).toInt
       math.max(maxNeededWithSpeculationLocalityOffset,
         executorMonitor.executorCountWithResourceProfile(rpId) + maxNeededForUnschedulables)
@@ -810,20 +811,25 @@ private[spark] class ExecutorAllocationManager(
       }
     }
 
-    override def onUnschedulableTaskSet
-      (unschedulableTaskSet: SparkListenerUnschedulableTaskSet): Unit = {
-      val stageId = unschedulableTaskSet.stageId
-      val stageAttemptId = unschedulableTaskSet.stageAttemptId
+    override def onUnschedulableTaskSetAdded(
+        unschedulableTaskSetAdded: SparkListenerUnschedulableTaskSetAdded): Unit = {
+      val stageId = unschedulableTaskSetAdded.stageId
+      val stageAttemptId = unschedulableTaskSetAdded.stageAttemptId
       allocationManager.synchronized {
-        (stageId, stageAttemptId) match {
-          case (Some(stageId), Some(stageAttemptId)) =>
-            val stageAttempt = StageAttempt(stageId, stageAttemptId)
-            unschedulableTaskSets.add(stageAttempt)
-          case (None, _) =>
-            // Clear unschedulableTaskSets since atleast one task becomes schedulable now
-            unschedulableTaskSets.clear()
-        }
-        allocationManager.onSchedulerBacklogged()
+        val stageAttempt = StageAttempt(stageId, stageAttemptId)
+        unschedulableTaskSets.add(stageAttempt)
+      }
+      allocationManager.onSchedulerBacklogged()
+    }
+
+    override def onUnschedulableTaskSetRemoved(
+        unschedulableTaskSetRemoved: SparkListenerUnschedulableTaskSetRemoved): Unit = {
+      val stageId = unschedulableTaskSetRemoved.stageId
+      val stageAttemptId = unschedulableTaskSetRemoved.stageAttemptId
+      allocationManager.synchronized {
+        // Clear unschedulableTaskSets since atleast one task becomes schedulable now
+        val stageAttempt = StageAttempt(stageId, stageAttemptId)
+        unschedulableTaskSets.remove(stageAttempt)
       }
     }
 
@@ -867,9 +873,9 @@ private[spark] class ExecutorAllocationManager(
       numTotalTasks - numRunning
     }
 
-    // Currently TaskSetManager.getCompletelyBlacklistedTaskIfAny only takes the first
-    // unschedulable task due to blacklisting. So keeping track of unschedulableTaskSets
-    // should be enough as we'll always have no more than a task unschedulable at any time.
+    // Currently TaskSetManager.getCompletelyBlacklistedTaskIfAny only takes the first unschedulable
+    // task found due to blacklisting. This way we only need to keep track of the unschedulable
+    // tasksets which is an indirect way to get the current number of unschedulable tasks.
     def pendingUnschedulableTaskSetsPerResourceProfile(rp: Int): Int = {
       val attempts = resourceProfileIdToStageAttempt.getOrElse(rp, Set.empty).toSeq
       attempts.filter(attempt => unschedulableTaskSets.contains(attempt)).size

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

@@ -312,13 +312,23 @@ private[spark] class DAGScheduler(
   }
 
   /**
-   * Called by the TaskSetManager when there is an unschedulable blacklist task and dynamic
+   * Called by the TaskSetManager when a taskset becomes unschedulable due to blacklisting and
+   * dynamic allocation is enabled
+   */
+  def unschedulableTaskSetAdded(
+       stageId: Int,
+       stageAttemptId: Int): Unit = {
+    eventProcessLoop.post(UnschedulableTaskSetAdded(stageId, stageAttemptId))
+  }
+
+  /**
+   * Called by the TaskSetManager when an unschedulable taskset becomes schedulable and dynamic
    * allocation is enabled
    */
-  def unschedulableBlacklistTaskSubmitted(
-       stageId: Option[Int],
-       stageAttemptId: Option[Int]): Unit = {
-    eventProcessLoop.post(UnschedulableBlacklistTaskSubmitted(stageId, stageAttemptId))
+  def unschedulableTaskSetRemoved(
+    stageId: Int,
+    stageAttemptId: Int): Unit = {
+    eventProcessLoop.post(UnschedulableTaskSetRemoved(stageId, stageAttemptId))
   }
 
   private[scheduler]
@@ -1024,11 +1034,18 @@ private[spark] class DAGScheduler(
     listenerBus.post(SparkListenerSpeculativeTaskSubmitted(task.stageId, task.stageAttemptId))
   }
 
-  private[scheduler] def handleUnschedulableBlacklistTaskSubmitted(
-      stageId: Option[Int],
-      stageAttemptId: Option[Int]): Unit = {
+  private[scheduler] def handleUnschedulableTaskSetAdded(
+      stageId: Int,
+      stageAttemptId: Int): Unit = {
     listenerBus.post(
-      SparkListenerUnschedulableTaskSet(stageId, stageAttemptId))
+      SparkListenerUnschedulableTaskSetAdded(stageId, stageAttemptId))
+  }
+
+  private[scheduler] def handleUnschedulableTaskSetRemoved(
+    stageId: Int,
+    stageAttemptId: Int): Unit = {
+    listenerBus.post(
+      SparkListenerUnschedulableTaskSetRemoved(stageId, stageAttemptId))
   }
 
   private[scheduler] def handleTaskSetFailed(
@@ -2304,8 +2321,11 @@ private[scheduler] class DAGSchedulerEventProcessLoop(dagScheduler: DAGScheduler
     case SpeculativeTaskSubmitted(task) =>
       dagScheduler.handleSpeculativeTaskSubmitted(task)
 
-    case UnschedulableBlacklistTaskSubmitted(stageId, stageAttemptId) =>
-      dagScheduler.handleUnschedulableBlacklistTaskSubmitted(stageId, stageAttemptId)
+    case UnschedulableTaskSetAdded(stageId, stageAttemptId) =>
+      dagScheduler.handleUnschedulableTaskSetAdded(stageId, stageAttemptId)
+
+    case UnschedulableTaskSetRemoved(stageId, stageAttemptId) =>
+      dagScheduler.handleUnschedulableTaskSetRemoved(stageId, stageAttemptId)
 
     case GettingResultEvent(taskInfo) =>
       dagScheduler.handleGetTaskResult(taskInfo)

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

@@ -98,5 +98,10 @@ private[scheduler]
 case class SpeculativeTaskSubmitted(task: Task[_]) extends DAGSchedulerEvent
 
 private[scheduler]
-case class UnschedulableBlacklistTaskSubmitted(stageId: Option[Int], stageAttemptId: Option[Int])
+case class UnschedulableTaskSetAdded(stageId: Int, stageAttemptId: Int)
   extends DAGSchedulerEvent
+
+private[scheduler]
+case class UnschedulableTaskSetRemoved(stageId: Int, stageAttemptId: Int)
+  extends DAGSchedulerEvent
+

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

@@ -158,9 +158,14 @@ case class SparkListenerNodeUnblacklisted(time: Long, hostId: String)
   extends SparkListenerEvent
 
 @DeveloperApi
-case class SparkListenerUnschedulableTaskSet(
-  stageId: Option[Int],
-  stageAttemptId: Option[Int]) extends SparkListenerEvent
+case class SparkListenerUnschedulableTaskSetAdded(
+  stageId: Int,
+  stageAttemptId: Int) extends SparkListenerEvent
+
+@DeveloperApi
+case class SparkListenerUnschedulableTaskSetRemoved(
+  stageId: Int,
+  stageAttemptId: Int) extends SparkListenerEvent
 
 @DeveloperApi
 case class SparkListenerBlockUpdated(blockUpdatedInfo: BlockUpdatedInfo) extends SparkListenerEvent
@@ -345,10 +350,18 @@ private[spark] trait SparkListenerInterface {
   def onNodeUnblacklisted(nodeUnblacklisted: SparkListenerNodeUnblacklisted): Unit
 
   /**
-   * Called when both dynamic allocation is enabled and there is an unschedulable blacklist task
+   * Called when a taskset becomes unschedulable due to blacklisting and dynamic allocation
+   * is enabled
+   */
+  def onUnschedulableTaskSetAdded(
+      unschedulableTaskSetAdded: SparkListenerUnschedulableTaskSetAdded): Unit
+
+  /**
+   * Called when an unschedulable taskset becomes schedulable and dynamic allocation
+   * is enabled
    */
-  def onUnschedulableTaskSet(
-    blacklistTask: SparkListenerUnschedulableTaskSet): Unit
+  def onUnschedulableTaskSetRemoved(
+      unschedulableTaskSetRemoved: SparkListenerUnschedulableTaskSetRemoved): Unit
 
   /**
    * Called when the driver receives a block update info.
@@ -436,8 +449,11 @@ abstract class SparkListener extends SparkListenerInterface {
   override def onNodeUnblacklisted(
       nodeUnblacklisted: SparkListenerNodeUnblacklisted): Unit = { }
 
-  override def onUnschedulableTaskSet(
-    unschedulableTaskSet: SparkListenerUnschedulableTaskSet): Unit = { }
+  override def onUnschedulableTaskSetAdded(
+      unschedulableTaskSetAdded: SparkListenerUnschedulableTaskSetAdded): Unit = { }
+
+  override def onUnschedulableTaskSetRemoved(
+      unschedulableTaskSetRemoved: SparkListenerUnschedulableTaskSetRemoved): Unit = { }
 
   override def onBlockUpdated(blockUpdated: SparkListenerBlockUpdated): Unit = { }
 

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

@@ -79,8 +79,10 @@ private[spark] trait SparkListenerBus
         listener.onBlockUpdated(blockUpdated)
       case speculativeTaskSubmitted: SparkListenerSpeculativeTaskSubmitted =>
         listener.onSpeculativeTaskSubmitted(speculativeTaskSubmitted)
-      case unschedulableTaskSet: SparkListenerUnschedulableTaskSet =>
-        listener.onUnschedulableTaskSet(unschedulableTaskSet)
+      case unschedulableTaskSetAdded: SparkListenerUnschedulableTaskSetAdded =>
+        listener.onUnschedulableTaskSetAdded(unschedulableTaskSetAdded)
+      case unschedulableTaskSetRemoved: SparkListenerUnschedulableTaskSetRemoved =>
+        listener.onUnschedulableTaskSetRemoved(unschedulableTaskSetRemoved)
       case resourceProfileAdded: SparkListenerResourceProfileAdded =>
         listener.onResourceProfileAdded(resourceProfileAdded)
       case _ => listener.onOtherEvent(event)

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

@@ -637,17 +637,20 @@ private[spark] class TaskSchedulerImpl(
         if (!launchedAnyTask) {
           taskSet.getCompletelyBlacklistedTaskIfAny(hostToExecutors).foreach { taskIndex =>
               // If the taskSet is unschedulable we try to find an existing idle blacklisted
-              // executor. If we cannot find one, we abort immediately. Else we kill the idle
-              // executor and kick off an abortTimer which if it doesn't schedule a task within the
-              // the timeout will abort the taskSet if we were unable to schedule any task from the
-              // taskSet.
+              // executor and kill the idle executor and kick off an abortTimer which if it doesn't
+              // schedule a task within the the timeout will abort the taskSet if we were unable to
+              // schedule any task from the taskSet.
               // Note 1: We keep track of schedulability on a per taskSet basis rather than on a per
               // task basis.
               // Note 2: The taskSet can still be aborted when there are more than one idle
               // blacklisted executors and dynamic allocation is on. This can happen when a killed
               // idle executor isn't replaced in time by ExecutorAllocationManager as it relies on
               // pending tasks and doesn't kill executors on idle timeouts, resulting in the abort
               // timer to expire and abort the taskSet.
+              //
+              // If there are no idle executors and dynamic allocation is enabled, then we would
+              // notify ExecutorAllocationManager to allocate more executors to schedule the
+              // unschedulable tasks else we will abort immediately.
               executorIdToRunningTaskIds.find(x => !isExecutorBusy(x._1)) match {
                 case Some ((executorId, _)) =>
                   if (!unschedulableTaskSetToExpiryTime.contains(taskSet)) {
@@ -659,11 +662,10 @@ private[spark] class TaskSchedulerImpl(
                   // in order to provision more executors to make them schedulable
                   if (Utils.isDynamicAllocationEnabled(conf)) {
                     if (!unschedulableTaskSetToExpiryTime.contains(taskSet)) {
-                      logInfo(s"Notifying ExecutorAllocationManager to add executors to" +
-                        s" schedule the unschedulable blacklisted task before aborting $taskSet.")
-                      dagScheduler.unschedulableBlacklistTaskSubmitted(
-                        Some(taskSet.taskSet.stageId),
-                        Some(taskSet.taskSet.stageAttemptId))
+                      logInfo(s"Notifying ExecutorAllocationManager to allocate more executors to" +
+                        s" schedule the unschedulable task before aborting $taskSet.")
+                      dagScheduler.unschedulableTaskSetAdded(taskSet.taskSet.stageId,
+                        taskSet.taskSet.stageAttemptId)
                       updateUnschedulableTaskSetTimeoutAndStartAbortTimer(taskSet, taskIndex)
                     }
                   } else {
@@ -686,7 +688,8 @@ private[spark] class TaskSchedulerImpl(
               "recently scheduled.")
             // Notify ExecutorAllocationManager as well as other subscribers that a task now
             // recently becomes schedulable
-            dagScheduler.unschedulableBlacklistTaskSubmitted(None, None)
+            dagScheduler.unschedulableTaskSetRemoved(taskSet.taskSet.stageId,
+              taskSet.taskSet.stageAttemptId)
             unschedulableTaskSetToExpiryTime.clear()
           }
         }

core/src/test/scala/org/apache/spark/ExecutorAllocationManagerSuite.scala

@@ -524,7 +524,7 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite {
     assert(maxNumExecutorsNeededPerResourceProfile(manager, defaultProfile) == 1)
 
     // Stage 0 becomes unschedulable due to blacklisting
-    post(SparkListenerUnschedulableTaskSet(Some(0), Some(0)))
+    post(SparkListenerUnschedulableTaskSetAdded(0, 0))
     clock.advance(1000)
     manager invokePrivate _updateAndSyncNumExecutorsTarget(clock.nanoTime())
     // Assert that we are getting additional executor to schedule unschedulable tasks
@@ -534,8 +534,7 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite {
     // Add a new executor
     onExecutorAddedDefaultProfile(manager, "1")
     // Now once the task becomes schedulable, clear the unschedulableTaskSets
-    // by posting unschedulable event with (None, None)
-    post(SparkListenerUnschedulableTaskSet(None, None))
+    post(SparkListenerUnschedulableTaskSetRemoved(0, 0))
     clock.advance(1000)
     manager invokePrivate _updateAndSyncNumExecutorsTarget(clock.nanoTime())
     assert(numExecutorsTarget(manager, defaultProfile.id) === 1)
@@ -581,8 +580,8 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite {
     post(SparkListenerStageCompleted(createStageInfo(0, 2)))
 
     // Stage 1 and 2 becomes unschedulable now due to blacklisting
-    post(SparkListenerUnschedulableTaskSet(Some(1), Some(0)))
-    post(SparkListenerUnschedulableTaskSet(Some(2), Some(0)))
+    post(SparkListenerUnschedulableTaskSetAdded(1, 0))
+    post(SparkListenerUnschedulableTaskSetAdded(2, 0))
 
     clock.advance(1000)
     manager invokePrivate _updateAndSyncNumExecutorsTarget(clock.nanoTime())
@@ -594,12 +593,11 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite {
     onExecutorAddedDefaultProfile(manager, "3")
 
     // Now once the task becomes schedulable, clear the unschedulableTaskSets
-    // by posting unschedulable event with (None, None)
-    post(SparkListenerUnschedulableTaskSet(None, None))
+    post(SparkListenerUnschedulableTaskSetRemoved(1, 0))
     clock.advance(1000)
     manager invokePrivate _updateAndSyncNumExecutorsTarget(clock.nanoTime())
-    assert(numExecutorsTarget(manager, defaultProfile.id) === 2)
-    assert(maxNumExecutorsNeededPerResourceProfile(manager, defaultProfile) == 2)
+    assert(numExecutorsTarget(manager, defaultProfile.id) === 4)
+    assert(maxNumExecutorsNeededPerResourceProfile(manager, defaultProfile) == 5)
   }
 
   test("SPARK-31418: remove executors after unschedulable tasks end") {
@@ -639,7 +637,7 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite {
     (0 to 3).foreach { i => onExecutorRemoved(manager, i.toString) }
 
     // Now due to blacklisting, the task becomes unschedulable
-    post(SparkListenerUnschedulableTaskSet(Some(0), Some(0)))
+    post(SparkListenerUnschedulableTaskSetAdded(0, 0))
     clock.advance(1000)
     manager invokePrivate _updateAndSyncNumExecutorsTarget(clock.nanoTime())
     assert(numExecutorsTarget(manager, defaultProfile.id) === 2)
@@ -649,8 +647,7 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite {
     onExecutorAddedDefaultProfile(manager, "5")
 
     // Now once the task becomes schedulable, clear the unschedulableTaskSets
-    // by posting unschedulable event with (None, None)
-    post(SparkListenerUnschedulableTaskSet(None, None))
+    post(SparkListenerUnschedulableTaskSetRemoved(0, 0))
     clock.advance(1000)
     manager invokePrivate _updateAndSyncNumExecutorsTarget(clock.nanoTime())
     assert(numExecutorsTarget(manager, defaultProfile.id) === 1)