Spark1.6.3 Stage划分算法

在分析Stage划分算法之前，需要知道RDD的宽依赖和窄依赖的概念。
窄依赖:Narrow Dependency,即每个RDD的partition，仅仅依赖于父RDD中的一个partition，父RDD和子RDD的关系是对应关系，即简单的一对一的依赖关系。
宽依赖：Shuffle Dependency，从它的名字就能看出，它的本质就是shuffle。即每个父RDD partition中的数据，都会传入一部分到子RDD中，这样就会形成错综复杂的关系，不再是单一的一对一关系。他们之间发生的操作就是shuffle了。

从源码中看出，一个action操作就会触发一个job，其内部调用了sparkContext的runJob 方法

//action 操作的主要入口
def runJob[T, U: ClassTag](
  rdd: RDD[T],
  func: (TaskContext, Iterator[T]) => U,
  partitions: Seq[Int],
  resultHandler: (Int, U) => Unit): Unit = {
   if (stopped.get()) {
    throw new IllegalStateException("SparkContext has been shutdown")
}
val callSite = getCallSite
val cleanedFunc = clean(func)
logInfo("Starting job: " + callSite.shortForm)
if (conf.getBoolean("spark.logLineage", false)) {
  logInfo("RDD's recursive dependencies:\n" + rdd.toDebugString)
}
// 使用SparkContext在初始化的时候会创建dagScheduler,这里调用其内部runJob方法。
dagScheduler.runJob(rdd, cleanedFunc, partitions, callSite, resultHandler, localProperties.get)
progressBar.foreach(_.finishAll())
rdd.doCheckpoint()}

DAGScheduler 发送 JobSubmitted 消息

eventProcessLoop.post(JobSubmitted(
  jobId, rdd, func2, partitions.toArray, callSite, waiter,
  SerializationUtils.clone(properties)))

接收消息进行处理

private def doOnReceive(event: DAGSchedulerEvent): Unit = event match {
case JobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties) =>
// handleJobSubmitted job调度核心入口。
  dagScheduler.handleJobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties)
}

核心方法分析：

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 {
 // 使用触发job的最后一个rdd创建finalStage 并将stage加入DAGScheduler内部的缓存中
  finalStage = newResultStage(finalRDD, func, partitions, jobId, callSite)
} catch {
  case e: Exception =>
    logWarning("Creating new stage failed due to exception - job: " + jobId, e)
    listener.jobFailed(e)
    return
}
// 用finalstage创建一个job
val job = new ActiveJob(jobId, finalStage, 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()
//  将job加入内存缓存中
jobIdToActiveJob(jobId) = job
activeJobs += job
finalStage.setActiveJob(job)
val stageIds = jobIdToStageIds(jobId).toArray
val stageInfos = stageIds.flatMap(id => stageIdToStage.get(id).map(_.latestInfo))
listenerBus.post(
  SparkListenerJobStart(job.jobId, jobSubmissionTime, stageInfos, properties))
//将finalstage提交
// 这个方法的调用 其实会导致第一个stage的提交
// 并且导致其他所有的stage，都被放入waitStage队列里
submitStage(finalStage)
// 提交等待的stage
   submitWaitingStages()   
}

// 提交stage，这是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)) {
    // 调用 该方法 去获得当前stage的父stage
    // 很关键的方法
    val missing = getMissingParentStages(stage).sortBy(_.id)
    logDebug("missing: " + missing)
    // 反复递归调用， 直到最初的stage 没有父stage了
    // 那么此时就会首先提交这个第一个stage。即stage0， 其余的stage 此时全部都在watingstages里面，这样优先提交父stage。
    if (missing.isEmpty) {
      logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
      submitMissingTasks(stage, jobId.get)
    } else {
      // 递归调用submit（）方法 去提交父stage
      for (parent <- missing) {
        submitStage(parent)
      }
      // 并且将当前的stage。放入等待执行的stage队列中
      waitingStages += stage
    }
  }
} else {
  abortStage(stage, "No active job for stage " + stage.id, 
None)
  }
}

/**
* 获取某个stage的父stage
* 这个方法的意思是，就是说 对一个stage
* 如果他的最后一个rdd的所有依赖，都是窄依赖，那么久不会创建新的stage
* 但是，只要发现stage的rdd宽依赖了某个rdd，那么就用宽依赖的那个rdd。创建一个新的stage
* 然后立即返回新的stage
*/
// 关键方法
private def getMissingParentStages(stage: Stage): List[Stage] = {
val missing = new HashSet[Stage]
val visited = new HashSet[RDD[_]]
val waitingForVisit = new Stack[RDD[_]]
def visit(rdd: RDD[_]) {
  if (!visited(rdd)) {
    visited += rdd
    val rddHasUncachedPartitions = getCacheLocs(rdd).contains(Nil)
    if (rddHasUncachedPartitions) {
      // 遍历rdd的依赖
      // 每一种shuffle操作比如groupByKey reduceBykey countByKey
      // 底层对应了三种RDD：MapPartitionRDD  ShuffleRDD MapPartitionRDD
      // 从第一个开始 划入新的stage

      for (dep <-  rdd.dependencies) {
        dep match {
          // 那么使用宽依赖的那个rdd 创建一个stage 并且将isShuffleMap设置为true
            // 默认最后一个stage ，不是shuffleMap stage
            // 但是finalStage之前所有的stage 都是shuffleMap stage
          case shufDep: ShuffleDependency[_, _, _] =>
            val mapStage = getShuffleMapStage(shufDep, stage.firstJobId)
            if (!mapStage.isAvailable) {
              missing += mapStage
            }
            // 如果是窄依赖 那么将依赖的rdd放入栈中
          case narrowDep: NarrowDependency[_] =>
            waitingForVisit.push(narrowDep.rdd)
        }
      }
    }
  }
}
// 首先往栈中 推入了stage的最后一个rdd
waitingForVisit.push(stage.rdd)
// 然后进行while循环
while (waitingForVisit.nonEmpty) {
  // 对stage的最后一个rdd，调用自己内部的visit方法
  visit(waitingForVisit.pop())
}
missing.toList
}

Stage划分算法：从触发action操作的rdd开始往前倒推，首先会为最后一个rdd创建一个stage，在这个过程中，如果发现对某个rdd是宽依赖，那么就会将宽依赖的那个rdd创建一个新的stage，此时该rdd就是新的stage的最后一个rdd。然后继续向前推，根据宽依赖和窄依赖，进行stage划分，直到所有的rdd全部遍历完为止。