DAGScheduler学习笔记

一、stage划分算法

先介绍stage的划分算法，会从出发action操作的那个rdd开始往前倒推，首先会为最后一个rdd创建一个stage，然后往前倒推，如果发现某个rdd是宽依赖，那么就会将宽依赖的那个rdd创建一个新的stage，那个rdd就是新的stage的最后一个rdd，然后依次类推，根据宽依赖和窄依赖，进行stage划分，直至所有rdd遍历完。

二、stage划分算法源码分析（非常重要）

从上一编blog，job的触发流程最后一步就是调用DAGScheduler的runJob方法，runjob方法中又调用了submitJob方法

/**
 * Submit a job to the job scheduler and get a JobWaiter object back. The JobWaiter object
 * can be used to block until the the job finishes executing or can be used to cancel the job.
 */
def submitJob[T, U](
    rdd: RDD[T],
    func: (TaskContext, Iterator[T]) => U,
    partitions: Seq[Int],
    callSite: CallSite,
    resultHandler: (Int, U) => Unit,
    properties: Properties): JobWaiter[U] = {
  // Check to make sure we are not launching a task on a partition that does not exist.
  val maxPartitions = rdd.partitions.length
  partitions.find(p => p >= maxPartitions || p < 0).foreach { p =>
    throw new IllegalArgumentException(
      "Attempting to access a non-existent partition: " + p + ". " +
        "Total number of partitions: " + maxPartitions)
  }

  val jobId = nextJobId.getAndIncrement()
  if (partitions.size == 0) {
    return new JobWaiter[U](this, jobId, 0, resultHandler)
  }

  assert(partitions.size > 0)
  val func2 = func.asInstanceOf[(TaskContext, Iterator[_]) => _]
  val waiter = new JobWaiter(this, jobId, partitions.size, resultHandler)
  eventProcessLoop.post(JobSubmitted(
    jobId, rdd, func2, partitions.toArray, callSite, waiter,
    SerializationUtils.clone(properties)))
  waiter
}

这段代码中有个很重要的东西，eventProcessLoop（DAGSchedulerEventProcessLoop是一个DAGScheduler的内部类），调用了post方法发送了JobSubmitted消息。从源码中可以看到，接收到消息后，调用了dagScheduler的handleJobSubmitted方法。这个方法是DAGScheduler的job调度的核心入口

private[scheduler] def handleJobSubmitted(jobId: Int,
    finalRDD: RDD[_],
    func: (TaskContext, Iterator[_]) => _,
    partitions: Array[Int],
    callSite: CallSite,
    listener: JobListener,
    properties: Properties) {
  var finalStage: ResultStage = null
  try {
    // New stage creation may throw an exception if, for example, jobs are run on a
    // HadoopRDD whose underlying HDFS files have been deleted.
    //第一步、使用触发job的最后一个RDD，创建finalStage，这个方法就是简单的创建了一个stage，
    //并且将stage加入到DAGScheduler的缓存（stage中有个重要的变量isShuffleMap）
    finalStage = newResultStage(finalRDD, partitions.length, jobId, callSite)
  } catch {
    case e: Exception =>
      logWarning("Creating new stage failed due to exception - job: " + jobId, e)
      listener.jobFailed(e)
      return
  }    
  if (finalStage != null) {
    //第二步、用finalStage创建一个job
    val job = new ActiveJob(jobId, finalStage, func, partitions, callSite, listener, properties)
    clearCacheLocs()
    logInfo("Got job %s (%s) with %d output partitions".format(
      job.jobId, callSite.shortForm, partitions.length))
    logInfo("Final stage: " + finalStage + "(" + finalStage.name + ")")
    logInfo("Parents of final stage: " + finalStage.parents)
    logInfo("Missing parents: " + getMissingParentStages(finalStage))
    val jobSubmissionTime = clock.getTimeMillis()
    jobIdToActiveJob(jobId) = job
    //第三步、将job加入到内存缓存中
    activeJobs += job
    finalStage.resultOfJob = Some(job)
    val stageIds = jobIdToStageIds(jobId).toArray
    val stageInfos = stageIds.flatMap(id => stageIdToStage.get(id).map(_.latestInfo))
    listenerBus.post(
      SparkListenerJobStart(job.jobId, jobSubmissionTime, stageInfos, properties))
    //第四步、（很关键）使用submitStage方法提交finalStage
    //这个方法会导致第一个stage提交，其他的stage放入waitingStages队列，使用递归优先提交父stage
    submitStage(finalStage)
  }
  //提交等待的stage队列
  submitWaitingStages()
}

接下来看下第四步调用的submitStage方法，这个是stage划分算法的入口，但是stage划分算法是有submitStage和getMissingParentStages方法共同组成的。

private def submitStage(stage: Stage) {
  val jobId = activeJobForStage(stage)
  if (jobId.isDefined) {
    logDebug("submitStage(" + stage + ")")
    if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {
      //很关键的一行，调用getMissingParentStage方法去获取这个stage的父stage
      val missing = getMissingParentStages(stage).sortBy(_.id)
      logDebug("missing: " + missing)
      //其实这里会循环递归调用，直到最初的stage没有父stage，其余的stage被放在waitingMissingStages
      if (missing.isEmpty) {
        logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
        //这个就是提交stage的方法。后面再分享
        submitMissingTasks(stage, jobId.get)
      } else {
        //如果不为空，就是有父Stage，递归调用submitStage方法去提交父Stage，这里是stage划分算法的精髓。
        for (parent <- missing) {
          submitStage(parent)
        }
        //并且将当前stage，放入等待执行的stage队列中
        waitingStages += stage
      }
    }
  } else {
    abortStage(stage, "No active job for stage " + stage.id, None)
  }
}

这里再来看下getMissingParentStage方法

private def getMissingParentStages(stage: Stage): List[Stage] = {
  val missing = new HashSet[Stage]
  val visited = new HashSet[RDD[_]]
  // We are manually maintaining a stack here to prevent StackOverflowError
  // caused by recursively visiting 先入后出
  val waitingForVisit = new Stack[RDD[_]]
  //定义visit方法，供后面代码中stage的RDD循环调用
  def visit(rdd: RDD[_]) {
    if (!visited(rdd)) {
      visited += rdd
      val rddHasUncachedPartitions = getCacheLocs(rdd).contains(Nil)
      if (rddHasUncachedPartitions) {
        //遍历RDD
        for (dep <- rdd.dependencies) {
          dep match {
            //如果是宽依赖，使用宽依赖的RDD创建一个新的stage，并且会把isShuffleMap变量设置为true
            //默认只有最后一个stage不是ShuffleMap stage
            case shufDep: ShuffleDependency[_, _, _] =>
              val mapStage = getShuffleMapStage(shufDep, stage.firstJobId)
              if (!mapStage.isAvailable) {
                //把stage放到缓存中
                missing += mapStage
              }
            //如果是窄依赖，就把rdd加入到stack中，虽然循环时调用了stack的pop方法，但是这里又push了一个进去。
            case narrowDep: NarrowDependency[_] =>
              waitingForVisit.push(narrowDep.rdd)
          }
        }
      }
    }
  }
  //首先，往stack中推入stage的最后一个rdd
  waitingForVisit.push(stage.rdd)
  //循环
  while (waitingForVisit.nonEmpty) {
    //对stage的最后一个rdd，调用自己内部定义方法（就是上面的visit方法），注意这里stack的pop方法取出rdd
    visit(waitingForVisit.pop())
  }
  //立刻返回新的stage
  missing.toList
}

三、提交stage的方法--submitsMissingTasks

这个源码有点长，就不在这里展示了。

这个方法的作用就是为stage创建一批task，task的数量和partition的数量相同

获取要partition的数量，将stage加入runningStage队列中

这里涉及到一个Task最佳位置的算法，就是调用getpreferredlocs方法。

原理就是：从stage的最后一个rdd开始，去找哪个rdd的partition被cache或checkPoint，那么Task的最佳位置就是cache或checkPoint的位置。因为Task在这个节点上就不用计算之前的RDD。如果没有，就有TaskScheduler来决定到哪个节点上运行。

最后，针对stage的task，创建TaskSet对象，调用TaskScheduler的submitTasks方法，提交TaskSet。