Address review comments + update docs

Author: dhruve

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -600,10 +600,10 @@ package object config {
   // Threshold above which we abort the TaskSet if a task could not be scheduled because of complete
   // blacklisting.
   private[spark] val UNSCHEDULABLE_TASKSET_TIMEOUT =
-    ConfigBuilder("spark.scheduler.unschedulableTaskSetTimeout")
-      .doc("The timeout in seconds to wait before aborting a TaskSet to acquire a new executor " +
-        "and schedule a task which was previously unschedulable because of being completely " +
-        "blacklisted.")
+    ConfigBuilder("spark.scheduler.blacklist.unschedulableTaskSetTimeout")
+      .doc("The timeout in seconds to wait to try to acquire a new executor and schedule a task " +
+        "before aborting a TaskSet which was previously unschedulable because of being " +
+        "completely blacklisted.")
       .timeConf(TimeUnit.SECONDS)
       .checkValue(v => v >= 0, "The value should be a non negative time value.")
       .createWithDefault(120)

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

@@ -22,10 +22,9 @@ import java.util.{Locale, Timer, TimerTask}
 import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
 import java.util.concurrent.atomic.AtomicLong
 
-import scala.collection.Set
+import scala.collection.{Set, mutable}
 import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet}
 import scala.util.Random
-
 import org.apache.spark._
 import org.apache.spark.TaskState.TaskState
 import org.apache.spark.executor.ExecutorMetrics
@@ -118,7 +117,6 @@ private[spark] class TaskSchedulerImpl(
   protected val executorIdToHost = new HashMap[String, String]
 
   private val abortTimer = new Timer(true)
-
   private val clock = new SystemClock
 
   protected val unschedulableTaskSetToExpiryTime = new HashMap[TaskSetManager, Long]
@@ -430,14 +428,13 @@ private[spark] class TaskSchedulerImpl(
               // executor. If we cannot find one, we abort immediately. Else we kill the idle
               // executor and kick off an abortTimer which after waiting will abort the taskSet if
               // we were unable to schedule any task from the taskSet.
-              // Note 1: We keep a track of schedulability on a per taskSet basis rather than on a
-              // per task basis.
+              // 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 is because we rely on the
-              // ExecutorAllocationManager to acquire a new executor based on the pending tasks and
-              // it won't release any blacklisted executors which idle timeout after we kill an
-              // executor to acquire a new one, resulting in the abort timer to expire and abort the
-              // taskSet.
+              // 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.
               executorIdToRunningTaskIds.find(x => !isExecutorBusy(x._1)) match {
                 case Some (x) =>
                   val executorId = x._1
@@ -465,18 +462,16 @@ private[spark] class TaskSchedulerImpl(
                   }
                 case _ => // Abort Immediately
                   logInfo("Cannot schedule any task because of complete blacklisting. No idle" +
-                  s" executors could be found. Aborting $taskSet." )
+                  s" executors can be found to kill. Aborting $taskSet." )
                   taskSet.abortSinceCompletelyBlacklisted(taskIndex.get)
               }
-            case _ => // Do nothing.
-            }
-          } else {
-            // If a task was scheduled, we clear the expiry time for the taskSet. The abort timer
-            // checks this entry to decide if we want to abort the taskSet.
-            if (unschedulableTaskSetToExpiryTime.contains(taskSet)) {
-              unschedulableTaskSetToExpiryTime.remove(taskSet)
-            }
+            case _ => // Do nothing if no tasks completely blacklisted.
           }
+        } else {
+          // If a task was scheduled, we clear the expiry time for the taskSet. The abort timer
+          // checks this entry to decide if we want to abort the taskSet.
+          unschedulableTaskSetToExpiryTime.remove(taskSet)
+        }
 
         if (launchedAnyTask && taskSet.isBarrier) {
           // Check whether the barrier tasks are partially launched.

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

@@ -530,8 +530,6 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
     )
     // Wait for the failed task to propagate.
     Thread.sleep(500)
-    //    taskScheduler.handleFailedTask(tsm, failedTask.taskId, TaskState.FAILED, TaskResultLost)
-    //    tsm.handleFailedTask(failedTask.taskId, TaskState.FAILED, TaskResultLost)
 
     when(stageToMockTaskSetBlacklist(0).isExecutorBlacklistedForTask("executor0", failedTask.index))
       .thenReturn(true)
@@ -548,7 +546,6 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
   }
 
   test("SPARK-22148 try to acquire a new executor when task is unschedulable with 1 executor") {
-
     taskScheduler = setupSchedulerWithMockTaskSetBlacklist(
       config.UNSCHEDULABLE_TASKSET_TIMEOUT.key -> "10")
 

docs/configuration.md

@@ -1583,6 +1583,14 @@ Apart from these, the following properties are also available, and may be useful
     driver using more memory.
   </td>
 </tr>
+<tr>
+  <td><code>spark.scheduler.blacklist.unschedulableTaskSetTimeout</code></td>
+  <td>120s</td>
+  <td>
+    The timeout in seconds to wait to try to acquire a new executor and schedule a task before
+    aborting a TaskSet which was previously unschedulable because of being completely blacklisted.
+  </td>
+</tr>
 <tr>
   <td><code>spark.blacklist.enabled</code></td>
   <td>