[WIP][SPARK-18886][CORE] Make locality wait time be the time since a TSM's available slots were fully utilized

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

@@ -119,4 +119,29 @@ private[spark] class Pool(
       parent.decreaseRunningTasks(taskNum)
     }
   }
+
+  override def updateAvailableSlots(numSlots: Float): Unit = {
+    schedulingMode match {
+      case SchedulingMode.FAIR =>
+        val usableWeights = schedulableQueue.asScala
+          .map(s => if (s.getSortedTaskSetQueue.nonEmpty) (s, s.weight) else (s, 0))
+        val totalWeights = usableWeights.map(_._2).sum
+        usableWeights.foreach { case (schedulable, usableWeight) =>
+          schedulable.updateAvailableSlots(
+            Math.max(numSlots * usableWeight / totalWeights, schedulable.minShare))
+        }
+      // for FIFO, allocate all slots to the first taskset in the queue
+      case SchedulingMode.FIFO =>
+        val sortedSchedulableQueue =
+          schedulableQueue.asScala.toSeq.sortWith(taskSetSchedulingAlgorithm.comparator)
+        var isFirst = true
+        for (schedulable <- sortedSchedulableQueue) {
+          schedulable.updateAvailableSlots(if (isFirst) numSlots else 0)
+          isFirst = false
+        }
+      case _ =>
+        val msg = s"Unsupported scheduling mode: $schedulingMode. Use FAIR or FIFO instead."
+        throw new IllegalArgumentException(msg)
+    }
+  }
 }

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

@@ -39,6 +39,7 @@ private[spark] trait Schedulable {
   def stageId: Int
   def name: String
 
+  def updateAvailableSlots(numSlots: Float): Unit
   def addSchedulable(schedulable: Schedulable): Unit
   def removeSchedulable(schedulable: Schedulable): Unit
   def getSchedulableByName(name: String): Schedulable

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

@@ -36,7 +36,7 @@ import org.apache.spark.rpc.RpcEndpoint
 import org.apache.spark.scheduler.SchedulingMode.SchedulingMode
 import org.apache.spark.scheduler.TaskLocality.TaskLocality
 import org.apache.spark.storage.BlockManagerId
-import org.apache.spark.util.{AccumulatorV2, SystemClock, ThreadUtils, Utils}
+import org.apache.spark.util.{AccumulatorV2, Clock, SystemClock, ThreadUtils, Utils}
 
 /**
  * Schedules tasks for multiple types of clusters by acting through a SchedulerBackend.
@@ -61,7 +61,8 @@ import org.apache.spark.util.{AccumulatorV2, SystemClock, ThreadUtils, Utils}
 private[spark] class TaskSchedulerImpl(
     val sc: SparkContext,
     val maxTaskFailures: Int,
-    isLocal: Boolean = false)
+    isLocal: Boolean = false,
+    clock: Clock = new SystemClock)
   extends TaskScheduler with Logging {
 
   import TaskSchedulerImpl._
@@ -125,10 +126,12 @@ private[spark] class TaskSchedulerImpl(
 
   protected val hostsByRack = new HashMap[String, HashSet[String]]
 
+  protected val executorIdToCores = new HashMap[String, Int]
+  protected var totalSlots = 0
+
   protected val executorIdToHost = new HashMap[String, String]
 
   private val abortTimer = new Timer(true)
-  private val clock = new SystemClock
   // Exposed for testing
   val unschedulableTaskSetToExpiryTime = new HashMap[TaskSetManager, Long]
 
@@ -403,6 +406,8 @@ private[spark] class TaskSchedulerImpl(
       if (!executorIdToRunningTaskIds.contains(o.executorId)) {
         hostToExecutors(o.host) += o.executorId
         executorAdded(o.executorId, o.host)
+        executorIdToCores(o.executorId) = o.cores
+        totalSlots += o.cores / CPUS_PER_TASK
         executorIdToHost(o.executorId) = o.host
         executorIdToRunningTaskIds(o.executorId) = HashSet[Long]()
         newExecAvail = true
@@ -439,6 +444,8 @@ private[spark] class TaskSchedulerImpl(
       }
     }
 
+    rootPool.updateAvailableSlots(totalSlots)
+
     // Take each TaskSet in our scheduling order, and then offer it each node in increasing order
     // of locality levels so that it gets a chance to launch local tasks on all of them.
     // NOTE: the preferredLocality order: PROCESS_LOCAL, NODE_LOCAL, NO_PREF, RACK_LOCAL, ANY
@@ -808,6 +815,7 @@ private[spark] class TaskSchedulerImpl(
       // happen below in the rootPool.executorLost() call.
       taskIds.foreach(cleanupTaskState)
     }
+    executorIdToCores.remove(executorId).foreach(totalSlots -= _ / CPUS_PER_TASK)
 
     val host = executorIdToHost(executorId)
     val execs = hostToExecutors.getOrElse(host, new HashSet)

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

@@ -75,6 +75,7 @@ private[spark] class TaskSetManager(
     .map { case (t, idx) => t.partitionId -> idx }.toMap
   val numTasks = tasks.length
   val copiesRunning = new Array[Int](numTasks)
+  var numAvailableSlot: Float = 0
 
   val speculationEnabled = conf.get(SPECULATION_ENABLED)
   // Quantile of tasks at which to start speculation
@@ -200,10 +201,13 @@ private[spark] class TaskSetManager(
   private[scheduler] var localityWaits = myLocalityLevels.map(getLocalityWait)
 
   // Delay scheduling variables: we keep track of our current locality level and the time we
-  // last launched a task at that level, and move up a level when localityWaits[curLevel] expires.
+  // last launched a task at that level and were also fully utilizing all slots
+  // Move up a level when localityWaits[curLevel] expires.
   // We then move down if we manage to launch a "more local" task.
   private var currentLocalityIndex = 0 // Index of our current locality level in validLocalityLevels
-  private var lastLaunchTime = clock.getTimeMillis()  // Time we last launched a task at this level
+  // Time we last launched a task at this level
+  // and number of running tasks was greater than or equal to numAvailableSlot
+  private var lastFullUtilizationTime = clock.getTimeMillis()
 
   override def schedulableQueue: ConcurrentLinkedQueue[Schedulable] = null
 
@@ -410,11 +414,13 @@ private[spark] class TaskSetManager(
           execId, host, taskLocality, speculative)
         taskInfos(taskId) = info
         taskAttempts(index) = info :: taskAttempts(index)
-        // Update our locality level for delay scheduling
+        // Update our locality level for delay scheduling if all allocated slots are being utilized
         // NO_PREF will not affect the variables related to delay scheduling
-        if (maxLocality != TaskLocality.NO_PREF) {
+        // Adding 1 to runningTasks, since runningTasks doesn't include this task yet
+        val utilizingAllSlots = runningTasks + 1 >= numAvailableSlot
+        if (maxLocality != TaskLocality.NO_PREF && utilizingAllSlots) {
           currentLocalityIndex = getLocalityIndex(taskLocality)
-          lastLaunchTime = curTime
+          lastFullUtilizationTime = curTime
         }
         // Serialize and return the task
         val serializedTask: ByteBuffer = try {
@@ -534,14 +540,14 @@ private[spark] class TaskSetManager(
         // This is a performance optimization: if there are no more tasks that can
         // be scheduled at a particular locality level, there is no point in waiting
         // for the locality wait timeout (SPARK-4939).
-        lastLaunchTime = curTime
+        lastFullUtilizationTime = curTime
         logDebug(s"No tasks for locality level ${myLocalityLevels(currentLocalityIndex)}, " +
           s"so moving to locality level ${myLocalityLevels(currentLocalityIndex + 1)}")
         currentLocalityIndex += 1
-      } else if (curTime - lastLaunchTime >= localityWaits(currentLocalityIndex)) {
-        // Jump to the next locality level, and reset lastLaunchTime so that the next locality
-        // wait timer doesn't immediately expire
-        lastLaunchTime += localityWaits(currentLocalityIndex)
+      } else if (curTime - lastFullUtilizationTime >= localityWaits(currentLocalityIndex)) {
+        // Jump to the next locality level, and reset lastFullUtilizationTime
+        // so that the next locality wait timer doesn't immediately expire
+        lastFullUtilizationTime += localityWaits(currentLocalityIndex)
         logDebug(s"Moving to ${myLocalityLevels(currentLocalityIndex + 1)} after waiting for " +
           s"${localityWaits(currentLocalityIndex)}ms")
         currentLocalityIndex += 1
@@ -1054,6 +1060,10 @@ private[spark] class TaskSetManager(
   def executorAdded(): Unit = {
     recomputeLocality()
   }
+
+  override def updateAvailableSlots(numSlots: Float): Unit = {
+    numAvailableSlot = numSlots
+  }
 }
 
 private[spark] object TaskSetManager {

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

@@ -51,6 +51,68 @@ class PoolSuite extends SparkFunSuite with LocalSparkContext {
     assert(nextTaskSetToSchedule.get.stageId === expectedStageId)
   }
 
+  test("validate FIFO slot distributions") {
+    sc = new SparkContext(LOCAL, APP_NAME)
+    val taskScheduler = new TaskSchedulerImpl(sc)
+
+    val rootPool = new Pool("", FIFO, 0, 0)
+    val schedulableBuilder = new FIFOSchedulableBuilder(rootPool)
+
+    val taskSetManager0 = createTaskSetManager(0, 2, taskScheduler)
+    val taskSetManager1 = createTaskSetManager(1, 2, taskScheduler)
+    val taskSetManager2 = createTaskSetManager(2, 2, taskScheduler)
+    schedulableBuilder.addTaskSetManager(taskSetManager0, null)
+    schedulableBuilder.addTaskSetManager(taskSetManager1, null)
+    schedulableBuilder.addTaskSetManager(taskSetManager2, null)
+
+    rootPool.updateAvailableSlots(9)
+    assert(taskSetManager0.numAvailableSlot === 9)
+    assert(taskSetManager1.numAvailableSlot === 0)
+    assert(taskSetManager2.numAvailableSlot === 0)
+  }
+
+  test("validate FAIR slot distributions") {
+    val xmlPath = getClass.getClassLoader.getResource("fairscheduler-with-valid-data.xml").getFile()
+    val conf = new SparkConf().set(SCHEDULER_ALLOCATION_FILE, xmlPath)
+    sc = new SparkContext(LOCAL, APP_NAME, conf)
+    val taskScheduler = new TaskSchedulerImpl(sc)
+
+    val rootPool = new Pool("", FAIR, 0, 0)
+    val schedulableBuilder = new FairSchedulableBuilder(rootPool, sc.conf)
+    schedulableBuilder.buildPools()
+
+    val taskSetManager0 = createTaskSetManager(0, 2, taskScheduler)
+    val taskSetManager1 = createTaskSetManager(1, 2, taskScheduler)
+    sc.setLocalProperty(SparkContext.SPARK_SCHEDULER_POOL, "pool1")
+    var localProperties = sc.getLocalProperties
+    schedulableBuilder.addTaskSetManager(taskSetManager0, localProperties)
+    schedulableBuilder.addTaskSetManager(taskSetManager1, localProperties)
+
+    sc.setLocalProperty(SparkContext.SPARK_SCHEDULER_POOL, "pool2")
+    localProperties = sc.getLocalProperties
+    val taskSetManager2 = createTaskSetManager(2, 2, taskScheduler)
+    val taskSetManager3 = createTaskSetManager(3, 2, taskScheduler)
+    schedulableBuilder.addTaskSetManager(taskSetManager2, localProperties)
+    schedulableBuilder.addTaskSetManager(taskSetManager3, localProperties)
+
+
+    rootPool.updateAvailableSlots(12)
+    // pool 1 weight is 1, pool 2 weight is 2, so they get 4 and 8 slots respectively
+    // since pool 1 is FIFO TSM 0 gets all 4 slots of pool 1
+    // since pool 2 is FAIR TSM 3 and 4 each get half (4) of pool 2's 8 slots
+    assert(taskSetManager0.numAvailableSlot === 4)
+    assert(taskSetManager1.numAvailableSlot === 0)
+    assert(taskSetManager2.numAvailableSlot === 4)
+    assert(taskSetManager3.numAvailableSlot === 4)
+
+    rootPool.updateAvailableSlots(6)
+    // similar to above but pool 1 has 3 slots due to having minShare of 3
+    assert(taskSetManager0.numAvailableSlot === 3)
+    assert(taskSetManager1.numAvailableSlot === 0)
+    assert(taskSetManager2.numAvailableSlot === 2)
+    assert(taskSetManager3.numAvailableSlot === 2)
+  }
+
   test("FIFO Scheduler Test") {
     sc = new SparkContext(LOCAL, APP_NAME)
     val taskScheduler = new TaskSchedulerImpl(sc)

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

@@ -195,6 +195,63 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
     assert(!failedTaskSet)
   }
 
+  test("executors should not sit idle for too long") {
+    val LOCALITY_WAIT_MS = 3000
+    val clock = new ManualClock
+    val conf = new SparkConf()
+    sc = new SparkContext("local", "TaskSchedulerImplSuite", conf)
+    val taskScheduler = new TaskSchedulerImpl(sc,
+      sc.conf.get(config.TASK_MAX_FAILURES),
+      clock = clock) {
+      override def createTaskSetManager(taskSet: TaskSet, maxTaskFailures: Int): TaskSetManager = {
+        new TaskSetManager(this, taskSet, maxTaskFailures, blacklistTrackerOpt, clock)
+      }
+      override def shuffleOffers(offers: IndexedSeq[WorkerOffer]): IndexedSeq[WorkerOffer] = {
+        // Don't shuffle the offers around for this test.  Instead, we'll just pass in all
+        // the permutations we care about directly.
+        offers
+      }
+    }
+    // Need to initialize a DAGScheduler for the taskScheduler to use for callbacks.
+    new DAGScheduler(sc, taskScheduler) {
+      override def taskStarted(task: Task[_], taskInfo: TaskInfo): Unit = {}
+
+      override def executorAdded(execId: String, host: String): Unit = {}
+    }
+    taskScheduler.initialize(new FakeSchedulerBackend)
+    val taskSet = FakeTask.createTaskSet(5, 1, 1,
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1"))
+    )
+    taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+    taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+    taskScheduler.submitTasks(taskSet)
+
+    // First offer host2, exec2: no task should be chosen due to bad data locality
+    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
+      .flatten.isEmpty)
+    val task0 = taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten.head
+    assert(task0.index === 0)
+    taskScheduler.statusUpdate(task0.taskId, TaskState.FINISHED, ByteBuffer.allocate(0))
+
+    clock.advance(LOCALITY_WAIT_MS * 2)
+    val task1 = taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten.head
+    assert(task1.index === 1)
+    val task2 = taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
+      .flatten.head
+    assert(task2.index === 2)
+    taskScheduler.statusUpdate(task1.taskId, TaskState.FINISHED, ByteBuffer.allocate(0))
+    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten.head.index === 3)
+    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
+      .flatten.isEmpty)
+  }
+
   test("Scheduler does not crash when tasks are not serializable") {
     val taskCpus = 2
     val taskScheduler = setupSchedulerWithMaster(
@@ -910,6 +967,7 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
   test("SPARK-16106 locality levels updated if executor added to existing host") {
     val taskScheduler = setupScheduler()
 
+    taskScheduler.resourceOffers(IndexedSeq(new WorkerOffer("executor0", "host0", 1)))
     taskScheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 0, stageAttemptId = 0,
       (0 until 2).map { _ => Seq(TaskLocation("host0", "executor2")) }: _*
     ))