Spark源码解读之Stage划分和提交

上一篇讲解了Spark源码解读之Job提交，这一篇主要讲解Stage划分和提交。

调用流程：

org.apache.spark.scheduler.DAGScheduler.handleJobSubmitted

org.apache.spark.scheduler.DAGScheduler.submitStage

org.apache.spark.scheduler.DAGScheduler.submitMissingTasks

org.apache.spark.scheduler.TaskScheduler.submitTasks

一、Stage划分

Spark中会根据RDD之间的依赖关系进行Stage划分，在遇到ShuffleDependency时，会将这两个RDD划分到不同的Stage。在调用DAGScheduler的handleJobSubmitted进行Job提交后，会先进行Stage划分，源码如下：

// 参数finalRDD为触发action操作时最后一个RDDprivate[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.    // 创建finalStage    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  }  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()  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  submitStage(finalStage)  submitWaitingStages()}

可以看出，在创建finalStage时初始化了newResultStage实例。最后调用submitStage方法（详情见Stage提交部分）。newResultStage源码如下：

/** * Create a ResultStage associated with the provided jobId. */private def newResultStage(    rdd: RDD[_],    func: (TaskContext, Iterator[_]) => _,    partitions: Array[Int],    jobId: Int,    callSite: CallSite): ResultStage = {  // 获取parent Stages和id  val (parentStages: List[Stage], id: Int) = getParentStagesAndId(rdd, jobId)  // 创建stage  val stage = new ResultStage(id, rdd, func, partitions, parentStages, jobId, callSite)  // 更新stageId和Stage、JobId和stageId之间的映射关系  stageIdToStage(id) = stage  updateJobIdStageIdMaps(jobId, stage)  // 返回stage  stage}

初始化newResultStage实例时会做两件事，一是调用getParentStagesAndId方法得到parentStages和id，二是更新stageId和Stage、JobId和stageId之间的映射关系。下面是getParentStagesAndId源码：

/** * Helper function to eliminate some code re-use when creating new stages. */private def getParentStagesAndId(rdd: RDD[_], firstJobId: Int): (List[Stage], Int) = {  // 获取parentStages  val parentStages = getParentStages(rdd, firstJobId)  // 获取一个唯一id  val id = nextStageId.getAndIncrement()  (parentStages, id)}

getParentStagesAndId会做两件事，一是调用getParentStages得到parentStages列表，二是获取一个Stage的唯一id。getParentStages源码如下：

/** * Get or create the list of parent stages for a given RDD.  The new Stages will be created with * the provided firstJobId. */private def getParentStages(rdd: RDD[_], firstJobId: Int): List[Stage] = {  val parents = 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[_]]  // 广度优先遍历方式  def visit(r: RDD[_]) {    if (!visited(r)) { // visited为空      visited += r      // Kind of ugly: need to register RDDs with the cache here since      // we can't do it in its constructor because # of partitions is unknown      for (dep <- r.dependencies) {        dep match {        // 依赖为ShuffleDependency类型时，则生成一个新的shuffle map Stage          case shufDep: ShuffleDependency[_, _, _] =>            parents += getShuffleMapStage(shufDep, firstJobId)          // 依赖为非ShuffleDependency类型时，则加入到waitingForVisit栈中          case _ =>            waitingForVisit.push(dep.rdd)        }      }    }  }  waitingForVisit.push(rdd)    while (waitingForVisit.nonEmpty) {  // 调用visit方法    visit(waitingForVisit.pop())  }  parents.toList}

getParentStages会先创建一个类型为RDD的栈waitingForVisit，然后遍历waitingForVisit，如果该RDD的依赖为ShuffleDependency类型，则调用getShuffleMapStage方法得到一个shuffle map stage，否则将该RDD的父RDD加入到waitingForVisit中。getShuffleMapStage源码如下：

/** * Get or create a shuffle map stage for the given shuffle dependency's map side. */private def getShuffleMapStage(    shuffleDep: ShuffleDependency[_, _, _],    firstJobId: Int): ShuffleMapStage = {  shuffleToMapStage.get(shuffleDep.shuffleId) match {    case Some(stage) => stage    case None =>      // We are going to register ancestor shuffle dependencies      getAncestorShuffleDependencies(shuffleDep.rdd).foreach { dep =>        shuffleToMapStage(dep.shuffleId) = newOrUsedShuffleStage(dep, firstJobId)      }      // Then register current shuffleDep      val stage = newOrUsedShuffleStage(shuffleDep, firstJobId)      shuffleToMapStage(shuffleDep.shuffleId) = stage      stage  }}

getShuffleMapStage会获取或者创建一个shuffle map stage。

二、Stage提交

Stage划分完成后，会进行Stage提交，Stage提交首先会调用submitStage方法，源码如下：

/** Submits stage, but first recursively submits any missing parents. */private def submitStage(stage: Stage) {  val jobId = activeJobForStage(stage)  if (jobId.isDefined) {    logDebug("submitStage(" + stage + ")")    if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {    // 获取未提交的父Stage      val missing = getMissingParentStages(stage).sortBy(_.id)      logDebug("missing: " + missing)      if (missing.isEmpty) { // 所有的父Stage都已提交        logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")        submitMissingTasks(stage, jobId.get) // 提交该Stage      } else {// 父Stage会提交        for (parent <- missing) {          submitStage(parent) /// 提交父Stage        }        waitingStages += stage      }    }  } else {    abortStage(stage, "No active job for stage " + stage.id, None)  }}

submitStage会检测该Stage的父Stage是否提交，如果有父Stage未提交，则会递归调用submitStage；如果父Stage都已提交，则会调用submitMissingTasks方法提交该Stage。submitMissingTasks源码如下：

  /** Called when stage's parents are available and we can now do its task. */  private def submitMissingTasks(stage: Stage, jobId: Int) {    logDebug("submitMissingTasks(" + stage + ")")    // Get our pending tasks and remember them in our pendingTasks entry    stage.pendingPartitions.clear()    // First figure out the indexes of partition ids to compute.    // 得到需要计算的partitions    val partitionsToCompute: Seq[Int] = stage.findMissingPartitions()    // Create internal accumulators if the stage has no accumulators initialized.    // Reset internal accumulators only if this stage is not partially submitted    // Otherwise, we may override existing accumulator values from some tasks    if (stage.internalAccumulators.isEmpty || stage.numPartitions == partitionsToCompute.size) {      stage.resetInternalAccumulators()    }    // Use the scheduling pool, job group, description, etc. from an ActiveJob associated    // with this Stage    val properties = jobIdToActiveJob(jobId).properties    runningStages += stage    // SparkListenerStageSubmitted should be posted before testing whether tasks are    // serializable. If tasks are not serializable, a SparkListenerStageCompleted event    // will be posted, which should always come after a corresponding SparkListenerStageSubmitted    // event.        stage match {      case s: ShuffleMapStage =>        outputCommitCoordinator.stageStart(stage = s.id, maxPartitionId = s.numPartitions - 1)      case s: ResultStage =>        outputCommitCoordinator.stageStart(          stage = s.id, maxPartitionId = s.rdd.partitions.length - 1)    }        // 创建一个Map：taskIdToLocations，存储的是id->Seq[TaskLocation]的映射关系，这里的id表示task所包含的RDD的partition id，TaskLocation表示任务位置    // 实现时，对stage中需要计算的RDD的分区调用PreferredLocations来获取优先位置信息，映射成id->Seq[TaskLocation]的关系      val taskIdToLocations: Map[Int, Seq[TaskLocation]] = try {      stage match {        case s: ShuffleMapStage =>          partitionsToCompute.map { id => (id, getPreferredLocs(stage.rdd, id))}.toMap        case s: ResultStage =>          val job = s.activeJob.get          partitionsToCompute.map { id =>            val p = s.partitions(id)            (id, getPreferredLocs(stage.rdd, p))          }.toMap      }    } catch {      case NonFatal(e) =>        stage.makeNewStageAttempt(partitionsToCompute.size)        listenerBus.post(SparkListenerStageSubmitted(stage.latestInfo, properties))        abortStage(stage, s"Task creation failed: $e\n${e.getStackTraceString}", Some(e))        runningStages -= stage        return    }    stage.makeNewStageAttempt(partitionsToCompute.size, taskIdToLocations.values.toSeq)    listenerBus.post(SparkListenerStageSubmitted(stage.latestInfo, properties))    // TODO: Maybe we can keep the taskBinary in Stage to avoid serializing it multiple times.    // Broadcasted binary for the task, used to dispatch tasks to executors. Note that we broadcast    // the serialized copy of the RDD and for each task we will deserialize it, which means each    // task gets a different copy of the RDD. This provides stronger isolation between tasks that    // might modify state of objects referenced in their closures. This is necessary in Hadoop    // where the JobConf/Configuration object is not thread-safe.    var taskBinary: Broadcast[Array[Byte]] = null    try {      // For ShuffleMapTask, serialize and broadcast (rdd, shuffleDep).      // 对于ShuffleMapTask，序列化并广播，广播的是rdd和shuffleDep        // For ResultTask, serialize and broadcast (rdd, func).      // 对于ResultTask，序列化并广播，广播的是rdd和func      val taskBinaryBytes: Array[Byte] = stage match {        case stage: ShuffleMapStage =>          closureSerializer.serialize((stage.rdd, stage.shuffleDep): AnyRef).array()        case stage: ResultStage =>          closureSerializer.serialize((stage.rdd, stage.func): AnyRef).array()      }      taskBinary = sc.broadcast(taskBinaryBytes)    } catch {      // In the case of a failure during serialization, abort the stage.      case e: NotSerializableException =>        abortStage(stage, "Task not serializable: " + e.toString, Some(e))        runningStages -= stage        // Abort execution        return      case NonFatal(e) =>        abortStage(stage, s"Task serialization failed: $e\n${e.getStackTraceString}", Some(e))        runningStages -= stage        return    }// 针对stage的每个分区构造task，形成tasks:ShuffleMapStage生成ShuffleMapTasks，ResultStage生成ResultTasks      val tasks: Seq[Task[_]] = try {      stage match {        case stage: ShuffleMapStage =>          partitionsToCompute.map { id =>            val locs = taskIdToLocations(id)            val part = stage.rdd.partitions(id)            new ShuffleMapTask(stage.id, stage.latestInfo.attemptId,              taskBinary, part, locs, stage.internalAccumulators)          }        case stage: ResultStage =>          val job = stage.activeJob.get          partitionsToCompute.map { id =>            val p: Int = stage.partitions(id)            val part = stage.rdd.partitions(p)            val locs = taskIdToLocations(id)            new ResultTask(stage.id, stage.latestInfo.attemptId,              taskBinary, part, locs, id, stage.internalAccumulators)          }      }    } catch {      case NonFatal(e) =>        abortStage(stage, s"Task creation failed: $e\n${e.getStackTraceString}", Some(e))        runningStages -= stage        return    }    // 如果存在tasks，则利用taskScheduler.submitTasks()提交task，否则标记stage已完成      if (tasks.size > 0) {      logInfo("Submitting " + tasks.size + " missing tasks from " + stage + " (" + stage.rdd + ")")      stage.pendingPartitions ++= tasks.map(_.partitionId)      logDebug("New pending partitions: " + stage.pendingPartitions)      // 调用taskScheduler.submitTasks提交task      taskScheduler.submitTasks(new TaskSet(        tasks.toArray, stage.id, stage.latestInfo.attemptId, jobId, properties))      // 记录提交时间      stage.latestInfo.submissionTime = Some(clock.getTimeMillis())    } else {      // Because we posted SparkListenerStageSubmitted earlier, we should mark      // the stage as completed here in case there are no tasks to run      markStageAsFinished(stage, None)      val debugString = stage match {        case stage: ShuffleMapStage =>          s"Stage ${stage} is actually done; " +            s"(available: ${stage.isAvailable}," +            s"available outputs: ${stage.numAvailableOutputs}," +            s"partitions: ${stage.numPartitions})"        case stage : ResultStage =>          s"Stage ${stage} is actually done; (partitions: ${stage.numPartitions})"      }      logDebug(debugString)    }  }

submitMissingTasks会做以下几个事：

1. 清空stage的pendingPartitions

2. 得到需要计算的partition id索引，放入partitionsToCompute

3. 将stage加入到runningStages中 

4. 启动一个stage

5. 得到task中执行的位置，即计算stage的每个RDD的partition的优先位置，存入taskIdToLocations

6. 对stage进行序列化并广播

7. （重要）针对stage的每个RDD的partition构造task，存入tasks

8. 存在tasks，则调用taskScheduler.submitTasks()提交task，否则标记stage已完成。