spark 笔记 14: spark中的delay scheduling实现

// Figure out which locality levels we have in our TaskSet, so we can do delay scheduling
var myLocalityLevels = computeValidLocalityLevels()
var localityWaits = myLocalityLevels.map(getLocalityWait) // Time to wait at each level

// 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.
// We then move down if we manage to launch a "more local" task.
var currentLocalityIndex = 0    // Index of our current locality level in validLocalityLevels

// Set of pending tasks for each executor. These collections are actually
// treated as stacks, in which new tasks are added to the end of the
// ArrayBuffer and removed from the end. This makes it faster to detect
// tasks that repeatedly fail because whenever a task failed, it is put
// back at the head of the stack. They are also only cleaned up lazily;
// when a task is launched, it remains in all the pending lists except
// the one that it was launched from, but gets removed from them later.
private val pendingTasksForExecutor = new HashMap[String, ArrayBuffer[Int]]

// Set of pending tasks for each host. Similar to pendingTasksForExecutor,
// but at host level.
private val pendingTasksForHost = new HashMap[String, ArrayBuffer[Int]]

// Set of pending tasks for each rack -- similar to the above.
private val pendingTasksForRack = new HashMap[String, ArrayBuffer[Int]]

// Set containing pending tasks with no locality preferences.
var pendingTasksWithNoPrefs = new ArrayBuffer[Int]

var lastLaunchTime = clock.getTime()  // Time we last launched a task at this level/**
 * Compute the locality levels used in this TaskSet. Assumes that all tasks have already been
 * added to queues using addPendingTask.
 *
 */
private def computeValidLocalityLevels(): Array[TaskLocality.TaskLocality] = {
  import TaskLocality.{PROCESS_LOCAL, NODE_LOCAL, NO_PREF, RACK_LOCAL, ANY}
  val levels = new ArrayBuffer[TaskLocality.TaskLocality]
  if (!pendingTasksForExecutor.isEmpty && getLocalityWait(PROCESS_LOCAL) != 0 &&
      pendingTasksForExecutor.keySet.exists(sched.isExecutorAlive(_))) {
    levels += PROCESS_LOCAL
  }
  if (!pendingTasksForHost.isEmpty && getLocalityWait(NODE_LOCAL) != 0 &&
      pendingTasksForHost.keySet.exists(sched.hasExecutorsAliveOnHost(_))) {
    levels += NODE_LOCAL
  }
  if (!pendingTasksWithNoPrefs.isEmpty) {
    levels += NO_PREF
  }
  if (!pendingTasksForRack.isEmpty && getLocalityWait(RACK_LOCAL) != 0 &&
      pendingTasksForRack.keySet.exists(sched.hasHostAliveOnRack(_))) {
    levels += RACK_LOCAL
  }
  levels += ANY
  logDebug("Valid locality levels for " + taskSet + ": " + levels.mkString(", "))
  levels.toArray
}

private def getLocalityWait(level: TaskLocality.TaskLocality): Long = {
  val defaultWait = conf.get("spark.locality.wait", "3000")
  level match {
    case TaskLocality.PROCESS_LOCAL =>
      conf.get("spark.locality.wait.process", defaultWait).toLong
    case TaskLocality.NODE_LOCAL =>
      conf.get("spark.locality.wait.node", defaultWait).toLong
    case TaskLocality.RACK_LOCAL =>
      conf.get("spark.locality.wait.rack", defaultWait).toLong
    case _ => 0L
  }
}

@DeveloperApi
object TaskLocality extends Enumeration {
  // Process local is expected to be used ONLY within TaskSetManager for now.
  val PROCESS_LOCAL, NODE_LOCAL, NO_PREF, RACK_LOCAL, ANY = Value

  type TaskLocality = Value

  def isAllowed(constraint: TaskLocality, condition: TaskLocality): Boolean = {
    condition <= constraint
  }
}

/**
 * Find the index in myLocalityLevels for a given locality. This is also designed to work with
 * localities that are not in myLocalityLevels (in case we somehow get those) by returning the
 * next-biggest level we have. Uses the fact that the last value in myLocalityLevels is ANY.
 */
def getLocalityIndex(locality: TaskLocality.TaskLocality): Int = {
  var index = 0
  while (locality > myLocalityLevels(index)) {
    index += 1
  }
  index
}

/**
 * Get the level we can launch tasks according to delay scheduling, based on current wait time.
 */
private def getAllowedLocalityLevel(curTime: Long): TaskLocality.TaskLocality = {
  while (curTime - lastLaunchTime >= localityWaits(currentLocalityIndex) &&
      currentLocalityIndex < myLocalityLevels.length - 1)
  {
    // Jump to the next locality level, and remove our waiting time for the current one since
    // we don‘t want to count it again on the next one
    lastLaunchTime += localityWaits(currentLocalityIndex)
    currentLocalityIndex += 1
  }
  myLocalityLevels(currentLocalityIndex)
}

def recomputeLocality() {
  val previousLocalityLevel = myLocalityLevels(currentLocalityIndex)
  myLocalityLevels = computeValidLocalityLevels()
  localityWaits = myLocalityLevels.map(getLocalityWait)
  currentLocalityIndex = getLocalityIndex(previousLocalityLevel)
}

/**
 * Compute the locality levels used in this TaskSet. Assumes that all tasks have already been
 * added to queues using addPendingTask.
 *
 */
private def computeValidLocalityLevels(): Array[TaskLocality.TaskLocality] = {
  import TaskLocality.{PROCESS_LOCAL, NODE_LOCAL, NO_PREF, RACK_LOCAL, ANY}
  val levels = new ArrayBuffer[TaskLocality.TaskLocality]
  if (!pendingTasksForExecutor.isEmpty && getLocalityWait(PROCESS_LOCAL) != 0 &&
      pendingTasksForExecutor.keySet.exists(sched.isExecutorAlive(_))) {
    levels += PROCESS_LOCAL
  }
  if (!pendingTasksForHost.isEmpty && getLocalityWait(NODE_LOCAL) != 0 &&
      pendingTasksForHost.keySet.exists(sched.hasExecutorsAliveOnHost(_))) {
    levels += NODE_LOCAL
  }
  if (!pendingTasksWithNoPrefs.isEmpty) {
    levels += NO_PREF
  }
  if (!pendingTasksForRack.isEmpty && getLocalityWait(RACK_LOCAL) != 0 &&
      pendingTasksForRack.keySet.exists(sched.hasHostAliveOnRack(_))) {
    levels += RACK_LOCAL
  }
  levels += ANY
  logDebug("Valid locality levels for " + taskSet + ": " + levels.mkString(", "))
  levels.toArray
}

/**
 * Dequeue a pending task for a given node and return its index and locality level.
 * Only search for tasks matching the given locality constraint.
 *
 * @return An option containing (task index within the task set, locality, is speculative?)
 */
private def findTask(execId: String, host: String, maxLocality: TaskLocality.Value)
  : Option[(Int, TaskLocality.Value, Boolean)] =
{
  for (index <- findTaskFromList(execId, getPendingTasksForExecutor(execId))) {
    return Some((index, TaskLocality.PROCESS_LOCAL, false))
  }
。。。
  // find a speculative task if all others tasks have been scheduled
  findSpeculativeTask(execId, host, maxLocality).map {
    case (taskIndex, allowedLocality) => (taskIndex, allowedLocality, true)}
}