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操作时最后一个RDD
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.
    // 创建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已完成。