本文介绍Spark任务调度框架中Stage的原理，并分析其实现机制。

Stage的基本概念

用户提交的计算任务是一个由RDD依赖构成的DAG，Spark会把RDD的依赖以shuffle依赖为边界划分成多个Stage，这些Stage之间也相互依赖，形成了Stage的DAG。然后，DAGScheduler会按依赖关系顺序执行这些Stage。

要是把RDD依赖构成的DAG看成是逻辑执行计划(logic plan)，那么，可以把Stage看成物理执行计划，为了更好的理解这个概念，我们来看一个例子。

下面的代码用来对README.md文件中包含整数值的单词进行计数，并打印RDD之间的依赖关系（Lineage）：

 scala> val counts = sc.textFile("README.md")
 .flatMap(x=>x.split("\\W+"))
 .filter(_.matches(".*\\d.*"))
 .map(x=>(x,1))
 .reduceByKey(_+_)
 // 调用一个action函数，用来触发任务的提交和执行
 scala> counts.collect()
 
 // 打印RDD的依赖关系(Lineage)
 scala> counts.toDebugString
 res7: String =
 (2) ShuffledRDD[17] at reduceByKey at <console>:24 []
  +-(2) MapPartitionsRDD[16] at map at <console>:24 []
     |  MapPartitionsRDD[15] at filter at <console>:24 []
     |  MapPartitionsRDD[14] at flatMap at <console>:24 []
     |  README.md MapPartitionsRDD[13] at textFile at <console>:24[]
     |  README.md HadoopRDD[12] at textFile at <console>:24 []

如图1所示，从图中可以看到从RDD依赖关系到Stage的转换：

Spark任务调度-Stage的划分与提交_java

图1 RDD依赖关系到Stage的转换

通过Spark UI的DAG Visualization页面可以证实这一点：

Spark任务调度-Stage的划分与提交_java_02

图2 Spark UI的Stage查看

从图1和图2中可以看到，由于flatMap、filter、map这些操作都不会产生shuffle，所以被划分到了一个Stage中，而reduceByKey会产生shuffle操作，所以被划分到了一个新的Stage中。

从概念上讲，Stage是一个并行执行的计算任务集，该任务集中的每个计算任务都会有相同的shuffle依赖，并会执行相同的处理函数，它们作为Spark Job执行的一部分。

Stage的划分

从上面的分析可知，DAGScheduler会根据RDD的依赖关系来划分Stage。那么，划分Stage的依据是什么呢？实际上，Spark是根据RDD的窄依赖和宽依赖来划分Stage的。

在《RDD转换操作原理》一节中分析过：在计算RDD时，宽依赖会产生Shuffle。为了最大程度提高计算性能，减少网络数据传输，Spark会把窄依赖划分到一个Stage，直到出现了宽依赖。

也就是说，每个Stage包含一系列RDD的窄转换（narrow transformations），这些转换操作可以在不进行Shuffling的情况下完成计算，这些Stage在Shuffle的边界(例如：shuffle发生的地方)处被分开，因此，可以说Stage是RDD依赖关系（RDD Lineage）在Shuffle阶段分裂的结果。

在每个Stage中，RDD的窄转换(例如：map()或filter()等)操作形成流水线（pipeline），中间结果可以保存在内存中，以便加快计算性能，这些任务形成一个任务集（TaskSet），但是Shuffle操作却需要依赖一个或多个Stage。

Spark任务调度-Stage的划分与提交_java_03图3 Stage的划分示意图

图3是一个如何根据RDD的依赖来划分Stage的示意图。为了更好的理解Stage的划分，下面分析一下图3的Stage的划分过程：

从rddF开始，rddF的计算依赖rddE，由于他们是窄依赖（map操作），不会产生Shuffle。所以rddF和rddE会被划分到一个Stage中，我们继续看rddE对其父RDD的依赖。

rddE依赖于两个RDD：rddB和rddD。由于rddE对这两个RDD的依赖都是宽依赖（intersection操作会产生Shuffle），此时就要作为Stage的分界点，所以会把rddE和rddF划分到一个Stage中（即：Stage 4）。

而由于rddE和rddB，rddD都是宽依赖，所以将不会和rddE划分到一个Stage中。也就是说，rddD和rddB分别在一个新的Stage中。下面继续分析rddB和rddD的依赖。

rddB依赖rddA，而rddB和rddA是宽依赖，因为在计算rddB时会调用rddA的reduceByKey函数，此时可能会产生Shuffle，所以，rddB会被单独划分成一个Stage（即：Stage3），而rddA不依赖任何其他的RDD，也会被划分到一个新的Stage中（即：Stage1）。

在看rddD，由于rddD是rddC通过map转换得到的，这个过程没有Shuffle产生，所以会把rddD和rddC的依赖划分到一个Stage中，而rddC不再和其他的RDD产生依赖，所以，rddC和rddD的依赖也就产生了Stage2。

另外，从图3中可以看出，Stage1和Stage2是相互独立的，可以并发执行，而Stage3必须要等到Stage1完成后才能执行，而Stage4必须要Stage3和Stage2都完成后才能执行。

Stage的分类和实现

在Spark中有两类Stage：

  • ResultStage：执行action操作的函数，得到最终的结果。

  • ShuffleMapStage：用来计算RDD中间结果。

这两类Stage分别由ResultStage和ShuffleMapStage类来实现，另外，这两种类型的Stage需要遵循Stage的实现合约(Stage抽象类)。类图关系如下：

Spark任务调度-Stage的划分与提交_java_04

Stage抽象类声明如下：

 private[scheduler] abstract class Stage

为了能够更好的理解后面的两种具体的Stage，下面对Stage的抽象类中重要的成员做一个说明：

成员名说明
IdStage的唯一标识。这是一个整数，每个Stage对象都是唯一的。
rdd运行该Stage的RDD。
numTasks该Stage中的任务总数。要注意的是：有些ResultStages可能不需要计算所有分区，例如：first()，lookup()，take()等。
parents这是一个List[Stage]，它是该Stage依赖的Stage列表。
firstJobId对于FIFO调度，此变量是此Stage属于的第一个Job的ID。
numPartitionsRDD的分区数。
jobIds该Stage属于的Job集。
findMissingPartitions返回需要计算(missing)但还没计算的分区id集合

ResultStage

ResultStage是Job的最后一个Stage，该Stage是基于执行action函数的rdd来创建的。该Stage用来计算一个action操作的结果。该类的声明如下：

 private[spark] class ResultStage(
     id: Int,
     rdd: RDD[_],
     val func: (TaskContext, Iterator[_]) => _,
     val partitions: Array[Int],
     parents: List[Stage], //依赖的父Stage
     firstJobId: Int,
     callSite: CallSite)
   extends Stage(id, rdd, partitions.length, parents, firstJobId,callSite) {

为了计算action操作的结果，ResultStage会在目标RDD的一个或多个分区上使用函数：func。需要计算的分区id集合保存在成员变量：partitions中。但对于有些action操作，比如：first()，take()等，函数:func可能不会在所有分区上使用。

另外，在提交Job时，会先创建ResultStage。但在提交Stage时，会先递归找到该Stage依赖的父级Stage，并先提交父级Stage。如下图所示：

Spark任务调度-Stage的划分与提交_java_05

ShuffleMapStage

ShuffleMapStages是在DAG执行过程中产生的Stage，用来为Shuffle产生数据。ShuffleMapStages发生在每个Shuffle操作之前，在Shuffle之前可能有多个窄转换操作，比如：map，filter，这些操作可以形成流水线（pipeline）。当执行ShuffleMapStages时，会产生Map的输出文件，这些文件会被随后的Reduce任务使用。

ShuffleMapStages也可以作为Jobs，通过DAGScheduler.submitMapStage函数单独进行提交。对于这样的Stages，会在变量mapStageJobs中跟踪提交它们的ActiveJobs。要注意的是，可能有多个ActiveJob尝试计算相同的ShuffleMapStages。

它为一个shuffle过程产生map操作的输出文件。它也可能是自适应查询规划/自适应调度工作的最后阶段。如下图所示：

Spark任务调度-Stage的划分与提交_java_06

Stage提交的实现

通过前面的分析我们知道，RDD的每个Action操作都会提交一个Job，而DAGScheduler会根据RDD的依赖关系把这个Job划分成一个或多个相互依赖的Stage，从而形成一个DAG。然后，根据Stage的DAG来执行Stage。那么，这个过程是如何实现的呢？下面就通过代码层面来分析一下实现过程：

Stage提交流程的函数调用顺序

 RDD.count()// 通过RDD的action操作来提交Job
 SparkContext#runJob(...)  // 调用SparkContext中的runJob函数
 DAGScheduler#runJob(...)  // 调用DAGScheduler中的runJob函数
 DAGScheduler#submitJob(...) // 提交Job。实际上是向事件总线提交JobSubmitted事件
 DAGSchedulerEventProcessLoop#post(JobSubmitted(...)) //发送Job提交事件
 DAGSchedulerEventProcessLoop#doOnReceive(event:DAGSchedulerEvent) //获取DAGSchedulerEvent
 DAGScheduler#handleJobSubmitted(...)  //处理Job提交事件
 DAGScheduler#submitStage(stage: Stage)  //提交Stage

提交过程的代码实现

从以上的函数调用流程来看，事件的处理是在handleJobSubmitted函数中进行的，我们来看一下该函数是如何处理Job提交事件的。代码如下（省去不相关代码）：

 private[scheduler] def handleJobSubmitted(...) {
     var finalStage: ResultStage = null
  ...
     // 首先，创建ResultStage和他依赖的Stage，形成一个Stage的DAG
     finalStage = createResultStage(finalRDD, func, partitions, jobId,callSite)
    ...
     // 提交finalStage
     submitStage(finalStage)  
 }

Stage DAG的构建

Stage（包括父Stage）的创建，都是在createResultStage函数中完成的。下面分析createResultStage的实现，该函数的代码如下：

   private def createResultStage(...): ResultStage = {
    ...
     // 递归创建Stage的父Stage，及其祖先Stage，最终形成Stage的DAG
     val parents = getOrCreateParentStages(rdd, jobId)
     val id = nextStageId.getAndIncrement()
     // 创建ResultStage
     val stage = new ResultStage(id, rdd, func, partitions, parents,jobId, callSite)
 ...
     stage
  }

createResultStage函数的基本逻辑如下图所示：

Spark任务调度-Stage的划分与提交_java_07

如上图所示：createResultStage函数创建Stage的过程是通过递归调用实现的。

在getOrCreateParentStages函数中，会基于stage的rdd的依赖关系向上查找其父rdd，若是窄依赖则继续向上查找；若是宽依赖（shuffle依赖）则会调用getOrCreateShuffleMapStage函数来创建一个ShuffleMapStage。

Stage的提交

通过函数createResultStage创建了一个Stage的DAG，并可以通过finalStage（它是一个ResultStage）来获取Stage的依赖关系（DAG）。执行到这里所有的Stage（包括Stage的依赖关系）就创建完成了，此时就可以开始提交Stage了。

Stage的提交时通过submitStage函数来实现的。该函数的主要实现逻辑如下图所示：

Spark任务调度-Stage的划分与提交_java_08

可见在提交Stage时，也是通过递归提交最先依赖的Stage，最后提交ResultStage。其实现代码如下：

 /** Submits stage, but first recursively submits any missing parents. */
   private def submitStage(stage: Stage) {
     val jobId = activeJobForStage(stage)
     
     if (jobId.isDefined) {
      ...
       if (!waitingStages(stage) && !runningStages(stage) &&!failedStages(stage)) {
         // 查找并获取依赖的父Stage
         val missing = getMissingParentStages(stage).sortBy(_.id)
        ...
         if (missing.isEmpty) {
            ...
           // 已经找到全部的依赖Stage并已提交，最后提交最后一个Stage
           submitMissingTasks(stage, jobId.get)
        } else {
           // 先提交依赖的父Stage
           for (parent <- missing) {
             submitStage(parent)
          }
           waitingStages += stage
        }
      }
    } else {
       abortStage(stage, "No active job for stage " + stage.id, None)
    }
  }

Stage和RDD

在创建Stage时，只会记录以shuffle依赖为边界的最后一个RDD。如下：

  Stage(
    val id: Int,
    val rdd: RDD[_], // rdd是每个Stage的最后一个RDD
    ...
  )

在创建任务时，会调用该RDD的计算函数，由于在每个Stage中RDD都是相互依赖的，而且同一个Stage中的RDD都是窄依赖，这意味着同一个Stage中RDD的分区计算的中间过程可以形成pipeline（不需要持久化到磁盘）。所以，当计算某个Stage中的最后一个RDD分区数据时，会根据依赖关系计算其依赖的父RDD的分区数据，并且会以pipeline的方式执行。

总结

本文说明了Stage的实现原理，并对Stage的提交过程进行了分析。

c