源码-DAGScheduler及Stage划分提交

DAGScheduler

DAGScheduler的主要任务是基于Stage构建DAG，决定每个任务的最佳位置

• 记录哪个RDD或者Stage输出被物化
• 面向stage的调度层，为job生成以stage组成的DAG，提交TaskSet给TaskScheduler执行
• 重新提交shuffle输出丢失的stage

每一个Stage内，都是独立的tasks，他们共同执行同一个computefunction，享有相同的shuffledependencies。DAG在切分stage的时候是依照出现shuffle为界限的。

DAGScheduler实例化

下面的代码是SparkContext实例化DAGScheduler的过程：

  @volatile private[spark] var dagScheduler: DAGScheduler = _
  try {
    dagScheduler = new DAGScheduler(this)
  } catch {
    case e: Exception => {
      try {
        stop()
      } finally {
        throw new SparkException("Error while constructing DAGScheduler", e)
      }
    }
  }

下面代码显示了DAGScheduler的构造函数定义中，通过绑定TaskScheduler的方式创建，其中次构造函数去调用主构造函数来将sc的字段填充入参：

private[spark]
class DAGScheduler(
    private[scheduler] val sc: SparkContext,
    private[scheduler] val taskScheduler: TaskScheduler,
    listenerBus: LiveListenerBus,
    mapOutputTracker: MapOutputTrackerMaster,
    blockManagerMaster: BlockManagerMaster,
    env: SparkEnv,
    clock: Clock = new SystemClock())
  extends Logging {

  def this(sc: SparkContext, taskScheduler: TaskScheduler) = {
    this(
      sc,
      taskScheduler,
      sc.listenerBus,
      sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster],
      sc.env.blockManager.master,
      sc.env)
  }

  def this(sc: SparkContext) = this(sc, sc.taskScheduler)

作业提交与DAGScheduler操作

Action的大部分操作会进行作业（job）的提交，源码1.0版的job提交过程的大致调用链是：sc.runJob()–>dagScheduler.runJob–>dagScheduler.submitJob—>dagSchedulerEventProcessActor.JobSubmitted–>dagScheduler.handleJobSubmitted–>dagScheduler.submitStage–>dagScheduler.submitMissingTasks–>taskScheduler.submitTasks。
具体的作业提交执行期的函数调用为：

1. sc.runJob->dagScheduler.runJob->submitJob
2. DAGScheduler::submitJob会创建JobSummitted的event发送给内嵌类eventProcessActor（在源码1.4中，submitJob函数中，使用DAGSchedulerEventProcessLoop类进行事件的处理）
3. eventProcessActor在接收到JobSubmmitted之后调用processEvent处理函数
4. job到stage的转换，生成finalStage并提交运行，关键是调用submitStage
5. 在submitStage中会计算stage之间的依赖关系，依赖关系分为宽依赖和窄依赖两种
6. 如果计算中发现当前的stage没有任何依赖或者所有的依赖都已经准备完毕，则提交task
7. 提交task是调用函数submitMissingTasks来完成
8. task真正运行在哪个worker上面是由TaskScheduler来管理，也就是上面的submitMissingTasks会调用TaskScheduler::submitTasks
9. TaskSchedulerImpl中会根据Spark的当前运行模式来创建相应的backend,如果是在单机运行则创建LocalBackend
10. LocalBackend收到TaskSchedulerImpl传递进来的ReceiveOffers事件
11. receiveOffers->executor.launchTask->TaskRunner.run

DAGScheduler的runJob函数

DAGScheduler.runjob最后把结果通过resultHandler保存返回。
这里DAGScheduler的runJob函数调用DAGScheduler的submitJob函数来提交任务：

  def runJob[T, U: ClassTag](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int],
      callSite: CallSite,
      allowLocal: Boolean,
      resultHandler: (Int, U) => Unit,
      properties: Properties): Unit = {
    val start = System.nanoTime
    val waiter = submitJob(rdd, func, partitions, callSite, allowLocal, resultHandler, properties)
    waiter.awaitResult() match {
      case JobSucceeded => {
        logInfo("Job %d finished: %s, took %f s".format
          (waiter.jobId, callSite.shortForm, (System.nanoTime - start) / 1e9))
      }
      case JobFailed(exception: Exception) =>
        logInfo("Job %d failed: %s, took %f s".format
          (waiter.jobId, callSite.shortForm, (System.nanoTime - start) / 1e9))
        throw exception
    }
  }

作业提交的调度

在Spark源码1.4.0中，DAGScheduler的submitJob函数不再使用DAGEventProcessActor进行事件处理和消息通信，而是使用DAGSchedulerEventProcessLoop类实例eventProcessLoop进行JobSubmitted事件的post动作。
下面是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,
      allowLocal: Boolean,
      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, allowLocal, callSite, waiter, properties))
    waiter
  }

当eventProcessLoop对象投递了JobSubmitted事件之后，对象内的eventThread线程实例对事件进行处理，不断从事件队列中取出事件，调用onReceive函数处理事件，当匹配到JobSubmitted事件后，调用DAGScheduler的handleJobSubmitted函数并传入jobid、rdd等参数来处理Job。

handleJobSubmitted函数

Job处理过程中handleJobSubmitted比较关键，该函数主要负责RDD的依赖性分析，生成finalStage，并根据finalStage来产生ActiveJob。
在handleJobSubmitted函数源码中，给出了部分注释：

  private[scheduler] def handleJobSubmitted(jobId: Int,
      finalRDD: RDD[_],
      func: (TaskContext, Iterator[_]) => _,
      partitions: Array[Int],
      allowLocal: Boolean,
      callSite: CallSite,
      listener: JobListener,
      properties: Properties) {
    var finalStage: Stage = 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.
      finalStage = newStage(finalRDD, partitions.size, None, jobId, callSite)
    } catch {
      //错误处理，告诉监听器作业失败，返回....
      case e: Exception =>
        logWarning("Creating new stage failed due to exception - job: " + jobId, e)
        listener.jobFailed(e)
        return
    }
    if (finalStage != null) {
      val job = new ActiveJob(jobId, finalStage, func, partitions, callSite, listener, properties)
      clearCacheLocs()
      logInfo("Got job %s (%s) with %d output partitions (allowLocal=%s)".format(
        job.jobId, callSite.shortForm, partitions.length, allowLocal))
      logInfo("Final stage: " + finalStage + "(" + finalStage.name + ")")
      logInfo("Parents of final stage: " + finalStage.parents)
      logInfo("Missing parents: " + getMissingParentStages(finalStage))
      val shouldRunLocally =
        localExecutionEnabled && allowLocal && finalStage.parents.isEmpty && partitions.length == 1
      val jobSubmissionTime = clock.getTimeMillis()
      if (shouldRunLocally) {
        // 很短、没有父stage的本地操作，比如 first() or take() 的操作本地执行
        // Compute very short actions like first() or take() with no parent stages locally.
        listenerBus.post(
          SparkListenerJobStart(job.jobId, jobSubmissionTime, Seq.empty, properties))
        runLocally(job)
      } else {
        // collect等操作走的是这个过程，更新相关的关系映射，用监听器监听，然后提交作业
        jobIdToActiveJob(jobId) = 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))
        // 提交stage
        submitStage(finalStage)
      }
    }
    // 提交stage
    submitWaitingStages()
  }

org.apache.spark.scheduler.DAGScheduler#handleJobSubmitted首先会根据RDD创建finalStage。finalStage，顾名思义，就是最后的那个Stage。然后创建job，最后提交。提交的job如果满足一下条件，那么它将以本地模式运行：

1）spark.localExecution.enabled设置为true  并且 2）用户程序显式指定可以本地运行 并且 3）finalStage的没有父Stage 并且 4）仅有一个partition

3）和 4）的话主要为了任务可以快速执行；如果有多个stage或者多个partition的话，本地运行可能会因为本机的计算资源的问题而影响任务的计算速度。

要理解什么是Stage，首先要搞明白什么是Task。Task是在集群上运行的基本单位。一个Task负责处理RDD的一个partition。RDD的多个patition会分别由不同的Task去处理。当然了这些Task的处理逻辑完全是一致的。这一组Task就组成了一个Stage。有两种Task：

1.  org.apache.spark.scheduler.ShuffleMapTask
2.  org.apache.spark.scheduler.ResultTask

ShuffleMapTask根据Task的partitioner将计算结果放到不同的bucket中。而ResultTask将计算结果发送回Driver Application。一个Job包含了多个Stage，而Stage是由一组完全相同的Task组成的。最后的Stage包含了一组ResultTask。

在用户触发了一个action后，比如count，collect，SparkContext会通过runJob的函数开始进行任务提交。最后会通过DAG的event processor 传递到DAGScheduler本身的handleJobSubmitted，它首先会划分Stage，提交Stage，提交Task。至此，Task就开始在运行在集群上了。

一个Stage的开始就是从外部存储或者shuffle结果中读取数据；一个Stage的结束就是由于发生shuffle或者生成结果时。

创建finalStage

handleJobSubmitted 通过调用newStage来创建finalStage：

1. finalStage = newStage(finalRDD, partitions.size, None, jobId, callSite)  

创建一个result stage，或者说finalStage，是通过调用org.apache.spark.scheduler.DAGScheduler#newStage完成的；而创建一个shuffle stage，需要通过调用org.apache.spark.scheduler.DAGScheduler#newOrUsedStage（1.6叫newOrUsedShuffleStage）。 

1. private def newStage(  //1.6版本叫newResultStage
2.       rdd: RDD[_],  
3.       numTasks: Int,  
4.       shuffleDep: Option[ShuffleDependency[_, _, _]],  
5.       jobId: Int,  
6.       callSite: CallSite)  
7.     : Stage =  
8.   {  
9.     val id = nextStageId.getAndIncrement()  
10.     val stage =  
11.       new Stage(id, rdd, numTasks, shuffleDep, getParentStages(rdd, jobId), jobId, callSite)  
12.     stageIdToStage(id) = stage  
13.     updateJobIdStageIdMaps(jobId, stage)  
14.     stage  
15.   } 
    

 //1.6版本

private def newResultStage(
      rdd: RDD[_],
      func: (TaskContext, Iterator[_]) => _,
      partitions: Array[Int],
      jobId: Int,
      callSite: CallSite): ResultStage = {
    val (parentStages: List[Stage], id: Int) = getParentStagesAndId(rdd, jobId)
    val stage = new ResultStage(id, rdd, func, partitions, parentStages, jobId, callSite)
    stageIdToStage(id) = stage
    updateJobIdStageIdMaps(jobId, stage)
    stage
  }
对于result 的final stage来说，传入的shuffleDep是None。

我们知道，RDD通过org.apache.spark.rdd.RDD#getDependencies可以获得它依赖的parent RDD。而Stage也可能会有parent Stage。看一个RDD论文的Stage划分吧：

一个stage的边界，输入是外部的存储或者一个stage shuffle的结果；输入则是Job的结果（result task对应的stage）或者shuffle的结果。

上图的话stage3的输入则是RDD A和RDD F shuffle的结果。而A和F由于到B和G需要shuffle，因此需要划分到不同的stage。

从源码实现的角度来看，通过触发action也就是最后一个RDD创建final stage（上图的stage 3），我们注意到new Stage的第五个参数就是该Stage的parent Stage：通过rdd和job id获取：

1. // 生成rdd的parent Stage。没遇到一个ShuffleDependency，就会生成一个Stage  
2.   private def getParentStages(rdd: RDD[_], jobId: Int): List[Stage] = {  
3.     val parents = new HashSet[Stage] //存储parent stage  
4.     val visited = new HashSet[RDD[_]] //存储已经被访问到得RDD  
5.     // We are manually maintaining a stack here to prevent StackOverflowError  
6.     // caused by recursively visiting // 存储需要被处理的RDD。Stack中得RDD都需要被处理。  
7.     val waitingForVisit = new Stack[RDD[_]]  
8.     def visit(r: RDD[_]) {  
9.       if (!visited(r)) {  
10.         visited += r  
11.         // Kind of ugly: need to register RDDs with the cache here since  
12.         // we can't do it in its constructor because # of partitions is unknown  
13.         for (dep <- r.dependencies) {  
14.           dep match {  
15.             case shufDep: ShuffleDependency[_, _, _] => // 在ShuffleDependency时需要生成新的stage  
16.               parents += getShuffleMapStage(shufDep, jobId)  
17.             case _ =>  
18.               waitingForVisit.push(dep.rdd) //不是ShuffleDependency，那么就属于同一个Stage  
19.           }  
20.         }  
21.       }  
22.     }  
23.     waitingForVisit.push(rdd) // 输入的rdd作为第一个需要处理的RDD。然后从该rdd开始，顺序访问其parent rdd  
24.     while (!waitingForVisit.isEmpty) { //只要stack不为空，则一直处理。  
25.       visit(waitingForVisit.pop()) //每次visit如果遇到了ShuffleDependency，那么就会形成一个Stage，否则这些RDD属于同一个Stage  
26.     }  
27.     parents.toList  
28.   }  

生成了finalStage后，就需要提交Stage了。

1. // 提交Stage，如果有parent Stage没有提交，那么递归提交它。  
2. private def submitStage(stage: Stage) {  
3.   val jobId = activeJobForStage(stage)  
4.   if (jobId.isDefined) {  
5.     logDebug("submitStage(" + stage + ")")  
6.     // 如果当前stage不在等待其parent stage的返回，并且 不在运行的状态， 并且 没有已经失败（失败会有重试机制，不会通过这里再次提交）  
7.     if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {  
8.       val missing = getMissingParentStages(stage).sortBy(_.id)  
9.       logDebug("missing: " + missing)  
10.       if (missing == Nil) { // 如果所有的parent stage都已经完成，那么提交该stage所包含的task  
11.         logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")  
12.         submitMissingTasks(stage, jobId.get)  
13.       } else {  
14.         for (parent <- missing) { // 有parent stage为完成，则递归提交它  
15.           submitStage(parent)  
16.         }  
17.         waitingStages += stage  
18.       }  
19.     }  
20.   } else {  
21.     abortStage(stage, "No active job for stage " + stage.id)  
22.   }  
23. }  

DAGScheduler将Stage划分完成后，提交实际上是通过把Stage转换为TaskSet，然后通过TaskScheduler将计算任务最终提交到集群。其所在的位置如下图所示。

接下来，将分析Stage是如何转换为TaskSet，并最终提交到Executor去运行的。

参考：http://blog.csdn.net/anzhsoft/article/details/39859463

http://blog.csdn.net/jasonding1354/article/details/46895397