diff --git a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala index 80805df256a15..6f17c27b08792 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala @@ -17,6 +17,7 @@ package org.apache.spark.scheduler +import java.util import java.util.concurrent.{ConcurrentHashMap, ConcurrentLinkedQueue} import scala.collection.JavaConverters._ @@ -119,4 +120,72 @@ private[spark] class Pool( parent.decreaseRunningTasks(taskNum) } } + + // Update the number of slots considered available for each TaskSetManager whose ancestor + // in the tree is this pool + // For FAIR scheduling, slots are distributed among pools based on weights and minshare. + // If a pool requires fewer slots than are available to it, the leftover slots are redistributed + // to the remaining pools using the remaining pools' weights. + // For FIFO scheduling, the schedulable queue is iterated over in FIFO order, + // giving each schedulable the remaining slots, + // up to the number of remaining tasks for that schedulable. + override def updateAvailableSlots(numSlots: Float): Unit = { + schedulingMode match { + case SchedulingMode.FAIR => + val queueCopy = new util.LinkedList[Schedulable](schedulableQueue) + var shouldRedistribute = true + var totalWeights = schedulableQueue.asScala.map(_.weight).sum + var totalSlots = numSlots + while (totalSlots > 0 && shouldRedistribute) { + shouldRedistribute = false + var nextWeights = totalWeights + var nextSlots = totalSlots + val iterator = queueCopy.iterator() + while (iterator.hasNext) { + val schedulable = iterator.next() + val numTasksRemaining = schedulable.getSortedTaskSetQueue + .map(tsm => tsm.tasks.length - tsm.tasksSuccessful).sum + val allocatedSlots = Math.max( + totalSlots * schedulable.weight / totalWeights, + schedulable.minShare) + if (numTasksRemaining < allocatedSlots) { + schedulable.updateAvailableSlots(numTasksRemaining) + nextWeights -= schedulable.weight + nextSlots -= numTasksRemaining + shouldRedistribute = true + iterator.remove() + } + } + totalWeights = nextWeights + totalSlots = nextSlots + } + + // All schedulables remaining have more or equal remaining tasks than their share of slots, + // so no need to recalculate remaining tasks. Just give them their total share of slots. + queueCopy.forEach(schedulable => { + val allocatedSlots = Math.max( + totalSlots * schedulable.weight / totalWeights, + schedulable.minShare) + schedulable.updateAvailableSlots(allocatedSlots) + }) + case SchedulingMode.FIFO => + val sortedSchedulableQueue = + schedulableQueue.asScala.toSeq.sortWith(taskSetSchedulingAlgorithm.comparator) + var remainingSlots = numSlots + for (schedulable <- sortedSchedulableQueue) { + if (remainingSlots == 0) { + schedulable.updateAvailableSlots(0) + } else { + val numTasksRemaining = schedulable.getSortedTaskSetQueue + .map(tsm => tsm.tasks.length - tsm.tasksSuccessful).sum + val toAssign = Math.min(numTasksRemaining, remainingSlots) + schedulable.updateAvailableSlots(toAssign) + remainingSlots -= toAssign + } + } + case _ => + val msg = s"Unsupported scheduling mode: $schedulingMode. Use FAIR or FIFO instead." + throw new IllegalArgumentException(msg) + } + } } diff --git a/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala b/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala index b6f88ed0a93aa..ed4b2fa4deaf0 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala +++ b/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 diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala index f25a36c7af22a..9074500733bec 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala +++ b/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,9 @@ private[spark] class TaskSchedulerImpl( if (!executorIdToRunningTaskIds.contains(o.executorId)) { hostToExecutors(o.host) += o.executorId executorAdded(o.executorId, o.host) + // Assumes the first offer will include all cores (free cores == all cores) + 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 +445,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 +816,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) diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala index 5c0bc497dd1b3..c46224c6eef09 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/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,14 @@ 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 "available" slots. + // See Pool.updateAvailableSlots() for how the number of available slots is calculated. + // 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 +415,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 +541,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 +1061,10 @@ private[spark] class TaskSetManager( def executorAdded(): Unit = { recomputeLocality() } + + override def updateAvailableSlots(numSlots: Float): Unit = { + numAvailableSlot = numSlots + } } private[spark] object TaskSetManager { diff --git a/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala index b953add9d58cb..4693e9d742356 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala @@ -51,6 +51,122 @@ 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, 5, taskScheduler) + val taskSetManager1 = createTaskSetManager(1, 3, taskScheduler) + val taskSetManager2 = createTaskSetManager(2, 2, taskScheduler) + schedulableBuilder.addTaskSetManager(taskSetManager0, null) + schedulableBuilder.addTaskSetManager(taskSetManager1, null) + schedulableBuilder.addTaskSetManager(taskSetManager2, null) + + rootPool.updateAvailableSlots(9) + // first task set only gets 5 slots as there are only 5 tasks + assert(taskSetManager0.numAvailableSlot === 5) + assert(taskSetManager1.numAvailableSlot === 3) + // last task set gets only 1 slot since there are 9 total and 8 have already been taken + assert(taskSetManager2.numAvailableSlot === 1) + + rootPool.removeSchedulable(taskSetManager1) + rootPool.updateAvailableSlots(6) + assert(taskSetManager0.numAvailableSlot === 5) + // TSM 2 gets only 1 slot since there are 6 total and 5 have been taken by TSM 0 + assert(taskSetManager2.numAvailableSlot === 1) + + val taskSetManager3 = createTaskSetManager(3, 3, taskScheduler) + schedulableBuilder.addTaskSetManager(taskSetManager3, null) + taskSetManager0.tasksSuccessful += 1 + taskSetManager0.addRunningTask(1) + rootPool.updateAvailableSlots(8) + // Gets only 4 slots since it has only 4 tasks remaining (1 complete) + assert(taskSetManager0.numAvailableSlot === 4) + assert(taskSetManager2.numAvailableSlot === 2) + // New TSM 3 gets single leftover slot: 8 - 4 (TSM 0) - 2 (TSM 2) = 2 + assert(taskSetManager3.numAvailableSlot === 2) + } + + 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, 4, taskScheduler) + taskSetManager0.tasksSuccessful += 1 + taskSetManager0.addRunningTask(1) + 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, 9, taskScheduler) + val taskSetManager3 = createTaskSetManager(3, 2, taskScheduler) + taskSetManager3.tasksSuccessful += 1 + 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 + // TSM 0 has 4 tasks total, but 1 is complete, so it needs only 3 slots. + assert(taskSetManager0.numAvailableSlot === 3) + // TSM 1 takes the remaining slot from pool 1. + assert(taskSetManager1.numAvailableSlot === 1) + // Since pool 2 is FAIR TSM 2 and 3 each could get half (4) of pool 2's 8 slots + // TSM 3 has only 1 incomplete task, so it takes only 1 out of 4 possible slots + assert(taskSetManager3.numAvailableSlot === 1) + // TSM 2 takes its original 4 slots + the 3 remaining from TSM 3 + assert(taskSetManager2.numAvailableSlot === 7) + + + 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 === 3) + assert(taskSetManager3.numAvailableSlot === 1) + + taskSetManager3.tasksSuccessful -= 1 + // same as above, but task 3 should use 2 slots since it has 0 complete tasks + rootPool.updateAvailableSlots(6) + assert(taskSetManager0.numAvailableSlot === 3) + assert(taskSetManager1.numAvailableSlot === 0) + assert(taskSetManager2.numAvailableSlot === 2) + assert(taskSetManager3.numAvailableSlot === 2) + + rootPool.getSchedulableByName("pool2").removeSchedulable(taskSetManager2) + rootPool.updateAvailableSlots(6) + // 4 slots go to pool 2, but TSM 3 only takes 2, remaining 2 can go to pool 1 + // Pool one has 2 slots + 2 from pool 2 = 4 + // TSM 0 takes 3 and TSM 1 takes 1 + assert(taskSetManager0.numAvailableSlot === 3) + assert(taskSetManager1.numAvailableSlot === 1) + assert(taskSetManager3.numAvailableSlot === 2) + + + // TSM 4 goes to pool 2 + val taskSetManager4 = createTaskSetManager(4, 1, taskScheduler) + schedulableBuilder.addTaskSetManager(taskSetManager4, localProperties) + rootPool.updateAvailableSlots(6) + // Same as above but new TSM 4 in pool 2 takes 1 slot, giving only 1 to pool 1 + assert(taskSetManager0.numAvailableSlot === 3) + assert(taskSetManager1.numAvailableSlot === 0) + assert(taskSetManager3.numAvailableSlot === 2) + assert(taskSetManager4.numAvailableSlot === 1) + } + test("FIFO Scheduler Test") { sc = new SparkContext(LOCAL, APP_NAME) val taskScheduler = new TaskSchedulerImpl(sc) diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala index e7ecf847ff4f4..61cd597ed413b 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala @@ -195,6 +195,172 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B assert(!failedTaskSet) } + test("SPARK-18886 Delay scheduling should not delay some executors indefinitely " + + "if one task is scheduled before delay timeout") { + 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 = {} + } + val valueSer = SparkEnv.get.serializer.newInstance() + taskScheduler.initialize(new FakeSchedulerBackend) + val taskSet = FakeTask.createTaskSet(9, 1, 1, + Seq(TaskLocation("host1", "exec1")), + Seq(TaskLocation("host1", "exec1")), + Seq(TaskLocation("host1", "exec1")), + Seq(TaskLocation("host1", "exec1")), + 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) + + // Next offer exec 1, which is local to all tasks. First task will take it + val taskDescriptions0 = taskScheduler + .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1))) + .flatten + assert(taskDescriptions0.length === 1) + val task0 = taskDescriptions0.head + assert(task0.index === 0) + val result0 = new DirectTaskResult[Int](valueSer.serialize(0), Seq(), Array()) + taskScheduler.statusUpdate(task0.taskId, TaskState.FINISHED, valueSer.serialize(result0)) + + // clock advances, increasing data locality level to ANY + clock.advance(LOCALITY_WAIT_MS * 2) + // Local resource is utilized again + val taskDescriptions1 = taskScheduler + .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1))) + .flatten + assert(taskDescriptions1.length === 1) + val task1 = taskDescriptions1.head + assert(task1.index === 1) + + // Even though we launched a local task, we still utilize non-local exec2 + // This is the behavior change to fix SPARK-18886. Also tested below in this same test + // This is because we have not utilized our 2 possible slots (exec1 + exec2) within + // the locality wait period. + val taskDescriptions2 = taskScheduler + .resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1))) + .flatten + assert(taskDescriptions2.length === 1) + val task2 = taskDescriptions2.head + assert(task2.index === 2) + + // Finish data local task on exec 1 (still 1 task running on exec 2) + val result1 = new DirectTaskResult[Int](valueSer.serialize(1), Seq(), Array()) + taskScheduler.statusUpdate(task1.taskId, TaskState.FINISHED, valueSer.serialize(result1)) + + // Local resource will again be utilized. Data locality is reset to PROCESS_LOCAL + val taskDescriptions3 = taskScheduler + .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1))) + .flatten + assert(taskDescriptions3.length === 1) + val task3 = taskDescriptions3.head + assert(task3.index === 3) + // Complete task 2 and task 3, no more tasks running + val result2 = new DirectTaskResult[Int](valueSer.serialize(2), Seq(), Array()) + taskScheduler.statusUpdate(task2.taskId, TaskState.FINISHED, valueSer.serialize(result2)) + val result3 = new DirectTaskResult[Int](valueSer.serialize(3), Seq(), Array()) + taskScheduler.statusUpdate(task3.taskId, TaskState.FINISHED, valueSer.serialize(result3)) + + // Non-local resource will not be utilized, since data locality was reset, + // and locality wait has not expired + assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1))) + .flatten.isEmpty) + + // Offer a total of 2 slots on exec 3, non are taken since non-local + assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec3", "host3", 2))) + .flatten.isEmpty) + + taskScheduler.executorLost("exec1", LossReasonPending) + taskScheduler.executorLost("exec2", LossReasonPending) + + // clock advances, increasing data locality level to ANY + clock.advance(LOCALITY_WAIT_MS * 2) + + // Accept non local resource, total utilized slots is 1 + val taskDescriptions4 = taskScheduler + .resourceOffers(IndexedSeq(WorkerOffer("exec3", "host3", 1))) + .flatten + assert(taskDescriptions4.length === 1) + val task4 = taskDescriptions4.head + assert(task4.index === 4) + + + // Offer data local exec 1. This increases total slots to 3, but only 2 are utilized now + val taskDescriptions5 = taskScheduler + .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1))) + .flatten + assert(taskDescriptions5.length === 1) + val task5 = taskDescriptions5.head + assert(task5.index === 5) + + // Data locality has not been reset, still accept non local resource. + // Now 3 tasks are running. 3 slots are being used. + // Locality timer is reset, but locality level is not + val taskDescriptions6 = taskScheduler + .resourceOffers(IndexedSeq(WorkerOffer("exec3", "host3", 1))) + .flatten + assert(taskDescriptions6.length === 1) + val task6 = taskDescriptions6.head + assert(task6.index === 6) + + // Finish data local task. Only 2 running tasks now + val result5 = new DirectTaskResult[Int](valueSer.serialize(5), Seq(), Array()) + taskScheduler.statusUpdate(task5.taskId, TaskState.FINISHED, valueSer.serialize(result5)) + + // Offer data local exec 1 again. 3 slots used and is data local, + // so locality is reset to PROCESS_LOCAL and timer is reset + val taskDescriptions7 = taskScheduler + .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1))) + .flatten + assert(taskDescriptions7.length === 1) + val task7 = taskDescriptions7.head + assert(task7.index === 7) + + // Non-local resource will not be utilized, since data locality was reset, + // and locality wait has not expired + assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec3", "host3", 1))) + .flatten.isEmpty) + + // Complete task on data local node (exec1) + val result7 = new DirectTaskResult[Int](valueSer.serialize(7), Seq(), Array()) + taskScheduler.statusUpdate(task7.taskId, TaskState.FINISHED, valueSer.serialize(result7)) + + // Complete the last task locally + val taskDescriptions8 = taskScheduler + .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1))) + .flatten + assert(taskDescriptions8.length === 1) + val task8 = taskDescriptions8.head + assert(task8.index === 8) + } + test("Scheduler does not crash when tasks are not serializable") { val taskCpus = 2 val taskScheduler = setupSchedulerWithMaster( @@ -910,6 +1076,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")) }: _* ))